Update app.py

app.py

@@ -216,51 +216,33 @@ async def predict_single_dog(image):
 #     iou = intersection / float(area1 + area2 - intersection)
 #     return iou
 
+
 def soft_nms(boxes, scores, sigma=0.5, thresh=0.001, method='gaussian'):
     N = len(boxes)
     
     for i in range(N):
-        maxscore = scores[i]
-        maxpos = i
-
-        tx1, ty1, tx2, ty2 = boxes[i]
-        ts = scores[i]
-
-        pos = i + 1
-        while pos < N:
-            if maxscore < scores[pos]:
-                maxscore = scores[pos]
-                maxpos = pos
-            pos += 1
-
+        maxscore, maxpos = torch.max(scores[i:], dim=0)
+        maxpos = maxpos.item() + i
+        
         # 交換最大分數和當前索引
-        boxes[i], boxes[maxpos] = boxes[maxpos].copy(), boxes[i].copy()
-        scores[i], scores[maxpos] = scores[maxpos], scores[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):
-            tx1, ty1, tx2, ty2 = boxes[pos]
-            ts = scores[pos]
-
-            iou = calculate_iou(boxes[i], boxes[pos])
+            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'
-                if iou > thresh:
-                    weight = 0
-                else:
-                    weight = 1
+                weight = 1 if iou <= thresh else 0
 
-            scores[pos] = weight * scores[pos]
+            scores[pos] *= weight
             
-            if scores[pos] < thresh:
-                scores[pos] = 0
-
-    # 刪除得分為零的框
-    keep = scores > 0
+    # 刪除得分低於閾值的框
+    keep = scores > thresh
     boxes = boxes[keep]
     scores = scores[keep]
 
@@ -274,16 +256,16 @@ async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.55):
     
     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()
+            xyxy = box.xyxy[0]
+            confidence = box.conf
             boxes.append(xyxy)
             scores.append(confidence)
     
     if not boxes:
         dogs.append((image, 1.0, [0, 0, image.width, image.height]))
     else:
-        boxes = torch.tensor(boxes)
-        scores = torch.tensor(scores)
+        boxes = torch.stack(boxes)
+        scores = torch.stack(scores)
         nms_boxes, nms_scores = soft_nms(boxes, scores, thresh=iou_threshold)
         
         for box, score in zip(nms_boxes, nms_scores):