Update app.py

app.py

@@ -167,22 +167,63 @@ async def predict_single_dog(image):
     return top1_prob, topk_breeds, topk_probs_percent
     
 
-async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.55):
+# async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.55):
+#     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()
+#             boxes.append((xyxy, confidence))
+    
+#     if not boxes:
+#         dogs.append((image, 1.0, [0, 0, image.width, image.height]))
+#     else:
+#         nms_boxes = non_max_suppression(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.05)
+#             y1 = max(0, y1 - h * 0.05)
+#             x2 = min(image.width, x2 + w * 0.05)
+#             y2 = min(image.height, y2 + h * 0.05)
+#             cropped_image = image.crop((x1, y1, x2, y2))
+#             dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))
+    
+#     return dogs
+
+async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.55, min_area_ratio=0.01, overlap_threshold=0.8):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
     boxes = []
+    image_area = image.width * image.height
+    
     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()
+            
+            # 計算檢測框面積
+            box_area = (xyxy[2] - xyxy[0]) * (xyxy[3] - xyxy[1])
+            
+            # 如果檢測框面積小於圖片面積的一定比例，則忽略
+            if box_area / image_area < min_area_ratio:
+                continue
+            
             boxes.append((xyxy, confidence))
     
     if not boxes:
         dogs.append((image, 1.0, [0, 0, image.width, image.height]))
     else:
+        # 使用非最大抑制
         nms_boxes = non_max_suppression(boxes, iou_threshold)
         
-        for box, confidence in nms_boxes:
+        # 進一步合併重疊嚴重的框
+        merged_boxes = merge_overlapping_boxes(nms_boxes, overlap_threshold)
+        
+        for box, confidence in merged_boxes:
             x1, y1, x2, y2 = box
             w, h = x2 - x1, y2 - y1
             x1 = max(0, x1 - w * 0.05)
@@ -194,6 +235,30 @@ async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.55):
     
     return dogs
 
+def merge_overlapping_boxes(boxes, overlap_threshold):
+    merged = []
+    while boxes:
+        base_box = boxes.pop(0)
+        i = 0
+        while i < len(boxes):
+            if calculate_iou(base_box[0], boxes[i][0]) > overlap_threshold:
+                # 合併框，取較大的置信度
+                merged_box = merge_boxes(base_box[0], boxes[i][0])
+                merged_conf = max(base_box[1], boxes[i][1])
+                base_box = (merged_box, merged_conf)
+                boxes.pop(i)
+            else:
+                i += 1
+        merged.append(base_box)
+    return merged
+
+def merge_boxes(box1, box2):
+    x1 = min(box1[0], box2[0])
+    y1 = min(box1[1], box2[1])
+    x2 = max(box1[2], box2[2])
+    y2 = max(box1[3], box2[3])
+    return [x1, y1, x2, y2]
+
 
 
 def non_max_suppression(boxes, iou_threshold):