Create main.py

main.py

@@ -0,0 +1,273 @@
+import gradio as gr
+from transformers import pipeline
+import os
+import cv2
+from ultralytics import YOLO
+import shutil # Import shutil for copying files
+import zipfile # Import zipfile for creating zip archives
+
+def multi_model_detection(image_paths_list: list, model_paths_list: list, output_dir: str = 'detection_results', conf_threshold: float = 0.25):
+    """
+    使用多個 YOLOv8 模型對多張圖片進行物件辨識，
+    並將結果繪製在圖片上，同時保存辨識資訊到文字檔案。
+
+    Args:
+        image_paths_list (list): 包含所有待辨識圖片路徑的列表。
+        model_paths_list (list): 包含所有模型 (.pt 檔案) 路徑的列表。
+        output_dir (str): 儲存結果圖片和文字檔案的目錄。
+                         如果不存在，函式會自動創建。
+        conf_threshold (float): 置信度閾值，只有高於此值的偵測結果會被標示。
+
+    Returns:
+        list: A list of paths to the annotated images.
+        list: A list of paths to the text files with detection information.
+    """
+
+    # 確保輸出目錄存在
+    if not os.path.exists(output_dir):
+        os.makedirs(output_dir)
+        print(f"已創建輸出目錄: {output_dir}")
+
+    # 載入所有模型
+    loaded_models = []
+    print("\n--- 載入模型 ---")
+    # If no models are uploaded, use the default yolov8n.pt
+    if not model_paths_list:
+        default_model_path = 'yolov8n.pt'
+        try:
+            model = YOLO(default_model_path)
+            loaded_models.append((default_model_path, model))
+            print(f"成功載入預設模型: {default_model_path}")
+        except Exception as e:
+            print(f"錯誤: 無法載入預設模型 '{default_model_path}' - {e}")
+            return [], []
+    else:
+        for model_path in model_paths_list:
+            try:
+                model = YOLO(model_path)
+                loaded_models.append((model_path, model)) # 儲存模型路徑和模型物件
+                print(f"成功載入模型: {model_path}")
+            except Exception as e:
+                print(f"錯誤: 無法載入模型 '{model_path}' - {e}")
+                continue # 如果模型載入失敗，跳過它
+
+
+    if not loaded_models:
+        print("沒有模型成功載入，請檢查模型路徑或預設模型。")
+        return [], []
+
+    annotated_image_paths = []
+    txt_output_paths = []
+
+    # 處理每張圖片
+    print("\n--- 開始圖片辨識 ---")
+    for image_path in image_paths_list:
+        if not os.path.exists(image_path):
+            print(f"警告: 圖片 '{image_path}' 不存在，跳過。")
+            continue
+
+        print(f"\n處理圖片: {os.path.basename(image_path)}")
+        original_image = cv2.imread(image_path)
+        if original_image is None:
+            print(f"錯誤: 無法讀取圖片 '{image_path}'，跳過。")
+            continue
+
+        # 複製圖片用於繪製，避免修改原始圖片
+        # 使用 NumPy 複製，而不是直接賦值
+        annotated_image = original_image.copy()
+
+        # 準備寫入文字檔的內容
+        txt_output_content = []
+        txt_output_content.append(f"檔案: {os.path.basename(image_path)}\n")
+
+        # 對每張圖片使用所有模型進行辨識
+        all_detections_for_image = [] # 儲存所有模型在當前圖片上的偵測結果
+
+        for model_path_str, model_obj in loaded_models:
+            model_name = os.path.basename(model_path_str) # 獲取模型檔案名
+            print(f"  使用模型 '{model_name}' 進行辨識...")
+
+            # 執行推論, device="cpu" ensures it runs on CPU if GPU is not available or preferred
+            results = model_obj(image_path, verbose=False, device="cpu")[0]
+
+            # 將辨識結果添加到 txt 輸出內容和繪圖列表
+            txt_output_content.append(f"\n--- 模型: {model_name} ---")
+
+            if results.boxes: # 檢查是否有偵測到物件
+                for box in results.boxes:
+                    # 取得邊界框座標和置信度
+                    conf = float(box.conf[0])
+                    if conf >= conf_threshold: # 檢查置信度是否達到閾值
+                        x1, y1, x2, y2 = map(int, box.xyxy[0])
+                        cls_id = int(box.cls[0])
+                        cls_name = model_obj.names[cls_id] # 取得類別名稱
+
+                        detection_info = {
+                            'model_name': model_name,
+                            'class_name': cls_name,
+                            'confidence': conf,
+                            'bbox': (x1, y1, x2, y2)
+                        }
+                        all_detections_for_image.append(detection_info)
+
+                        # 加入到文字檔內容
+                        txt_output_content.append(f"  - {cls_name} (Conf: {conf:.2f}) [x1:{x1}, y1:{y1}, x2:{x2}, y2:{y2}]")
+            else:
+                txt_output_content.append("  沒有偵測到���何物件。")
+
+        # 繪製所有模型在當前圖片上的偵測結果
+        # 我們會根據模型來源給予不同的顏色或樣式，讓結果更容易區分
+
+        # 定義一個顏色循環列表，方便給不同模型分配不同顏色
+        colors = [
+            (255, 0, 0),    # 紅色 (例如給模型 A)
+            (0, 255, 0),    # 綠色 (例如給模型 B)
+            (0, 0, 255),    # 藍色
+            (255, 255, 0),  # 黃色
+            (255, 0, 255),  # 紫色
+            (0, 255, 255),  # 青色
+            (128, 0, 0),    # 深紅
+            (0, 128, 0)     # 深綠
+        ]
+        color_map = {} # 用來映射模型名稱到顏色
+
+        for idx, (model_path_str, _) in enumerate(loaded_models):
+            model_name = os.path.basename(model_path_str)
+            color_map[model_name] = colors[idx % len(colors)] # 確保顏色循環使用
+
+        for det in all_detections_for_image:
+            x1, y1, x2, y2 = det['bbox']
+            conf = det['confidence']
+            cls_name = det['class_name']
+            model_name = det['model_name']
+
+            color = color_map.get(model_name, (200, 200, 200)) # 預設灰色
+
+            # 繪製邊界框
+            cv2.rectangle(annotated_image, (x1, y1), (x2, y2), color, 2)
+
+            # 繪製標籤 (類別名稱 + 置信度 + 模型名稱縮寫)
+            # 為了避免標籤過長，模型名稱只取前幾個字母
+            model_abbr = "".join([s[0] for s in model_name.split('.')[:-1]]) # 例如 'a.pt' -> 'a'
+            label = f'{cls_name} {conf:.2f} ({model_abbr})'
+            cv2.putText(annotated_image, label, (x1, y1 - 10),
+                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
+
+        # 保存繪製後的圖片
+        image_base_name = os.path.basename(image_path)
+        image_name_without_ext = os.path.splitext(image_base_name)[0]
+        output_image_path = os.path.join(output_dir, f"{image_name_without_ext}_detected.jpg")
+        cv2.imwrite(output_image_path, annotated_image)
+        annotated_image_paths.append(output_image_path)
+        print(f"  結果圖片保存至: {output_image_path}")
+
+        # 保存辨識資訊到文字檔案
+        output_txt_path = os.path.join(output_dir, f"{image_name_without_ext}.txt")
+        with open(output_txt_path, 'w', encoding='utf-8') as f:
+            f.write("\n".join(txt_output_content))
+        txt_output_paths.append(output_txt_path)
+        print(f"  辨識資訊保存至: {output_txt_path}")
+
+
+    print("\n--- 所有圖片處理完成 ---")
+    return annotated_image_paths, txt_output_paths
+
+def create_zip_archive(files, zip_filename):
+    """Creates a zip archive from a list of files."""
+    with zipfile.ZipFile(zip_filename, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        for file in files:
+            if os.path.exists(file):
+                zipf.write(file, os.path.basename(file))
+            else:
+                print(f"警告: 檔案 '{file}' 不存在，無法加入壓縮檔。")
+    return zip_filename
+
+
+# --- Gradio Interface ---
+def gradio_multi_model_detection(image_files, model_files, conf_threshold, output_subdir):
+    """
+    Gradio 的主要處理函式。
+    接收上傳的檔案和參數，呼叫後端辨識函式，並返回結果。
+
+    Args:
+        image_files (list): Gradio File 元件回傳的圖片檔案列表 (暫存路徑)。
+        model_files (list): Gradio File 元件回傳的模型檔案列表 (暫存路徑)。
+        conf_threshold (float): 置信度閾值。
+        output_subdir (str): 用於儲存本次執行結果的子目錄名稱。
+
+    Returns:
+        tuple: 更新 Gradio 介面所需的多個輸出。
+    """
+    if not image_files:
+        return None, "請上傳圖片檔案。", None, None
+
+    # Get the temporary file paths from Gradio File objects
+    image_paths = [file.name for file in image_files]
+    # Use uploaded model paths or an empty list if none are uploaded
+    model_paths = [file.name for file in model_files] if model_files else []
+
+
+    # Define the output directory for this run within the main results directory
+    base_output_dir = 'gradio_detection_results'
+    run_output_dir = os.path.join(base_output_dir, output_subdir)
+
+    # Perform detection
+    annotated_images, detection_texts = multi_model_detection(
+        image_paths_list=image_paths,
+        model_paths_list=model_paths,
+        output_dir=run_output_dir,
+        conf_threshold=conf_threshold
+    )
+
+    if not annotated_images:
+         return None, "辨識失敗，請檢查輸入或模型。", None, None
+
+    # Combine detection texts for display in one textbox
+    combined_detection_text = "--- 辨識結果 ---\n\n"
+    for txt_path in detection_texts:
+        with open(txt_path, 'r', encoding='utf-8') as f:
+            combined_detection_text += f.read() + "\n\n"
+
+    # Create a zip file containing both annotated images and text files
+    all_result_files = annotated_images + detection_texts
+    zip_filename = os.path.join(run_output_dir, f"{output_subdir}_results.zip")
+    created_zip_path = create_zip_archive(all_result_files, zip_filename)
+
+
+    # Return annotated images and combined text for Gradio output
+    # Gradio Gallery expects a list of image paths
+    return annotated_images, combined_detection_text, f"結果儲存於: {os.path.abspath(run_output_dir)}", created_zip_path
+
+
+# Create the Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# 支援多模型YOLO物件辨識(demo)")
+    gr.Markdown("上傳您的圖片和模型，並設定置信度閾值進行物件辨識。若未上傳模型，將使用預設的 yolov8n.pt 進行辨識。")
+
+    with gr.Row():
+        with gr.Column():
+            image_input = gr.File(label="上傳圖片", file_count="multiple", file_types=["image"])
+            model_input = gr.File(label="上傳模型 (.pt)", file_count="multiple", file_types=[".pt"])
+            conf_slider = gr.Slider(minimum=0, maximum=1, value=0.25, step=0.05, label="置信度閾值")
+            output_subdir_input = gr.Textbox(label="結果子目錄名稱", value="run_1", placeholder="請輸入儲存結果的子目錄名稱")
+            run_button = gr.Button("開始辨識")
+
+        with gr.Column():
+            # show_label=False hides the class name label below each image
+            # allow_preview=True enables double-clicking to zoom
+            # allow_download=True adds a download button for each image in the gallery
+            output_gallery = gr.Gallery(label="辨識結果圖片", height=400, allow_preview=True, object_fit="contain")
+            output_text = gr.Textbox(label="辨識資訊", lines=10)
+            output_status = gr.Textbox(label="狀態/儲存路徑")
+            download_button = gr.File(label="下載所有結果 (.zip)", file_count="single")
+
+
+    # Link the button click to the function
+    run_button.click(
+        fn=gradio_multi_model_detection,
+        inputs=[image_input, model_input, conf_slider, output_subdir_input],
+        outputs=[output_gallery, output_text, output_status, download_button]
+    )
+
+# Launch the interface
+demo.launch(debug=True)