Update app.py

app.py

@@ -308,13 +308,16 @@ async def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        # First, use YOLO to detect multiple dogs
+        # Always use YOLO to detect dogs
         dogs = await detect_multiple_dogs(image)
 
-        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 await process_single_dog_result(top1_prob, topk_breeds, topk_probs_percent, image)
+        if len(dogs) == 0:
+            return "No dogs detected in the image. Please upload a clear image of a dog.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
+        elif len(dogs) == 1:
+            # Single dog detected
+            cropped_image, _, _ = dogs[0]
+            top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
+            return await process_single_dog_result(top1_prob, topk_breeds, topk_probs_percent, image, dogs[0][2])
         else:
             # Multiple dogs detected
             return await process_multiple_dogs_result(dogs, image)
@@ -322,25 +325,6 @@ async def predict(image):
     except Exception as e:
         return f"An error occurred: {str(e)}", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
-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)
@@ -358,7 +342,7 @@ async def process_multiple_dogs_result(dogs, image):
         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))
+            explanations.append(f"Dog {i}: **{breed}** ({top1_prob*100:.2f}% confidence)\n\n{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:
@@ -377,6 +361,30 @@ Dog {i}: Detected with moderate confidence. Here are the top 3 possible breeds:
     else:
         return final_explanation, annotated_image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
+async def process_single_dog_result(top1_prob, topk_breeds, topk_probs_percent, image, box):
+    annotated_image = image.copy()
+    draw = ImageDraw.Draw(annotated_image)
+    draw.rectangle(box, outline="red", width=3)
+    draw.text((box[0], box[1]), "Dog", fill="yellow", font=ImageFont.load_default())
+
+    if top1_prob >= 0.5:
+        breed = topk_breeds[0]
+        description = get_dog_description(breed)
+        formatted_description = format_description(description, breed)
+        return formatted_description, annotated_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, annotated_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.", annotated_image, 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."