Update app.py

app.py

@@ -250,22 +250,25 @@ def get_akc_breeds_link():
 # if __name__ == "__main__":
 #     iface.launch()
 
-
 def format_description(description, breed):
     if isinstance(description, dict):
         formatted_description = "\n".join([f"**{key}**: {value}" for key, value in description.items()])
     else:
         formatted_description = description
 
-    akc_link = get_akc_breeds_link()
-    formatted_description += f"\n\n**Want to learn more about dog breeds?**\n[Visit the AKC dog breeds page]({akc_link}) and search for {breed} to find detailed information."
+    formatted_description = f"""
+**Breed**: {breed}
+
+{formatted_description}
 
-    disclaimer = ("\n\n*Disclaimer: The external link provided leads to the American Kennel Club (AKC) dog breeds page. "
-                  "You may need to search for the specific breed on that page. "
-                  "I am not responsible for the content on external sites. "
-                  "Please refer to the AKC's terms of use and privacy policy.*")
-    formatted_description += disclaimer
+**Want to learn more about dog breeds?**
+[Visit the AKC dog breeds page]({get_akc_breeds_link()}) and search for {breed} to find detailed information.
 
+*Disclaimer: The external link provided leads to the American Kennel Club (AKC) dog breeds page. 
+You may need to search for the specific breed on that page. 
+I am not responsible for the content on external sites. 
+Please refer to the AKC's terms of use and privacy policy.*
+"""
     return formatted_description
 
 def predict_single_dog(image):
@@ -280,7 +283,7 @@ def predict_single_dog(image):
         topk_probs_percent = [f"{prob.item() * 100:.2f}%" for prob in topk_probs[0]]
     return top1_prob, topk_breeds, topk_probs_percent
 
-def detect_dogs(image):
+def detect_multiple_dogs(image):
     results = model_yolo(image)
     dogs = []
     for result in results:
@@ -292,19 +295,6 @@ def detect_dogs(image):
                 dogs.append((cropped_image, confidence, xyxy))
     return dogs
 
-
-def predict_breed(cropped_image):
-    image_tensor = preprocess_image(cropped_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 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)
@@ -313,21 +303,22 @@ def predict(image):
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
-        # First, try single dog prediction
-        top1_prob, topk_breeds, topk_probs_percent = predict_single_dog(image)
-        
-        if top1_prob >= 0.5:
-            # If confident enough, use single dog prediction
+        # First, check if there are multiple dogs using YOLO
+        dogs = detect_multiple_dogs(image)
+
+        if len(dogs) <= 1:
+            # Single dog or no dog detected, use direct classification
+            top1_prob, topk_breeds, topk_probs_percent = predict_single_dog(image)
             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)
         
-        # If not confident, use YOLO for multiple dog detection
-        dogs = detect_dogs(image)
-        if len(dogs) == 0:
-            return "No dogs detected or the image is unclear. Please upload a clearer image of a dog.", None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False)
-
+        # Multiple dogs detected, process each dog
         explanations = []
         visible_buttons = []
         annotated_image = image.copy()
@@ -342,14 +333,15 @@ def predict(image):
             if top1_prob >= 0.5:
                 breed = topk_breeds[0]
                 description = get_dog_description(breed)
-                explanations.append(f"Dog {i+1}:\n**Breed**: {breed}\n{format_description(description, breed)}")
+                explanations.append(f"Dog {i+1}:\n{format_description(description, breed)}")
             elif 0.2 <= top1_prob < 0.5:
-                explanation = (
-                    f"Dog {i+1}: Detected with moderate confidence. Here are the top 3 possible breeds:\n"
-                    f"1. **{topk_breeds[0]}** ({topk_probs_percent[0]})\n"
-                    f"2. **{topk_breeds[1]}** ({topk_probs_percent[1]})\n"
-                    f"3. **{topk_breeds[2]}** ({topk_probs_percent[2]})\n"
-                )
+                explanation = f"""
+Dog {i+1}: 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]}"])
             else:
@@ -361,12 +353,12 @@ def predict(image):
     except Exception as e:
         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]
     description = get_dog_description(breed_name)
     return format_description(description, breed_name)
 
+
 with gr.Blocks(css="""
     .container {
         max-width: 900px;