Update app.py

app.py

@@ -193,17 +193,62 @@ async def predict_single_dog(image):
         topk_probs_percent = [f"{prob.item() * 100:.2f}%" for prob in topk_probs[0]]
     return top1_prob, topk_breeds, topk_probs_percent
 
-async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.45):
+async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.5):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
+    boxes = []
     for box in results.boxes:
         if box.cls == 16:  # COCO dataset class for dog is 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))
+            boxes.append(xyxy)
+
+    # 如果沒有檢測到狗，使用整張圖片
+    if not boxes:
+        dogs.append((image, 1.0, [0, 0, image.width, image.height]))
+    else:
+        # 合併重疊的框
+        merged_boxes = merge_boxes(boxes)
+        for box in merged_boxes:
+            cropped_image = image.crop((box[0], box[1], box[2], box[3]))
+            dogs.append((cropped_image, 1.0, box))
+
     return dogs
 
+def merge_boxes(boxes, iou_threshold=0.5):
+    merged = []
+    while boxes:
+        base_box = boxes.pop(0)
+        i = 0
+        while i < len(boxes):
+            if calculate_iou(base_box, boxes[i]) > iou_threshold:
+                base_box = merge_two_boxes(base_box, boxes.pop(i))
+            else:
+                i += 1
+        merged.append(base_box)
+    return merged
+
+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
+
+def merge_two_boxes(box1, box2):
+    return [
+        min(box1[0], box2[0]),
+        min(box1[1], box2[1]),
+        max(box1[2], box2[2]),
+        max(box1[3], box2[3])
+    ]
+
 
 async def process_single_dog(image):
     top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
@@ -412,9 +457,6 @@ async def predict(image):
 
         dogs = await detect_multiple_dogs(image)
         
-        if len(dogs) == 0:
-            dogs = [(image, 1.0, [0, 0, image.width, image.height])]
-        
         color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
         explanations = []
         buttons = []
@@ -452,7 +494,7 @@ async def predict(image):
                 "is_multi_dog": len(dogs) > 1,
                 "dogs_info": explanations
             }
-            return final_explanation, annotated_image, gr.update(visible=True, choices=buttons), initial_state
+            return final_explanation, annotated_image, gr.update(visible=true, choices=buttons), initial_state
         else:
             initial_state = {
                 "explanation": final_explanation,
@@ -469,6 +511,7 @@ async def predict(image):
         print(error_msg)
         return error_msg, None, gr.update(visible=False, choices=[]), None
 
+
 def show_details(choice, previous_output, initial_state):
     if not choice:
         return previous_output, gr.update(visible=True), initial_state