Update app.py

app.py

@@ -8,7 +8,6 @@ from torchvision.ops import nms, box_iou
 import torch.nn.functional as F
 from torchvision import transforms
 from PIL import Image, ImageDraw, ImageFont, ImageFilter
-from sklearn.cluster import KMeans
 from data_manager import get_dog_description
 from urllib.parse import quote
 from ultralytics import YOLO
@@ -168,39 +167,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.4, 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,37 +183,66 @@ async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.5):
         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]
-    num_dogs = len(dogs)
 
-    # 根據檢測到的狗的數量動態調整閾值
-    if num_dogs > 5:
-        min_ratio, max_ratio = 0.003, 0.5
-    elif num_dogs > 2:
-        min_ratio, max_ratio = 0.005, 0.6
-    else:
-        min_ratio, max_ratio = 0.01, 0.7
+# 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]
+#     num_dogs = len(dogs)
+
+#     # 根據檢測到的狗的數量動態調整閾值
+#     if num_dogs > 5:
+#         min_ratio, max_ratio = 0.003, 0.5
+#     elif num_dogs > 2:
+#         min_ratio, max_ratio = 0.005, 0.6
+#     else:
+#         min_ratio, max_ratio = 0.01, 0.7
 
-    for dog in dogs:
-        _, confidence, box = dog
-        dog_area = (box[2] - box[0]) * (box[3] - box[1])
-        area_ratio = dog_area / image_area
+#     for dog in dogs:
+#         _, confidence, box = dog
+#         dog_area = (box[2] - box[0]) * (box[3] - box[1])
+#         area_ratio = dog_area / image_area
         
-        if min_ratio < area_ratio < max_ratio:
-            filtered_dogs.append(dog)
+#         if min_ratio < area_ratio < max_ratio:
+#             filtered_dogs.append(dog)
     
-    return filtered_dogs
+#     return filtered_dogs
 
 
 def non_max_suppression(boxes, iou_threshold):