Update app.py

app.py

@@ -194,7 +194,7 @@ async def predict_single_dog(image):
     return top1_prob, topk_breeds, topk_probs_percent
     
 
-async def detect_multiple_dogs(image, conf_threshold=0.2, iou_threshold=0.45):
+async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.3):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
     boxes = []
@@ -208,22 +208,47 @@ async def detect_multiple_dogs(image, conf_threshold=0.2, iou_threshold=0.45):
     if not boxes:
         dogs.append((image, 1.0, [0, 0, image.width, image.height]))
     else:
-        # 按置信度排序並選擇所有框
+        # 按置信度排序
         sorted_boxes = sorted(boxes, key=lambda x: x[1], reverse=True)
         
-        for box, confidence in sorted_boxes:
+        # 使用非極大值抑制（NMS）來合併重疊的框
+        nms_boxes = non_max_suppression(sorted_boxes, iou_threshold)
+        
+        for box, confidence in nms_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)
+            x1 = max(0, x1 - w * 0.15)
+            y1 = max(0, y1 - h * 0.15)
+            x2 = min(image.width, x2 + w * 0.15)
+            y2 = min(image.height, y2 + h * 0.15)
             cropped_image = image.crop((x1, y1, x2, y2))
             dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))
 
     return dogs
 
+def non_max_suppression(boxes, iou_threshold):
+    keep = []
+    boxes = sorted(boxes, key=lambda x: x[1], reverse=True)
+    while boxes:
+        current = boxes.pop(0)
+        keep.append(current)
+        boxes = [box for box in boxes if calculate_iou(current[0], box[0]) < iou_threshold]
+    return keep
+
+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 process_single_dog(image):
     top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)