Update app.py

app.py

@@ -6,7 +6,7 @@ import gradio as gr
 from torchvision.models import efficientnet_v2_m, EfficientNet_V2_M_Weights
 import torch.nn.functional as F
 from torchvision import transforms
-from PIL import Image, ImageDraw, ImageFont
+from PIL import Image, ImageDraw, ImageFont, ImageFilter
 from data_manager import get_dog_description
 from urllib.parse import quote
 from ultralytics import YOLO
@@ -312,7 +312,7 @@ def _predict_single_dog(image):
 #     return dogs
 # 此為如果後面調不好 使用的版本
 
-async def detect_multiple_dogs(image, conf_threshold=0.2, iou_threshold=0.3):
+async def detect_multiple_dogs(image, conf_threshold=0.1, iou_threshold=0.3):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
     for box in results.boxes:
@@ -321,7 +321,7 @@ async def detect_multiple_dogs(image, conf_threshold=0.2, iou_threshold=0.3):
             confidence = box.conf.item()
             area = (xyxy[2] - xyxy[0]) * (xyxy[3] - xyxy[1])
             image_area = image.width * image.height
-            if area > 0.01 * image_area:  # 過濾掉太小的檢測框，但使用相對面積
+            if area > 0.005 * image_area:  # 降低面積閾值以檢測更多狗
                 cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
                 dogs.append((cropped_image, confidence, xyxy))
     
@@ -334,7 +334,7 @@ async def detect_multiple_dogs(image, conf_threshold=0.2, iou_threshold=0.3):
                 confidence = box.conf.item()
                 area = (xyxy[2] - xyxy[0]) * (xyxy[3] - xyxy[1])
                 image_area = image.width * image.height
-                if area > 0.01 * image_area and not is_box_duplicate(xyxy, [d[2] for d in dogs]):
+                if area > 0.005 * image_area and not is_box_duplicate(xyxy, [d[2] for d in dogs]):
                     cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
                     dogs.append((cropped_image, confidence, xyxy))
     
@@ -347,7 +347,6 @@ def is_box_duplicate(new_box, existing_boxes, iou_threshold=0.5):
     return False
 
 def calculate_iou(box1, box2):
-    # 計算兩個邊界框的交集面積
     x1 = max(box1[0], box2[0])
     y1 = max(box1[1], box2[1])
     x2 = min(box1[2], box2[2])
@@ -473,7 +472,6 @@ def calculate_iou(box1, box2):
 # if __name__ == "__main__":
 #     iface.launch()
 
-
 async def predict(image):
     if image is None:
         return "Please upload an image to start.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), None
@@ -482,65 +480,15 @@ async def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        dogs = await detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.5)
-          
-        if len(dogs) == 0:
+        # 首先使用YOLO檢測是否有多隻狗
+        dogs = await detect_multiple_dogs(image)
+        
+        if len(dogs) <= 1:
+            # 單狗情境或沒有檢測到狗，使用整張圖片進行預測
             return await process_single_dog(image)
-        elif len(dogs) == 1:
-            # 如果只檢測到一隻狗，但圖像可能包含多隻狗，再次嘗試檢測
-            if has_multiple_dogs(image):
-                dogs = await detect_multiple_dogs(image, conf_threshold=0.1, iou_threshold=0.2)
-            if len(dogs) == 1:
-                return await process_single_dog(dogs[0][0])
         else:
-             # 多狗情境
-            color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
-            explanations = []
-            buttons = []
-            annotated_image = image.copy()
-            draw = ImageDraw.Draw(annotated_image)
-            font = ImageFont.load_default()
-
-            for i, (cropped_image, confidence, box) in enumerate(dogs):
-                top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
-                color = color_list[i % len(color_list)]
-                draw.rectangle(box, outline=color, width=3)
-                draw.text((box[0], box[1]), f"Dog {i+1}", fill=color, font=font)
-
-                breed = topk_breeds[0]
-                if top1_prob >= 0.5:
-                    description = get_dog_description(breed)
-                    formatted_description = format_description(description, breed)
-                    explanations.append(f"Dog {i+1}: {formatted_description}")
-                elif top1_prob >= 0.2:
-                    dog_explanation = f"Dog {i+1}: Top 3 possible breeds:\n"
-                    dog_explanation += "\n".join([f"{j+1}. **{breed}** ({prob} confidence)" for j, (breed, prob) in enumerate(zip(topk_breeds[:3], topk_probs_percent[:3]))])
-                    explanations.append(dog_explanation)
-                    buttons.extend([gr.update(visible=True, value=f"Dog {i+1}: More about {breed}") for breed in topk_breeds[:3]])
-                else:
-                    explanations.append(f"Dog {i+1}: The image is unclear or the breed is not in the dataset.")
-
-            final_explanation = "\n\n".join(explanations)
-            if buttons:
-                final_explanation += "\n\nClick on a button to view more information about the breed."
-                initial_state = {
-                    "explanation": final_explanation,
-                    "buttons": buttons,
-                    "show_back": True
-                }
-                return (final_explanation, annotated_image, 
-                        buttons[0] if len(buttons) > 0 else gr.update(visible=False),
-                        buttons[1] if len(buttons) > 1 else gr.update(visible=False),
-                        buttons[2] if len(buttons) > 2 else gr.update(visible=False),
-                        gr.update(visible=True),
-                        initial_state)
-            else:
-                initial_state = {
-                    "explanation": final_explanation,
-                    "buttons": [],
-                    "show_back": False
-                }
-                return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
+            # 多狗情境
+            return await process_multiple_dogs(image, dogs)
 
     except Exception as e:
         error_msg = f"An error occurred: {str(e)}"
