Update app.py

app.py

@@ -195,6 +195,31 @@ async def predict_single_dog(image):
 #     return dogs
 
 
+async def detect_multiple_dogs(image, conf_threshold=0.4, iou_threshold=0.5):
+    # 提高conf_threshold來減少無效框
+    results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
+    dogs = []
+    boxes = []
+    for box in results.boxes:
+        if box.cls == 16:  # 確保是狗的類別
+            xyxy = box.xyxy[0].tolist()
+            confidence = box.conf.item()
+            # 只保存高信心的框，降低不必要框的數量
+            if confidence >= conf_threshold:
+                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
+            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)