Update app.py

app.py

@@ -193,7 +193,7 @@ 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.25, iou_threshold=0.5):
+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 = []
     boxes = []
@@ -201,17 +201,25 @@ async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.5):
         if box.cls == 16:  # COCO dataset class for dog is 16
             xyxy = box.xyxy[0].tolist()
             confidence = box.conf.item()
-            boxes.append(xyxy)
+            boxes.append((xyxy, confidence))
 
     # 如果沒有檢測到狗，使用整張圖片
     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))
+        # 按置信度排序並選擇前4個框（如果有的話）
+        sorted_boxes = sorted(boxes, key=lambda x: x[1], reverse=True)[:4]
+        
+        for box, confidence in sorted_boxes:
+            x1, y1, x2, y2 = box
+            # 擴大框的大小
+            w, h = x2 - x1, y2 - y1
+            x1 = max(0, x1 - w * 0.1)
+            y1 = max(0, y1 - h * 0.1)
+            x2 = min(image.width, x2 + w * 0.1)
+            y2 = min(image.height, y2 + h * 0.1)
+            cropped_image = image.crop((x1, y1, x2, y2))
+            dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))
 
     return dogs
 
@@ -221,7 +229,7 @@ def merge_boxes(boxes, iou_threshold=0.5):
         base_box = boxes.pop(0)
         i = 0
         while i < len(boxes):
-            if calculate_iou(base_box, boxes[i]) > iou_threshold:
+            if calculate_iou(base_box[0], boxes[i][0]) > iou_threshold:
                 base_box = merge_two_boxes(base_box, boxes.pop(i))
             else:
                 i += 1
@@ -242,12 +250,13 @@ def calculate_iou(box1, box2):
     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])
-    ]
+    return (
+        [min(box1[0][0], box2[0][0]),
+         min(box1[0][1], box2[0][1]),
+         max(box1[0][2], box2[0][2]),
+         max(box1[0][3], box2[0][3])],
+        max(box1[1], box2[1])  # 取較高的置信度
+    )
 
 
 async def process_single_dog(image):
@@ -494,7 +503,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,