Update app.py

app.py

@@ -312,7 +312,7 @@ def _predict_single_dog(image):
 #     return dogs
 # 此為如果後面調不好 使用的版本
 
-async def detect_multiple_dogs(image, conf_threshold=0.1, iou_threshold=0.3):
+async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.5):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
     for box in results.boxes:
@@ -320,33 +320,45 @@ async def detect_multiple_dogs(image, conf_threshold=0.1, iou_threshold=0.3):
             xyxy = box.xyxy[0].tolist()
             confidence = box.conf.item()
             area = (xyxy[2] - xyxy[0]) * (xyxy[3] - xyxy[1])
-            image_area = image.width * image.height
-            if area > 0.005 * image_area:  # 降低面積閾值以檢測更多狗
+            if area > 1000:  # 過濾掉太小的檢測框
                 cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
                 dogs.append((cropped_image, confidence, xyxy))
     
-    # 如果檢測到的狗太少，嘗試降低閾值再次檢測
-    if len(dogs) < 2:
-        results = model_yolo(image, conf=conf_threshold/2, iou=iou_threshold)[0]
-        for box in results.boxes:
-            if box.cls == 16:
-                xyxy = box.xyxy[0].tolist()
-                confidence = box.conf.item()
-                area = (xyxy[2] - xyxy[0]) * (xyxy[3] - xyxy[1])
-                image_area = image.width * image.height
-                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))
+    # 合併重疊的檢測框
+    dogs = merge_overlapping_boxes(dogs, iou_threshold=0.6)
     
     return dogs
 
-def is_box_duplicate(new_box, existing_boxes, iou_threshold=0.5):
-    for box in existing_boxes:
-        if calculate_iou(new_box, box) > iou_threshold:
-            return True
-    return False
+def merge_overlapping_boxes(dogs, iou_threshold=0.6):
+    merged_dogs = []
+    while dogs:
+        base = dogs.pop(0)
+        i = 0
+        while i < len(dogs):
+            if calculate_iou(base[2], dogs[i][2]) > iou_threshold:
+                # 合併重疊的框
+                base = merge_boxes(base, dogs.pop(i))
+            else:
+                i += 1
+        merged_dogs.append(base)
+    return merged_dogs
+
+def merge_boxes(box1, box2):
+    xyxy1, conf1, _ = box1
+    xyxy2, conf2, _ = box2
+    merged_xyxy = [
+        min(xyxy1[0], xyxy2[0]),
+        min(xyxy1[1], xyxy2[1]),
+        max(xyxy1[2], xyxy2[2]),
+        max(xyxy1[3], xyxy2[3])
+    ]
+    merged_conf = max(conf1, conf2)
+    merged_image = Image.new('RGB', (int(merged_xyxy[2] - merged_xyxy[0]), int(merged_xyxy[3] - merged_xyxy[1])))
+    merged_image.paste(box1[0], (0, 0))
+    return (merged_image, merged_conf, merged_xyxy)
 
 def calculate_iou(box1, box2):
+    # 計算兩個邊界框的交集面積
     x1 = max(box1[0], box2[0])
     y1 = max(box1[1], box2[1])
     x2 = min(box1[2], box2[2])
@@ -474,48 +486,72 @@ def calculate_iou(box1, box2):
 
 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
+        return "Please upload an image to start.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
     try:
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        # 首先使用YOLO檢測是否有多隻狗
-        dogs = await detect_multiple_dogs(image)
+        # 嘗試檢測多隻狗，進一步降低閾值以提高檢測率
+        dogs = await detect_multiple_dogs(image, conf_threshold=0.05)  # 降低閾值以檢測更多狗
         
         if len(dogs) <= 1:
-            # 單狗情境或沒有檢測到狗，使用整張圖片進行預測
+            # 單狗情境
             return await process_single_dog(image)
+
+        # 多狗情境
+        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."
+            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))  # 顯示 back 按鈕
         else:
-            # 多狗情境
-            return await process_multiple_dogs(image, dogs)
+            return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
     except Exception as e:
-        error_msg = f"An error occurred: {str(e)}"
-        print(error_msg)  # 添加日誌輸出
-        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
+        return f"An error occurred: {str(e)}", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
 async def process_single_dog(image):
     top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
-    
+    if top1_prob < 0.2:
+        return "The image is unclear or the breed is not in the dataset. Please upload a clearer image of a dog.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+
     breed = topk_breeds[0]
     description = get_dog_description(breed)
 
     if top1_prob >= 0.5:
         formatted_description = format_description(description, breed)
-        initial_state = {
-            "explanation": formatted_description,
-            "buttons": [],
-            "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
-    elif top1_prob >= 0.2:
+        return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+    else:
         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"
@@ -523,80 +559,18 @@ async def process_single_dog(image):
             f"3. **{topk_breeds[2]}** ({topk_probs_percent[2]} confidence)\n\n"
             "Click on a button to view more information about the breed."
         )
-        buttons = [
-            gr.update(visible=True, value=f"More about {topk_breeds[0]}"),
-            gr.update(visible=True, value=f"More about {topk_breeds[1]}"),
-            gr.update(visible=True, value=f"More about {topk_breeds[2]}")
-        ]
-        initial_state = {
-            "explanation": explanation,
-            "buttons": buttons,
-            "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
+        return (explanation, image, 
+                gr.update(visible=True, value=f"More about {topk_breeds[0]}"),
+                gr.update(visible=True, value=f"More about {topk_breeds[1]}"),
+                gr.update(visible=True, value=f"More about {topk_breeds[2]}"),
+                gr.update(visible=True))  # 顯示 back 按鈕
 
 def show_details(choice, previous_output):
     if not choice:
         return previous_output, gr.update(visible=True)
 
     try:
-        breed = choice.split("More about ")[-1].split(": ")[-1]
+        breed = choice.split("More about ")[-1]
         description = get_dog_description(breed)
         return format_description(description, breed), gr.update(visible=True)
     except Exception as e: