Update app.py

app.py

@@ -250,15 +250,16 @@ def get_akc_breeds_link():
 # if __name__ == "__main__":
 #     iface.launch()
 
-def format_description(description, breed):
+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()])
+        formatted_description = "\n\n".join([f"**{key}**: {value}" for key, value in description.items() if key != "Breed"])
     else:
         formatted_description = description
 
+    header = f"**Dog {dog_number}: {breed}**\n\n" if is_multi_dog else f"**Breed: {breed}**\n\n"
+    
     formatted_description = f"""
-**Breed**: {breed}
-
+{header}
 {formatted_description}
 
 **Want to learn more about dog breeds?**
@@ -271,30 +272,6 @@ Please refer to the AKC's terms of use and privacy policy.*
 """
     return formatted_description
 
-def predict_single_dog(image):
-    image_tensor = preprocess_image(image)
-    with torch.no_grad():
-        output = model(image_tensor)
-        logits = output[0] if isinstance(output, tuple) else output
-        probabilities = F.softmax(logits, dim=1)
-        topk_probs, topk_indices = torch.topk(probabilities, k=3)
-        top1_prob = topk_probs[0][0].item()
-        topk_breeds = [dog_breeds[idx.item()] for idx in topk_indices[0]]
-        topk_probs_percent = [f"{prob.item() * 100:.2f}%" for prob in topk_probs[0]]
-    return top1_prob, topk_breeds, topk_probs_percent
-
-def detect_multiple_dogs(image):
-    results = model_yolo(image)
-    dogs = []
-    for result in results:
-        for box in result.boxes:
-            if box.cls == 16:  # COCO dataset class for dog is 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))
-    return dogs
-
 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)
@@ -318,7 +295,7 @@ def predict(image):
                 return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
             elif 0.2 <= top1_prob < 0.5:
                 explanation = f"""
-Detected with moderate confidence. Here are the top 3 possible breeds:
+**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]})
@@ -336,28 +313,28 @@ Click on a button below to view more information about each breed.
         annotated_image = image.copy()
         draw = ImageDraw.Draw(annotated_image)
 
-        for i, (cropped_image, _, box) in enumerate(dogs):
+        for i, (cropped_image, _, box) in enumerate(dogs, 1):
             top1_prob, topk_breeds, topk_probs_percent = predict_single_dog(cropped_image)
 
             draw.rectangle(box, outline="red", width=3)
-            draw.text((box[0], box[1]), f"Dog {i+1}", fill="red")
+            draw.text((box[0], box[1]), f"Dog {i}", fill="red", font=ImageFont.truetype("arial.ttf", 20))
 
             if top1_prob >= 0.5:
                 breed = topk_breeds[0]
                 description = get_dog_description(breed)
-                explanations.append(f"Dog {i+1}:\n{format_description(description, breed)}")
+                explanations.append(format_description(description, breed, is_multi_dog=True, dog_number=i))
             elif 0.2 <= top1_prob < 0.5:
                 explanation = f"""
-Dog {i+1}: Detected with moderate confidence. Here are the top 3 possible breeds:
+**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)
-                visible_buttons.extend([f"More about {topk_breeds[0]}", f"More about {topk_breeds[1]}", f"More about {topk_breeds[2]}"])
+                visible_buttons.extend([f"More about Dog {i}: {topk_breeds[0]}", f"More about Dog {i}: {topk_breeds[1]}", f"More about Dog {i}: {topk_breeds[2]}"])
             else:
-                explanations.append(f"Dog {i+1}: The image is unclear or the breed is not in the dataset.")
+                explanations.append(f"**Dog {i}**: 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=True, value=visible_buttons[1] if len(visible_buttons) >= 2 else ""), gr.update(visible=True, value=visible_buttons[2] if len(visible_buttons) >= 3 else "")
@@ -366,7 +343,7 @@ Dog {i+1}: Detected with moderate confidence. Here are the top 3 possible breeds
         return f"An error occurred: {e}", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
 
 def show_details(breed):
-    breed_name = breed.split("More about ")[-1]
+    breed_name = breed.split("More about ")[-1].split(": ")[-1]  # Handle both single and multi-dog cases
     description = get_dog_description(breed_name)
     return format_description(description, breed_name)