Update app.py

app.py

@@ -244,12 +244,12 @@ def _predict_single_dog(image):
 #         return error_msg, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), None
 
 
-async def detect_multiple_dogs(image, conf_threshold=0.1, iou_threshold=0.3, merge_threshold=0.5):
+async def detect_multiple_dogs(image, conf_threshold=0.1, iou_threshold=0.4, merge_threshold=0.7):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
     
     image_area = image.width * image.height
-    min_area_ratio = 0.005  # 降低最小面積比例以捕獲更多小型狗狗
+    min_area_ratio = 0.005  # 最小檢測面積佔整個圖像的比例
     
     for box in results.boxes:
         if box.cls == 16:  # COCO 數據集中狗的類別是 16
@@ -259,17 +259,57 @@ async def detect_multiple_dogs(image, conf_threshold=0.1, iou_threshold=0.3, mer
                 confidence = box.conf.item()
                 dogs.append((xyxy, confidence))
     
-    # 使用 NMS 進行後處理
     if dogs:
         boxes = torch.tensor([dog[0] for dog in dogs])
         scores = torch.tensor([dog[1] for dog in dogs])
+        
+        # 應用 NMS
         keep = nms(boxes, scores, iou_threshold)
         
         merged_dogs = []
         for i in keep:
             xyxy = boxes[i].tolist()
             confidence = scores[i].item()
-            # 擴大邊界框以包含更多上下文
+            merged_dogs.append((xyxy, confidence))
+        
+        # 後處理：分離過於接近的檢測框
+        final_dogs = []
+        while merged_dogs:
+            base_dog = merged_dogs.pop(0)
+            to_merge = [base_dog]
+            
+            i = 0
+            while i < len(merged_dogs):
+                iou = box_iou(torch.tensor([base_dog[0]]), torch.tensor([merged_dogs[i][0]]))[0][0].item()
+                if iou > merge_threshold:
+                    to_merge.append(merged_dogs.pop(i))
+                else:
+                    i += 1
+            
+            if len(to_merge) == 1:
+                final_dogs.append(base_dog)
+            else:
+                # 如果檢測到多個重疊框，嘗試分離它們
+                centers = torch.tensor([[((box[0] + box[2]) / 2, (box[1] + box[3]) / 2)] for box, _ in to_merge])
+                distances = torch.cdist(centers, centers)
+                
+                if torch.any(distances > 0):  # 確保不是完全重疊
+                    max_distance = distances.max()
+                    if max_distance > (base_dog[0][2] - base_dog[0][0]) * 0.5:  # 如果最大距離大於框寬度的一半
+                        final_dogs.extend(to_merge)
+                    else:
+                        # 合併為一個框
+                        merged_box = torch.tensor([box for box, _ in to_merge]).mean(dim=0)
+                        merged_confidence = max(conf for _, conf in to_merge)
+                        final_dogs.append((merged_box.tolist(), merged_confidence))
+                else:
+                    # 完全重疊的情況，保留置信度最高的
+                    best_dog = max(to_merge, key=lambda x: x[1])
+                    final_dogs.append(best_dog)
+        
+        # 擴展邊界框並創建剪裁的圖像
+        expanded_dogs = []
+        for xyxy, confidence in final_dogs:
             expanded_xyxy = [
                 max(0, xyxy[0] - 20),
                 max(0, xyxy[1] - 20),
@@ -277,9 +317,9 @@ async def detect_multiple_dogs(image, conf_threshold=0.1, iou_threshold=0.3, mer
                 min(image.height, xyxy[3] + 20)
             ]
             cropped_image = image.crop(expanded_xyxy)
-            merged_dogs.append((cropped_image, confidence, expanded_xyxy))
+            expanded_dogs.append((cropped_image, confidence, expanded_xyxy))
         
-        return merged_dogs
+        return expanded_dogs
     
     # 如果沒有檢測到狗狗，返回整張圖片
     return [(image, 1.0, [0, 0, image.width, image.height])]