Update app.py

app.py

@@ -243,12 +243,13 @@ 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.2, iou_threshold=0.4, merge_threshold=0.3):
+
+async def detect_multiple_dogs(image, conf_threshold=0.1, iou_threshold=0.3, merge_threshold=0.5):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
     
     image_area = image.width * image.height
-    min_area_ratio = 0.01  # 最小檢測面積佔整個圖像的比例
+    min_area_ratio = 0.005  # 降低最小面積比例以捕獲更多小型狗狗
     
     for box in results.boxes:
         if box.cls == 16:  # COCO 數據集中狗的類別是 16
@@ -256,34 +257,32 @@ async def detect_multiple_dogs(image, conf_threshold=0.2, iou_threshold=0.4, mer
             area = (xyxy[2] - xyxy[0]) * (xyxy[3] - xyxy[1])
             if area / image_area >= min_area_ratio:
                 confidence = box.conf.item()
-                cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
-                dogs.append((cropped_image, confidence, xyxy))
+                dogs.append((xyxy, confidence))
     
-    # 合併重疊的框
-    merged_dogs = []
-    while dogs:
-        base_dog = dogs.pop(0)
-        base_box = torch.tensor(base_dog[2])
-        to_merge = [base_dog]
+    # 使用 NMS 進行後處理
+    if dogs:
+        boxes = torch.tensor([dog[0] for dog in dogs])
+        scores = torch.tensor([dog[1] for dog in dogs])
+        keep = nms(boxes, scores, iou_threshold)
         
-        i = 0
-        while i < len(dogs):
-            compare_box = torch.tensor(dogs[i][2])
-            iou = box_iou(base_box.unsqueeze(0), compare_box.unsqueeze(0)).item()
-            if iou > merge_threshold:
-                to_merge.append(dogs.pop(i))
-            else:
-                i += 1
+        merged_dogs = []
+        for i in keep:
+            xyxy = boxes[i].tolist()
+            confidence = scores[i].item()
+            # 擴大邊界框以包含更多上下文
+            expanded_xyxy = [
+                max(0, xyxy[0] - 20),
+                max(0, xyxy[1] - 20),
+                min(image.width, xyxy[2] + 20),
+                min(image.height, xyxy[3] + 20)
+            ]
+            cropped_image = image.crop(expanded_xyxy)
+            merged_dogs.append((cropped_image, confidence, expanded_xyxy))
         
-        if len(to_merge) == 1:
-            merged_dogs.append(base_dog)
-        else:
-            merged_box = torch.tensor([dog[2] for dog in to_merge]).mean(0)
-            merged_confidence = max(dog[1] for dog in to_merge)
-            merged_image = image.crop(merged_box.tolist())
-            merged_dogs.append((merged_image, merged_confidence, merged_box.tolist()))
+        return merged_dogs
     
-    return merged_dogs if merged_dogs else [(image, 1.0, [0, 0, image.width, image.height])]
+    # 如果沒有檢測到狗狗，返回整張圖片
+    return [(image, 1.0, [0, 0, image.width, image.height])]
 
 async def predict(image):
     if image is None:
@@ -295,10 +294,15 @@ async def predict(image):
 
         dogs = await detect_multiple_dogs(image)
         
-        if len(dogs) == 1:
-            return await process_single_dog(dogs[0][0])
+        # 如果沒有檢測到狗狗或只檢測到一隻，使用整張圖像進行分類
+        if len(dogs) <= 1:
+            top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
+            if top1_prob >= 0.5:
+                return await process_single_dog(image)
+            else:
+                dogs = [(image, 1.0, [0, 0, image.width, image.height])]
 
-        # 多狗情境
+        # 多狗情境處理保持不變
         color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
         explanations = []
         buttons = []