Update app.py

app.py

@@ -217,38 +217,12 @@ async def predict_single_dog(image):
 #     return iou
 
 
-def soft_nms(boxes, scores, sigma=0.5, thresh=0.001, method='gaussian'):
-    N = len(boxes)
-    
-    for i in range(N):
-        maxscore, maxpos = torch.max(scores[i:], dim=0)
-        maxpos = maxpos.item() + i
-        
-        # 交換最大分數和當前索引
-        boxes[i], boxes[maxpos] = boxes[maxpos].clone(), boxes[i].clone()
-        scores[i], scores[maxpos] = scores[maxpos].clone(), scores[i].clone()
-
-        # 對剩餘的框應用 Soft NMS
-        for pos in range(i + 1, N):
-            iou = box_iou(boxes[i].unsqueeze(0), boxes[pos].unsqueeze(0)).squeeze()
-            
-            if method == 'linear':
-                weight = 1 - iou if iou > thresh else 1
-            elif method == 'gaussian':
-                weight = torch.exp(-(iou * iou) / sigma)
-            else:  # 'original'
-                weight = 1 if iou <= thresh else 0
-
-            scores[pos] *= weight
-            
-    # 刪除得分低於閾值的框
-    keep = scores > thresh
-    boxes = boxes[keep.squeeze(-1)]  # 確保 keep 的維度與 boxes 匹配
-    scores = scores[keep]
-
-    return boxes, scores
-
-async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.55):
+def weighted_nms(boxes, scores, iou_threshold=0.5, score_threshold=0.05):
+    keep = nms(boxes, scores, iou_threshold)
+    keep = keep[scores[keep] > score_threshold]
+    return boxes[keep], scores[keep]
+
+async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.45):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
     boxes = []
@@ -265,8 +239,8 @@ async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.55):
         dogs.append((image, 1.0, [0, 0, image.width, image.height]))
     else:
         boxes = torch.stack(boxes)
-        scores = torch.stack(scores).squeeze(-1)  # 確保 scores 是一維張量
-        nms_boxes, nms_scores = soft_nms(boxes, scores, thresh=iou_threshold)
+        scores = torch.stack(scores).squeeze(-1)
+        nms_boxes, nms_scores = weighted_nms(boxes, scores, iou_threshold=iou_threshold, score_threshold=0.1)
         
         for box, score in zip(nms_boxes, nms_scores):
             x1, y1, x2, y2 = box.tolist()