Update app.py

app.py

@@ -20,7 +20,7 @@ logger = logging.getLogger(__name__)
 
 
 # 下載YOLOv8預訓練模型
-model_yolo = YOLO('yolov8m.pt')  # 使用 YOLOv8 預訓練模型
+model_yolo = YOLO('yolov8s.pt')  # 使用 YOLOv8 預訓練模型
 
 
 dog_breeds = ["Afghan_Hound", "African_Hunting_Dog", "Airedale", "American_Staffordshire_Terrier", 
@@ -167,56 +167,34 @@ async def predict_single_dog(image):
     return top1_prob, topk_breeds, topk_probs_percent
     
 
-# async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.6):
-#     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.25, iou_threshold=0.6):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
     boxes = []
     for box in results.boxes:
-        if box.cls == 16:  # 狗類別
+        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)