ui interface with llm and user define prompt

main.py

@@ -1,285 +1,262 @@
-# Required packages:
-# - gradio
-# - transformers
-# - opencv-python
-# - ultralytics
-# - Pillow # Pillow is a common dependency for image processing libraries like OpenCV and Gradio
+# -*- coding: utf-8 -*-
+"""
+系統需求:
+- gradio: 用於建立 Web UI
+- opencv-python: 用於圖片處理
+- ultralytics: YOLOv8 官方函式庫
+- Pillow: 圖片處理基礎庫
+- transformers: (可選，若YOLO模型需要)
+"""
 
 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
-import gemini_ai  as genai  # Assuming gemini_ai is a custom module for Gemini API interactions
+import shutil
+import zipfile
+import uuid  # 匯入 uuid 以生成唯一的執行 ID
+from pathlib import Path # 匯入 Path 以更方便地操作路徑
+import gemini_ai as genai
 
-def multi_model_detection(image_paths_list: list, model_paths_list: list, output_dir: str = 'detection_results', conf_threshold: float = 0.40):
+def create_zip_archive(files, zip_filename):
     """
-    使用多個 YOLO 模型對多張圖片進行物件辨識，
-    並將結果繪製在圖片上，同時保存辨識資訊到文字檔案。
+    將一系列檔案壓縮成一個 zip 檔案。
 
     Args:
-        image_paths_list (list): 包含所有待辨識圖片路徑的列表。
-        model_paths_list (list): 包含所有模型 (.pt 檔案) 路徑的列表。
-        output_dir (str): 儲存結果圖片和文字檔案的目錄。
-                         如果不存在，函式會自動創建。
-        conf_threshold (float): 置信度閾值，只有高於此值的偵測結果會被標示。
+        files (list): 要壓縮的檔案路徑列表。
+        zip_filename (str): 產生的 zip 檔案路徑。
 
     Returns:
-        list: A list of paths to the annotated images.
-        list: A list of paths to the text files with detection information.
+        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
 
-    # 確保輸出目錄存在
-    if not os.path.exists(output_dir):
-        os.makedirs(output_dir)
-        print(f"已創建輸出目錄: {output_dir}")
+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 介面元件的字典。
+    """
+    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 = []
-    print("\n--- 載入模型 ---")
-    # If no models are uploaded, use the default yolov8n.pt
-    if not model_paths_list:
-        default_model_path = 'yoloe-11s-seg-pf.pt' #'yolov8n.pt'
+    if not model_paths:
+        # 如果沒有上傳模型，使用預設模型
+        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 [], []
+            yield {output_status: gr.update(value=f"錯誤: 無法載入預設模型 '{default_model_path}' - {e}")}
+            return
     else:
-        for model_path in model_paths_list:
+        for model_path in model_paths:
             try:
                 model = YOLO(model_path)
-                loaded_models.append((model_path, model)) # 儲存模型路徑和模型物件
-                print(f"成功載入模型: {model_path}")
+                loaded_models.append((model_path, model))
             except Exception as e:
-                print(f"錯誤: 無法載入模型 '{model_path}' - {e}")
-                continue # 如果模型載入失敗，跳過它
-
+                print(f"警告: 無法載入模型 '{model_path}' - {e}，將跳過此模型。")
+                continue
 
     if not loaded_models:
-        print("沒有模型成功載入，請檢查模型路徑或預設模型。")
-        return [], []
+        yield {output_status: gr.update(value="錯誤: 沒有任何模型成功載入。")}
+        return
 
+    # --- 3. 逐一處理圖片 ---
+    total_images = len(image_paths)
     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)
+    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}'，跳過。")
+            print(f"警告: 無法讀取圖片 '{image_path}'，跳過。")
             continue
-
-        # 複製圖片用於繪製，避免修改原始圖片
-        # 使用 NumPy 複製，而不是直接賦值
+        
         annotated_image = original_image.copy()
+        image_base_name = image_path.stem
 
-        # 準備寫入文字檔的內容
-        txt_output_content = []
-        txt_output_content.append(f"檔案: {os.path.basename(image_path)}\n")
-
-        # 對每張圖片使用所有模型進行辨識
-        all_detections_for_image = [] # 儲存所有模型在當前圖片上的偵測結果
-
+        # --- 3a. YOLO 物件偵測 ---
+        yolo_output_content = [f"--- 檔案: {image_path.name} ---"]
+        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} ---")
+            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]
 
-            # Example usage of Gemini API
-            if results.boxes: # 檢查是否有偵測到物件
+            if results.boxes:
                 for box in results.boxes:
-                    # 取得邊界框座標和置信度
                     conf = float(box.conf[0])
-                    if conf >= conf_threshold: # 檢查置信度是否達到閾值
+                    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)
-                        }
+                        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}]")
+                        yolo_output_content.append(f"  - {cls_name} (信賴度: {conf:.2f}) [座標: {x1},{y1},{x2},{y2}]")
             else:
-                txt_output_content.append("  沒有偵測到任何物件。")
-        
-        MLLM_str = genai.analyze_content_with_gemini(image_path) 
-        txt_output_content.append("-MLLM 分析結果為 : " + MLLM_str)
-        
-        # 繪製所有模型在當前圖片上的偵測結果
-        # 我們會根據模型來源給予不同的顏色或樣式，讓結果更容易區分
-
-        # 定義一個顏色循環列表，方便給不同模型分配不同顏色
-        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)] # 確保顏色循環使用
+                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']
-            conf = det['confidence']
-            cls_name = det['class_name']
-            model_name = det['model_name']
-
-            color = color_map.get(model_name, (200, 200, 200)) # 預設灰色
-
-            # 繪製邊界框
+            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)
-
-            # 繪製標籤 (類別名稱 + 置信度 + 模型名稱縮寫)
-            # 為了避免標籤過長，模型名稱只取前幾個字母
-            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):
+            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
+        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))
+
+        # 將本次圖片的結果加入到總列表中
+        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):
     """
-    Gradio 的主要處理函式。
-    接收上傳的檔案和參數，呼叫後端辨識函式，並返回結果。
-
-    Args:
-        image_files (list): Gradio File 元件回傳的圖片檔案列表 (暫存路徑)。
-        model_files (list): Gradio File 元件回傳的模型檔案列表 (暫存路徑)。
-        conf_threshold (float): 置信度閾值。
-        output_subdir (str): 用於儲存本次執行結果的子目錄名稱。
-
-    Returns:
-        tuple: 更新 Gradio 介面所需的多個輸出。
+    根據 Checkbox 狀態，顯示或隱藏 MLLM prompt 輸入框。
     """
-    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
-    )
+    return gr.update(visible=is_enabled)
 
-    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物件辨識+MLLM(demo)")
-    gr.Markdown("上傳您的圖片和模型，並設定置信度閾值進行物件辨識。若未上傳模型，將使用預設模型進行辨識。")
+# --- Gradio Interface ---
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 智慧影像分析工具 (YOLO + MLLM)")
+    gr.Markdown("上傳圖片與YOLO模型進行物件偵測，並可選用MLLM進行進階圖像理解。")
 
     with gr.Row():
-        with gr.Column():
+        with gr.Column(scale=1):
+            # 輸入元件
             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.40, 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
+            #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, output_subdir_input],
-        outputs=[output_gallery, output_text, output_status, download_button]
+        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
     )
 
-# Launch the interface
-demo.launch(debug=True)
+# 啟動 Gradio 應用
+if __name__ == "__main__":
+    demo.launch(debug=True)
+    #demo.launch(share=True)