Update app.py

app.py

@@ -249,7 +249,6 @@ def get_akc_breeds_link():
 # if __name__ == "__main__":
 #     iface.launch()
 
-# 格式化狗的品種描述函數
 def format_description(description, breed, is_multi_dog=False, dog_number=None):
     if isinstance(description, dict):
         formatted_description = "\n\n".join([f"**{key}**: {value}" for key, value in description.items() if key != "Breed"])
@@ -302,7 +301,6 @@ async def detect_multiple_dogs(image):
         print(f"Error in detect_multiple_dogs: {e}")
         return []
 
-# 主預測函數
 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)
@@ -311,23 +309,34 @@ async def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        # 使用 YOLO 偵測有無多隻狗
-        dogs = await detect_multiple_dogs(image)
+        # First, try to predict without YOLO
+        top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
 
-        if len(dogs) == 1:
-            # 如果只有一隻狗，直接預測品種，不使用 YOLO
-            top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
+        if top1_prob >= 0.5:
             breed = topk_breeds[0]
             description = get_dog_description(breed)
             formatted_description = format_description(description, breed)
-            
-            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)
-            
             return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
         
-        elif len(dogs) > 1:
-            # 如果偵測到多隻狗，使用 YOLO 進行分割並預測
+        elif top1_prob >= 0.2:
+            explanation = (
+                f"The model couldn't confidently identify the breed. Here are the top 3 possible breeds:\n\n"
+                f"1. **{topk_breeds[0]}** ({topk_probs_percent[0]} confidence)\n"
+                f"2. **{topk_breeds[1]}** ({topk_probs_percent[1]} confidence)\n"
+                f"3. **{topk_breeds[2]}** ({topk_probs_percent[2]} confidence)\n\n"
+                "Click on a button to view more information about the breed."
+            )
+            breed_buttons = [f"More about {breed}" for breed in topk_breeds[:3]]
+            return explanation, image, gr.update(visible=True, choices=breed_buttons), gr.update(visible=False), gr.update(visible=False)
+        
+        elif 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)
+
+        # If confidence is low, use YOLO to check for multiple dogs
+        dogs = await detect_multiple_dogs(image)
+
+        if len(dogs) > 1:
+            # Multiple dogs detected, process each one
             explanations = []
             buttons = []
             annotated_image = image.copy()
@@ -343,7 +352,7 @@ async def predict(image):
                 if top1_prob >= 0.5:
                     breed = topk_breeds[0]
                     description = get_dog_description(breed)
-                    explanations.append(f"Dog {i}: **{breed}**\n\n{format_description(description, breed, is_multi_dog=True, dog_number=i)}")
+                    explanations.append(format_description(description, breed, is_multi_dog=True, dog_number=i))
                 else:
                     explanation = f"""
 Dog {i}: Detected with moderate confidence. Here are the top 3 possible breeds:
@@ -363,19 +372,21 @@ Dog {i}: Detected with moderate confidence. Here are the top 3 possible breeds:
                 return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
         else:
-            # 未偵測到狗
-            return "No dogs detected. Please upload a clear image of a dog.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible(False))
+            # No dogs detected or only one dog (already processed)
+            return "No additional dogs detected.", image, gr.update(visible=False), 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))
+        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:
-        _, breed = choice.split(": ", 1)
+        if "Dog" in choice:
+            _, breed = choice.split(": ", 1)
+        else:
+            _, breed = choice.split("More about ", 1)
         description = get_dog_description(breed)
         return format_description(description, breed)
     except Exception as e: