Update app.py

app.py

@@ -185,63 +185,10 @@ async def predict_single_dog(image):
         return probabilities[0], breeds[:3], relative_probs
 
 
-# async def detect_multiple_dogs(image, conf_threshold=0.3, 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
-
-
-async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.55, sigma=0.5):
+async def detect_multiple_dogs(image, conf_threshold=0.3, 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()
@@ -251,69 +198,31 @@ async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.55, s
     if not boxes:
         dogs.append((image, 1.0, [0, 0, image.width, image.height]))
     else:
-        # 使用SoftNMS替代原有的NMS
-        nms_boxes = soft_nms(boxes, iou_threshold, sigma)
+        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.1)  # 增加到10%的margin
-            y1 = max(0, y1 - h * 0.1)
-            x2 = min(image.width, x2 + w * 0.1)
-            y2 = min(image.height, y2 + h * 0.1)
+            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 soft_nms(boxes, iou_threshold=0.55, sigma=0.5, score_threshold=0.25):
-    """
-    SoftNMS with Gaussian decay
-    """
-    if not boxes:
-        return []
-    
-    # 轉換格式以便處理
-    box_coords = np.array([box[0] for box in boxes])
-    scores = np.array([box[1] for box in boxes])
-    
-    # 按照confidence排序
-    indices = np.argsort(scores)[::-1]
-    box_coords = box_coords[indices]
-    scores = scores[indices]
-    
-    keep_boxes = []
-    keep_scores = []
-    
-    while len(scores) > 0:
-        # 保留最高分數的框
-        keep_boxes.append(box_coords[0].tolist())
-        keep_scores.append(scores[0])
-        
-        if len(scores) == 1:
-            break
-        
-        # 計算當前最高分框與其他所有框的IoU
-        ious = np.array([calculate_iou(box_coords[0], box) for box in box_coords[1:]])
-        
-        # 使用高斯衰減更新分數
-        scores[1:] = scores[1:] * np.exp(-(ious * ious) / sigma)
-        
-        # 移除最高分的框並過濾低於閾值的框
-        box_coords = box_coords[1:]
-        scores = scores[1:]
-        mask = scores > score_threshold
-        box_coords = box_coords[mask]
-        scores = scores[mask]
-    
-    return list(zip(keep_boxes, keep_scores))
+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):
-    """
-    IoU 計算
-    """
     x1 = max(box1[0], box2[0])
     y1 = max(box1[1], box2[1])
     x2 = min(box1[2], box2[2])
@@ -327,7 +236,6 @@ def calculate_iou(box1, box2):
     return iou
 
 
-
 def create_breed_comparison(breed1: str, breed2: str) -> dict:
     breed1_info = get_dog_description(breed1)
     breed2_info = get_dog_description(breed2)