Update app.py

app.py

@@ -283,21 +283,32 @@ def _predict_single_dog(image):
         topk_probs_percent = [f"{prob.item() * 100:.2f}%" for prob in topk_probs[0]]
     return top1_prob, topk_breeds, topk_probs_percent
 
-async def detect_multiple_dogs(image, conf_threshold=0.3):
-    # 調整 YOLO 模型的置信度閾值
-    return await asyncio.to_thread(_detect_multiple_dogs, image, conf_threshold)
-
-def _detect_multiple_dogs(image, conf_threshold):
-    results = model_yolo(image, conf=conf_threshold)
+# async def detect_multiple_dogs(image, conf_threshold=0.3):
+#     # 調整 YOLO 模型的置信度閾值
+#     return await asyncio.to_thread(_detect_multiple_dogs, image, conf_threshold)
+
+# def _detect_multiple_dogs(image, conf_threshold):
+#     results = model_yolo(image, conf=conf_threshold)
+#     dogs = []
+#     for result in results:
+#         for box in result.boxes:
+#             if box.cls == 16:  # COCO 資料集中狗的類別是 16
+#                 xyxy = box.xyxy[0].tolist()
+#                 confidence = box.conf.item()
+#                 if confidence >= conf_threshold:  # 只保留置信度高於閾值的框
+#                     cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
+#                     dogs.append((cropped_image, confidence, xyxy))
+#     return dogs
+
+async def detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.5):
+    results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
     dogs = []
-    for result in results:
-        for box in result.boxes:
-            if box.cls == 16:  # COCO 資料集中狗的類別是 16
-                xyxy = box.xyxy[0].tolist()
-                confidence = box.conf.item()
-                if confidence >= conf_threshold:  # 只保留置信度高於閾值的框
-                    cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
-                    dogs.append((cropped_image, confidence, xyxy))
+    for box in results.boxes:
+        if box.cls == 16:  # COCO 資料集中狗的類別是 16
+            xyxy = box.xyxy[0].tolist()
+            confidence = box.conf.item()
+            cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
+            dogs.append((cropped_image, confidence, xyxy))
     return dogs
 
 
@@ -421,10 +432,11 @@ async def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        # 嘗試檢測多隻狗，進一步降低閾值以提高檢測率
-        dogs = await detect_multiple_dogs(image, conf_threshold=0.1)
-        if len(dogs) == 0:
-            # 單狗情境
+        # 嘗試檢測多隻狗，使用較高的閾值以避免錯誤檢測
+        dogs = await detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.5)
+        
+        # 如果只檢測到一隻狗，使用單狗處理邏輯
+        if len(dogs) <= 1:
             return await process_single_dog(image)
 
         # 多狗情境
@@ -461,9 +473,9 @@ async def predict(image):
                     buttons[0] if len(buttons) > 0 else gr.update(visible=False),
                     buttons[1] if len(buttons) > 1 else gr.update(visible=False),
                     buttons[2] if len(buttons) > 2 else gr.update(visible=False),
-                    gr.update(visible=False))
+                    gr.update(visible=True))
         else:
-            return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+            return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=True)
 
     except Exception as e:
         return f"An error occurred: {str(e)}", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
@@ -495,7 +507,7 @@ async def process_single_dog(image):
 
 def show_details(choice, previous_output):
     if not choice:
-        return previous_output, gr.update(visible=False)
+        return previous_output, gr.update(visible=True)
 
     try:
         breed = choice.split("More about ")[-1]