Update app.py

app.py

@@ -168,39 +168,10 @@ 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
-
-
-# async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.5):
+# 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()
@@ -213,28 +184,57 @@ async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.55):
 #         nms_boxes = non_max_suppression(boxes, iou_threshold)
         
 #         for box, confidence in nms_boxes:
-#             x1, y1, x2, y2 = [int(coord) for coord in box]
+#             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]))
     
-#     # 應用過濾器來移除可能的錯誤檢測
-#     dogs = filter_detections(dogs, (image.width, image.height))
-    
 #     return dogs
 
-# def filter_detections(dogs, image_size):
-#     filtered_dogs = []
-#     image_area = image_size[0] * image_size[1]
+
+async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.5):
+    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 = [int(coord) for coord in box]
+            cropped_image = image.crop((x1, y1, x2, y2))
+            dogs.append((cropped_image, confidence, [x1, y1, x2, y2]))
+    
+    # 應用過濾器來移除可能的錯誤檢測
+    dogs = filter_detections(dogs, (image.width, image.height))
+    
+    return dogs
+
+def filter_detections(dogs, image_size):
+    filtered_dogs = []
+    image_area = image_size[0] * image_size[1]
     
-#     for dog in dogs:
-#         _, _, box = dog
-#         dog_area = (box[2] - box[0]) * (box[3] - box[1])
-#         area_ratio = dog_area / image_area
+    for dog in dogs:
+        _, _, box = dog
+        dog_area = (box[2] - box[0]) * (box[3] - box[1])
+        area_ratio = dog_area / image_area
         
-#         if 0.01 < area_ratio < 0.9:  # 過濾掉太小或太大的檢測框
-#             filtered_dogs.append(dog)
+        if 0.01 < area_ratio < 0.9:  # 過濾掉太小或太大的檢測框
+            filtered_dogs.append(dog)
     
-#     return filtered_dogs
+    return filtered_dogs
 
 
 def non_max_suppression(boxes, iou_threshold):