Update app.py

app.py

@@ -300,17 +300,72 @@ def _predict_single_dog(image):
 #                     dogs.append((cropped_image, confidence, xyxy))
 #     return dogs
 
-async def detect_multiple_dogs(image, conf_threshold=0.2, iou_threshold=0.5):
+# async def detect_multiple_dogs(image, conf_threshold=0.2, iou_threshold=0.5):
+#     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
+#     dogs = []
+#     for box in results.boxes:
+#         if box.cls == 16:  # COCO 資料集中狗的類別是 16
+#             xyxy = box.xyxy[0].tolist()
+#             confidence = box.conf.item()
+#             cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
+#             dogs.append((cropped_image, confidence, xyxy))
+#     return dogs
+# 此為如果後面調不好 使用的版本
+
+async def detect_multiple_dogs(image, conf_threshold=0.2, iou_threshold=0.45):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
+    all_boxes = []
+    
+    # 首先收集所有可能的狗的邊界框
     for box in results.boxes:
         if box.cls == 16:  # COCO 資料集中狗的類別是 16
             xyxy = box.xyxy[0].tolist()
             confidence = box.conf.item()
-            cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
-            dogs.append((cropped_image, confidence, xyxy))
+            all_boxes.append((xyxy, confidence))
+    
+    # 按置信度排序
+    all_boxes.sort(key=lambda x: x[1], reverse=True)
+    
+    # 應用非最大抑制
+    for box, confidence in all_boxes:
+        if not is_box_overlapping(box, [d[2] for d in dogs], iou_threshold):
+            cropped_image = image.crop((box[0], box[1], box[2], box[3]))
+            dogs.append((cropped_image, confidence, box))
+    
+    # 如果沒有檢測到狗，嘗試降低閾值再次檢測
+    if len(dogs) == 0:
+        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()
+                cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
+                dogs.append((cropped_image, confidence, xyxy))
+    
     return dogs
 
+def is_box_overlapping(box, existing_boxes, iou_threshold):
+    for existing_box in existing_boxes:
+        if calculate_iou(box, existing_box) > iou_threshold:
+            return True
+    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])
+    y2 = min(box1[3], box2[3])
+    
+    intersection = max(0, x2 - x1) * max(0, y2 - y1)
+    
+    area1 = (box1[2] - box1[0]) * (box1[3] - box1[1])
+    area2 = (box2[2] - box2[0]) * (box2[3] - box2[1])
+    
+    iou = intersection / float(area1 + area2 - intersection)
+    return iou
+
 
 # async def predict(image):
 #     if image is None:
@@ -432,59 +487,63 @@ async def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        dogs = await detect_multiple_dogs(image, conf_threshold=0.05)
+        dogs = await detect_multiple_dogs(image, conf_threshold=0.05, iou_threshold=0.45)
         
-        if len(dogs) <= 1:
+        if len(dogs) == 0:
+            # 沒有檢測到狗，使用原始圖像進行單狗處理
             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."
-            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)
+        elif len(dogs) == 1:
+            # 只檢測到一隻狗，使用裁剪後的圖像進行處理
+            return await process_single_dog(dogs[0][0])
         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
+            # 多狗情境
+            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
 
     except Exception as e:
         error_msg = f"An error occurred: {str(e)}"