Update app.py

app.py

@@ -166,57 +166,139 @@ 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):
+#     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
+
+
+# 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
+
+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
+
+        # 交換最大分數和當前索引
+        boxes[i], boxes[maxpos] = boxes[maxpos].copy(), boxes[i].copy()
+        scores[i], scores[maxpos] = scores[maxpos], scores[i]
+
+        # 對剩餘的框應用 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])
+            
+            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
+
+            scores[pos] = weight * scores[pos]
+            
+            if scores[pos] < thresh:
+                scores[pos] = 0
+
+    # 刪除得分為零的框
+    keep = scores > 0
+    boxes = boxes[keep]
+    scores = scores[keep]
+
+    return boxes, scores
+
 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 = []
+    scores = []
+    
     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))
+            boxes.append(xyxy)
+            scores.append(confidence)
     
     if not boxes:
         dogs.append((image, 1.0, [0, 0, image.width, image.height]))
     else:
-        nms_boxes = non_max_suppression(boxes, iou_threshold)
+        boxes = torch.tensor(boxes)
+        scores = torch.tensor(scores)
+        nms_boxes, nms_scores = soft_nms(boxes, scores, thresh=iou_threshold)
         
-        for box, confidence in nms_boxes:
-            x1, y1, x2, y2 = box
+        for box, score in zip(nms_boxes, nms_scores):
+            x1, y1, x2, y2 = box.tolist()
             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]))
+            dogs.append((cropped_image, score.item(), [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)
     if top1_prob < 0.15: