Update app.py

app.py

@@ -330,50 +330,48 @@ async def predict(image):
             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()
-            draw = ImageDraw.Draw(annotated_image)
-            font = ImageFont.load_default()
-
-            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))
+            # Only use YOLO if the confidence is very low
+            dogs = await detect_multiple_dogs(image)
+
+            if len(dogs) > 1:
+                # Multiple dogs detected, process each one
+                explanations = []
+                buttons = []
+                annotated_image = image.copy()
+                draw = ImageDraw.Draw(annotated_image)
+                font = ImageFont.load_default()
+
+                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:
+                        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:
-                    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}")
+                    return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
-            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)
-
-        else:
-            # 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)
+                # No dogs detected or only one dog with low confidence
+                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)
             
     except Exception as e:
         return f"An error occurred: {e}", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)