Update app.py

app.py

@@ -243,25 +243,30 @@ def _predict_single_dog(image):
 #         print(error_msg)  # 添加日誌輸出
 #         return error_msg, None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), None
 
-async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4, merge_threshold=0.3):
+async def detect_multiple_dogs(image, conf_threshold=0.1, iou_threshold=0.5, merge_threshold=0.2):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
     boxes = []
     confidences = []
     
+    image_area = image.width * image.height
+    min_area_ratio = 0.005  # 最小檢測面積佔整個圖像的比例
+    
     for box in results.boxes:
         if box.cls == 16:  # COCO 數據集中狗的類別是 16
             xyxy = box.xyxy[0].tolist()
-            confidence = box.conf.item()
-            boxes.append(torch.tensor(xyxy))
-            confidences.append(confidence)
+            area = (xyxy[2] - xyxy[0]) * (xyxy[3] - xyxy[1])
+            if area / image_area >= min_area_ratio:
+                confidence = box.conf.item()
+                boxes.append(torch.tensor(xyxy))
+                confidences.append(confidence)
     
     if boxes:
         boxes = torch.stack(boxes)
         confidences = torch.tensor(confidences)
         
-        # 應用非極大值抑制 (NMS)
-        keep = nms(boxes, confidences, iou_threshold)
+        # 應用軟 NMS
+        keep = soft_nms(boxes, confidences, iou_threshold=iou_threshold, sigma=0.5)
         
         for i in keep:
             xyxy = boxes[i].tolist()
@@ -293,9 +298,32 @@ async def detect_multiple_dogs(image, conf_threshold=0.25, iou_threshold=0.4, me
                 merged_image = image.crop(merged_box.tolist())
                 merged_dogs.append((merged_image, merged_confidence, merged_box.tolist()))
         
+        # 後處理：限制檢測到的狗狗數量
+        if len(merged_dogs) > 5:
+            merged_dogs = sorted(merged_dogs, key=lambda x: x[1], reverse=True)[:5]
+        
         return merged_dogs
     
-    return []
+    # 如果沒有檢測到狗狗，使用備用分類器
+    return await fallback_classifier(image)
+
+async def fallback_classifier(image):
+    # 使用預訓練的 ResNet 或其他適合的分類器
+    transform = transforms.Compose([
+        transforms.Resize((224, 224)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+    ])
+    img_tensor = transform(image).unsqueeze(0)
+    
+    with torch.no_grad():
+        output = fallback_model(img_tensor)
+        confidence, predicted = torch.max(output, 1)
+    
+    if confidence.item() > 0.5:  # 設置一個合適的閾值
+        return [(image, confidence.item(), [0, 0, image.width, image.height])]
+    else:
+        return []
 
 async def predict(image):
     if image is None: