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)