@@ -548,23 +496,14 @@ async def predict(image):
         return error_msg, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), None
 
 def has_multiple_dogs(image):
-    # 使用簡單的啟發式方法來檢查圖像是否可能包含多隻狗
-    # 這裡可以使用更複雜的方法，如特徵提取或輕量級模型
     gray = image.convert('L')
     edges = gray.filter(ImageFilter.FIND_EDGES)
     edge_pixels = np.array(edges)
-    return np.sum(edge_pixels > 128) > image.width * image.height * 0.1  # 假設邊緣像素比例大於 10% 表示可能有多隻狗
+    return np.sum(edge_pixels > 128) > image.width * image.height * 0.1
 
 async def process_single_dog(image):
     top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
-    if top1_prob < 0.2:
-        initial_state = {
-            "explanation": "The image is unclear or the breed is not in the dataset. Please upload a clearer image of a dog.",
-            "buttons": [],
-            "show_back": False
-        }
-        return initial_state["explanation"], None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
-
+    
     breed = topk_breeds[0]
     description = get_dog_description(breed)
 
@@ -576,7 +515,7 @@ async def process_single_dog(image):
             "show_back": False
         }
         return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
-    else:
+    elif top1_prob >= 0.2:
         explanation = (
             f"The model couldn't confidently identify the breed. Here are the top 3 possible breeds:\n\n"
             f"1. **{topk_breeds[0]}** ({topk_probs_percent[0]} confidence)\n"
@@ -595,20 +534,74 @@ async def process_single_dog(image):
             "show_back": True
         }
         return explanation, image, buttons[0], buttons[1], buttons[2], gr.update(visible=True), initial_state
+    else:
+        initial_state = {
+            "explanation": "The image is unclear or the breed is not in the dataset. Please upload a clearer image of a dog.",
+            "buttons": [],
+            "show_back": False
+        }
+        return initial_state["explanation"], None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
+
+async def process_multiple_dogs(image, dogs):
+    color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
+    explanations = []
+    buttons = []
+    annotated_image = image.copy()
+    draw = ImageDraw.Draw(annotated_image)
+    font = ImageFont.load_default()
+
+    for i, (cropped_image, _, box) in enumerate(dogs):
+        top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
+        color = color_list[i % len(color_list)]
+        draw.rectangle(box, outline=color, width=3)
+        draw.text((box[0], box[1]), f"Dog {i+1}", fill=color, font=font)
+
+        breed = topk_breeds[0]
+        if top1_prob >= 0.5:
+            description = get_dog_description(breed)
+            formatted_description = format_description(description, breed)
+            explanations.append(f"Dog {i+1}: {formatted_description}")
+        elif top1_prob >= 0.2:
+            dog_explanation = f"Dog {i+1}: Top 3 possible breeds:\n"
+            dog_explanation += "\n".join([f"{j+1}. **{breed}** ({prob} confidence)" for j, (breed, prob) in enumerate(zip(topk_breeds[:3], topk_probs_percent[:3]))])
+            explanations.append(dog_explanation)
+            buttons.extend([gr.update(visible=True, value=f"Dog {i+1}: More about {breed}") for breed in topk_breeds[:3]])
+        else:
+            explanations.append(f"Dog {i+1}: The image is unclear or the breed is not in the dataset.")
+
+    final_explanation = "\n\n".join(explanations)
+    if buttons:
+        final_explanation += "\n\nClick on a button to view more information about the breed."
+        initial_state = {
+            "explanation": final_explanation,
+            "buttons": buttons,
+            "show_back": True
+        }
+        return (final_explanation, annotated_image, 
+                buttons[0] if len(buttons) > 0 else gr.update(visible=False),
+                buttons[1] if len(buttons) > 1 else gr.update(visible=False),
+                buttons[2] if len(buttons) > 2 else gr.update(visible=False),
+                gr.update(visible=True),
+                initial_state)
+    else:
+        initial_state = {
+            "explanation": final_explanation,
+            "buttons": [],
+            "show_back": False
+        }
+        return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
 
-def show_details(choice, previous_output, initial_state):
+def show_details(choice, previous_output):
     if not choice:
-        return previous_output, gr.update(visible=True), initial_state
+        return previous_output, gr.update(visible=True)
 
     try:
-        breed = choice.split("More about ")[-1]
+        breed = choice.split("More about ")[-1].split(": ")[-1]
         description = get_dog_description(breed)
-        formatted_description = format_description(description, breed)
-        return formatted_description, gr.update(visible=True), initial_state
+        return format_description(description, breed), gr.update(visible=True)
     except Exception as e:
-        error_msg = f"An error occurred while showing details: {e}"
-        print(error_msg)  # 添加日誌輸出
-        return error_msg, gr.update(visible=True), initial_state
+        return f"An error occurred while showing details: {e}", gr.update(visible=True)
+
 
 # 介面部分
 with gr.Blocks() as iface: