Update app.py

app.py

@@ -250,8 +250,9 @@ def get_akc_breeds_link():
 #     iface.launch()
 
 
+
+# Update the format_description to handle descriptions more cleanly
 def format_description(description, breed):
-    # 分別將不同的屬性分開來顯示，保持結果的可讀性
     if isinstance(description, dict):
         formatted_description = "\n".join([f"**{key}**: {value}" for key, value in description.items()])
     else:
@@ -259,12 +260,9 @@ def format_description(description, breed):
 
     formatted_description = f"""
 **Breed**: {breed}
-
 {formatted_description}
-
 **Want to learn more about dog breeds?**  
 [Visit the AKC dog breeds page]({get_akc_breeds_link()}) and search for {breed} to find detailed information.
-
 *Disclaimer: The external link provided leads to the American Kennel Club (AKC) dog breeds page. 
 You may need to search for the specific breed on that page. 
 I am not responsible for the content on external sites. 
@@ -274,9 +272,9 @@ Please refer to the AKC's terms of use and privacy policy.*
 
 
 async def predict_single_dog(image):
-    # 單一品種直接預測，避免 YOLO
     return await asyncio.to_thread(_predict_single_dog, image)
 
+
 def _predict_single_dog(image):
     image_tensor = preprocess_image(image)
     with torch.no_grad():
@@ -289,22 +287,26 @@ 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):
-    # 檢測多隻狗時使用 YOLO
     return await asyncio.to_thread(_detect_multiple_dogs, image)
 
-def _detect_multiple_dogs(image):
-    results = model_yolo(image)
+
+def _detect_multiple_dogs(image, conf_threshold=0.3):
+    # 調整 YOLO 模型的置信度閾值
+    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()
-                cropped_image = image.crop((xyxy[0], xyxy[1], xyxy[2], xyxy[3]))
-                dogs.append((cropped_image, confidence, xyxy))
+                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 predict(image):
     if image is None:
         return "Please upload an image to start.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
@@ -313,11 +315,9 @@ async def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        # 首先檢查圖片中是否有多隻狗
         dogs = await detect_multiple_dogs(image)
 
         if len(dogs) == 0:
-            # 單狗直接分類
             top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
             if top1_prob < 0.2:
                 return "The image is unclear or the breed is not in the dataset. Please upload a clearer image of a dog.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
@@ -326,19 +326,28 @@ async def predict(image):
             formatted_description = format_description(description, breed)
             return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
+        if len(dogs) == 1:
+            top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
+            breed = topk_breeds[0]
+            description = get_dog_description(breed)
+            formatted_description = format_description(description, breed)
+            return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+
+        # Colors and fonts for multiple dogs
+        color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
         explanations = []
         visible_buttons = []
         annotated_image = image.copy()
         draw = ImageDraw.Draw(annotated_image)
+        font = ImageFont.load_default()
 
-        # 多狗情況
+        # Process each dog detected by YOLO
         for i, (cropped_image, _, box) in enumerate(dogs):
             top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
 
-            # 更改框的顏色以增加辨識度
-            color = ["red", "blue", "green", "yellow"][i % 4]
+            color = color_list[i % len(color_list)]  # Cycle through colors
             draw.rectangle(box, outline=color, width=3)
-            draw.text((box[0], box[1]), f"Dog {i+1}", fill=color)
+            draw.text((box[0], box[1]), f"Dog {i+1}", fill=color, font=font)
 
             if top1_prob >= 0.5:
                 breed = topk_breeds[0]
@@ -352,21 +361,21 @@ Dog {i+1}: Detected with moderate confidence. Here are the top 3 possible breeds
 3. **{topk_breeds[2]}** ({topk_probs_percent[2]})
 """
                 explanations.append(explanation)
-                # 將 Top3 選項顯示在該狗的下方
                 visible_buttons.extend([f"More about Dog {i+1}: {topk_breeds[0]}", f"More about Dog {i+1}: {topk_breeds[1]}", f"More about Dog {i+1}: {topk_breeds[2]}"])
             else:
                 explanations.append(f"Dog {i+1}: The image is unclear or the breed is not in the dataset.")
 
         final_explanation = "\n\n".join(explanations)
-        return final_explanation, annotated_image, gr.update(visible=len(visible_buttons) >= 1, value=visible_buttons[0] if visible_buttons else ""), gr.update(visible=len(visible_buttons) >= 2, value=visible_buttons[1] if len(visible_buttons) >= 2 else ""), gr.update(visible=len(visible_buttons) >= 3, value=visible_buttons[2] if len(visible_buttons) >= 3 else "")
+        return final_explanation, annotated_image, gr.update(visible=True, choices=visible_buttons), gr.update(visible=False), gr.update(visible=False)
 
     except Exception as e:
         return f"An error occurred: {e}", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
+
 async def show_details(choice):
     if not choice:
         return "Please select a breed to view details."
-    
+
     try:
         if "Dog" in choice:
             _, breed = choice.split(": ", 1)
@@ -378,6 +387,7 @@ async def show_details(choice):
         return f"An error occurred while showing details: {e}"
 
 
+
 with gr.Blocks(css="""
     .container { max-width: 900px; margin: auto; padding: 20px; }
     .gr-box { border-radius: 15px; }