# -*- coding: utf-8 -*- """ 系統需求: - gradio: 用於建立 Web UI - opencv-python: 用於圖片處理 - ultralytics: YOLOv8 官方函式庫 - Pillow: 圖片處理基礎庫 - transformers: (可選，若YOLO模型需要) """ import gradio as gr import os import cv2 from ultralytics import YOLO import shutil import zipfile import uuid # 匯入 uuid 以生成唯一的執行 ID from pathlib import Path # 匯入 Path 以更方便地操作路徑 import gemini_ai as genai from datetime import datetime import mongo_lib as mongo def create_zip_archive(files, zip_filename): """ 將一系列檔案壓縮成一個 zip 檔案。 Args: files (list): 要壓縮的檔案路徑列表。 zip_filename (str): 產生的 zip 檔案路徑。 Returns: str: 產生的 zip 檔案路徑。 """ with zipfile.ZipFile(zip_filename, 'w', zipfile.ZIP_DEFLATED) as zipf: for file in files: if os.path.exists(file): # 使用 os.path.basename 確保只寫入檔案名稱，而非完整路徑 zipf.write(file, os.path.basename(file)) else: print(f"警告: 檔案 '{file}' 不存在，無法加入壓縮檔。") return zip_filename def gradio_multi_model_detection( image_files, model_files, conf_threshold, enable_mllm, mllm_prompt, progress=gr.Progress(track_tqdm=True) ): """ Gradio 的主要處理函式，使用生成器 (yield) 實現流式輸出。 Args: image_files (list): Gradio File 元件回傳的圖片檔案列表。 model_files (list): Gradio File 元件回傳的模型檔案列表。 conf_threshold (float): 置信度閾值。 enable_mllm (bool): 是否啟用 MLLM 分析。 mllm_prompt (str): 使用者自訂的 MLLM prompt。 progress (gr.Progress): Gradio 的進度條元件。 Yields: dict: 用於更新 Gradio 介面元件的字典。 """ global_datetime = datetime.now() #寫主表log document = {"log_style":"master", "create_datetime": str(global_datetime), "image_files": image_files, "model_files": model_files, "conf_threshold":conf_threshold, "enable_mllm":enable_mllm, "mllm_prompt":mllm_prompt } mongo.insert_mongodb_log("multi_model_detection",document) #寫入log方便日後查驗 if not image_files: yield { output_status: gr.update(value="錯誤：請至少上傳一張圖片。"), output_gallery: None, output_text: None, download_button: None } return # --- 1. 初始化設定 --- # 為本次執行創建一個唯一的子目錄 run_id = str(uuid.uuid4()) base_output_dir = Path('gradio_detection_results') run_output_dir = base_output_dir / f"run_{run_id[:8]}" run_output_dir.mkdir(parents=True, exist_ok=True) image_paths = [file.name for file in image_files] model_paths = [file.name for file in model_files] if model_files else [] # --- 2. 載入模型 --- yield {output_status: gr.update(value="正在載入模型...")} loaded_models = [] if not model_paths: # 如果沒有上傳模型，使用預設模型 default_model_path = 'yolov8n.pt' try: model = YOLO(default_model_path) loaded_models.append((default_model_path, model)) except Exception as e: yield {output_status: gr.update(value=f"錯誤: 無法載入預設模型 '{default_model_path}' - {e}")} return else: for model_path in model_paths: try: model = YOLO(model_path) loaded_models.append((model_path, model)) except Exception as e: print(f"警告: 無法載入模型 '{model_path}' - {e}，將跳過此模型。") continue if not loaded_models: yield {output_status: gr.update(value="錯誤: 沒有任何模型成功載入。")} return # --- 3. 逐一處理圖片 --- total_images = len(image_paths) annotated_image_paths = [] all_result_files = [] # results_map 儲存圖片路徑與其對應的文字檔路徑，用於後續點擊查詢 results_map = {} # all_texts 用於收集所有圖片的辨識結果文字 all_texts = [] for i, image_path_str in enumerate(image_paths): image_path = Path(image_path_str) progress(i / total_images, desc=f"處理中: {image_path.name}") yield { output_status: gr.update(value=f"處理中... ({i+1}/{total_images}) - {image_path.name}"), output_gallery: gr.update(value=annotated_image_paths) } original_image = cv2.imread(str(image_path)) if original_image is None: print(f"警告: 無法讀取圖片 '{image_path}'，跳過。") continue annotated_image = original_image.copy() image_base_name = image_path.stem # --- 3a. YOLO 物件偵測 --- yolo_output_content = [f"--- 檔案: {image_path.name} ---"] all_detections_for_image = [] for model_path_str, model_obj in loaded_models: model_name = Path(model_path_str).name yolo_output_content.append(f"--- 模型: {model_name} ---") results = model_obj(str(image_path), verbose=False, device="cpu")[0] 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) yolo_output_content.append(f" - {cls_name} (信賴度: {conf:.2f}) [座標: {x1},{y1},{x2},{y2}]") else: yolo_output_content.append(" 未偵測到任何物件。") # 繪製偵測框 colors = [(255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255), (0, 255, 255)] color_map = {Path(p).name: colors[idx % len(colors)] for idx, (p, _) in enumerate(loaded_models)} for det in all_detections_for_image: x1, y1, x2, y2 = det['bbox'] color = color_map.get(det['model_name'], (200, 200, 200)) label = f"{det['class_name']} {det['confidence']:.2f}" cv2.rectangle(annotated_image, (x1, y1), (x2, y2), color, 2) cv2.putText(annotated_image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2) # 儲存 YOLO 標註圖 output_image_path = run_output_dir / f"{image_base_name}_yolo_detected.jpg" cv2.imwrite(str(output_image_path), annotated_image) annotated_image_paths.append(str(output_image_path)) all_result_files.append(str(output_image_path)) # 儲存 YOLO 辨識資訊 output_yolo_txt_path = run_output_dir / f"{image_base_name}_yolo_objects.txt" output_yolo_txt_path.write_text("\n".join(yolo_output_content), encoding='utf-8') all_result_files.append(str(output_yolo_txt_path)) # --- 3b. MLLM 分析 (如果啟用) --- output_mllm_txt_path = None mllm_result_content = "" if enable_mllm: try: prompt_to_use = mllm_prompt if mllm_prompt and mllm_prompt.strip() else None mllm_str = genai.analyze_content_with_gemini(str(image_path), prompt_to_use) mllm_result_content = f"--- MLLM 分析結果 ---\n{mllm_str}" except Exception as e: mllm_result_content = f"--- MLLM 分析失敗 ---\n原因: {e}" output_mllm_txt_path = run_output_dir / f"{image_base_name}_mllm_result.txt" output_mllm_txt_path.write_text(mllm_result_content, encoding='utf-8') all_result_files.append(str(output_mllm_txt_path)) #寫明細表log document = {"log_style":"detail", "create_datetime": str(global_datetime), "image_path": str(image_path), "yolo_result": yolo_output_content, "enable_mllm": enable_mllm, "mllm_prompt": mllm_prompt, "mllm_result": mllm_result_content} mongo.insert_mongodb_log("multi_model_detection",document) #寫入log方便日後查驗 # 將本次圖片的結果加入到總列表中 all_texts.append("\n".join(yolo_output_content)) if output_mllm_txt_path: all_texts.append(output_mllm_txt_path.read_text(encoding='utf-8')) # --- 4. 完成處理，打包並更新最終結果 --- progress(1, desc="打包結果中...") zip_filename = run_output_dir / f"run_{run_id[:8]}_results.zip" created_zip_path = create_zip_archive(all_result_files, str(zip_filename)) final_status = f"處理完成！共 {total_images} 張圖片。結果儲存於: {run_output_dir.absolute()}" combined_text_output = "\n\n".join(all_texts) yield { output_status: gr.update(value=final_status), download_button: gr.update(value=created_zip_path, visible=True), output_text: gr.update(value=combined_text_output), output_gallery: gr.update(value=annotated_image_paths) # 確保最終 gallery 也被更新 } def toggle_mllm_prompt(is_enabled): """ 根據 Checkbox 狀態，顯示或隱藏 MLLM prompt 輸入框。 """ return gr.update(visible=is_enabled) # --- Gradio Interface --- with gr.Blocks(theme=gr.themes.Soft()) as demo: gr.Markdown("# 智慧影像分析工具 (YOLO + MLLM)") gr.Markdown("上傳圖片與YOLO模型進行物件偵測，並可選用MLLM進行進階圖像理解。 ver.250824.1") # mongo_uri = os.getenv('mongo_uri') # gr.Markdown(mongo_uri) with gr.Row(): with gr.Column(scale=1): # 輸入元件 image_input = gr.File(label="上傳圖片", file_count="multiple", file_types=["image"]) #model_input = gr.File(label="上傳YOLO模型 (.pt)", file_count="multiple", file_types=[".pt"], info="若不提供，將使用預設的 yolov8n.pt 模型。") model_input = gr.File(label="上傳YOLO模型 (.pt)", file_count="multiple", file_types=[".pt"]) with gr.Accordion("進階設定", open=False): conf_slider = gr.Slider(minimum=0.1, maximum=1, value=0.40, step=0.05, label="信賴度閾值") mllm_enabled_checkbox = gr.Checkbox(label="開啟MLLM辨識", value=False) mllm_prompt_input = gr.Textbox(label="自訂 MLLM Prompt (選填)", placeholder="例如：請描述圖中人物的穿著與場景。", visible=False) run_button = gr.Button("開始辨識", variant="primary") with gr.Column(scale=2): # 輸出元件 output_gallery = gr.Gallery(label="辨識結果預覽", height=500, object_fit="contain", allow_preview=True) output_text = gr.Textbox(label="詳細辨識資訊", lines=15, placeholder="辨識完成後，所有結果將顯示於此。") output_status = gr.Textbox(label="執行狀態", interactive=False) download_button = gr.File(label="下載所有結果 (.zip)", file_count="single", visible=False) # --- 事件綁定 --- # 點擊 "開始辨識" 按鈕 run_button.click( fn=gradio_multi_model_detection, inputs=[image_input, model_input, conf_slider, mllm_enabled_checkbox, mllm_prompt_input], outputs=[output_gallery, output_status, download_button, output_text] ) # 勾選/取消 "開啟MLLM辨識" mllm_enabled_checkbox.change( fn=toggle_mllm_prompt, inputs=mllm_enabled_checkbox, outputs=mllm_prompt_input ) # 啟動 Gradio 應用 if __name__ == "__main__": demo.launch(debug=True) #demo.launch(share=True)