Update app.py

app.py

@@ -308,45 +308,74 @@ async def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        # Always start with single dog prediction
-        top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
+        # First, use YOLO to detect multiple dogs
+        dogs = await detect_multiple_dogs(image)
 
-        if top1_prob >= 0.5:
-            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)
-        
-        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)
-        
-        else:  # 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 len(dogs) <= 1:
+            # Single dog or no dog detected, use the original method
+            top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
+            return process_single_dog_result(top1_prob, topk_breeds, topk_probs_percent, image)
+        else:
+            # Multiple dogs detected
+            return process_multiple_dogs_result(dogs, image)
 
     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."
+async def process_single_dog_result(top1_prob, topk_breeds, topk_probs_percent, image):
+    if top1_prob >= 0.5:
+        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)
+    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)
+    else:
+        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)
+
+async def process_multiple_dogs_result(dogs, image):
+    annotated_image = image.copy()
+    draw = ImageDraw.Draw(annotated_image)
+    font = ImageFont.load_default()
     
-    try:
-        if "Dog" in choice:
-            _, breed = choice.split(": ", 1)
+    explanations = []
+    buttons = []
+    
+    for i, (cropped_image, _, box) in enumerate(dogs, 1):
+        top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
+        
+        draw.rectangle(box, outline="red", width=3)
+        draw.text((box[0], box[1]), f"Dog {i}", fill="yellow", font=font)
+        
+        if top1_prob >= 0.5:
+            breed = topk_breeds[0]
+            description = get_dog_description(breed)
+            explanations.append(format_description(description, breed, is_multi_dog=True, dog_number=i))
         else:
-            _, breed = choice.split("More about ", 1)
-        description = get_dog_description(breed)
-        return format_description(description, breed)
-    except Exception as e:
-        return f"An error occurred while showing details: {e}"
+            explanation = f"""
+Dog {i}: Detected with moderate confidence. Here are the top 3 possible breeds:
+1. **{topk_breeds[0]}** ({topk_probs_percent[0]})
+2. **{topk_breeds[1]}** ({topk_probs_percent[1]})
+3. **{topk_breeds[2]}** ({topk_probs_percent[2]})
+"""
+            explanations.append(explanation)
+            for breed in topk_breeds:
+                buttons.append(f"More about Dog {i}: {breed}")
+
+    final_explanation = "\n\n---\n\n".join(explanations)
+    
+    if buttons:
+        return final_explanation, annotated_image, gr.update(visible=True, choices=buttons), gr.update(visible=False), gr.update(visible=False)
+    else:
+        return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
 # Gradio 介面設置
 with gr.Blocks(css="""