Update app.py

app.py

@@ -132,6 +132,18 @@ def get_akc_breeds_link(breed):
     return f"{base_url}{breed_url}/"
 
 
+# async 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
+
 
 async def predict_single_dog(image):
     image_tensor = preprocess_image(image)
@@ -142,8 +154,13 @@ async def predict_single_dog(image):
         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
+        
+        # Calculate relative probabilities for display
+        raw_probs = [prob.item() for prob in topk_probs[0]]
+        sum_probs = sum(raw_probs)
+        relative_probs = [f"{(prob/sum_probs * 100):.2f}%" for prob in raw_probs]
+        
+    return top1_prob, topk_breeds, relative_probs
     
 
 async def detect_multiple_dogs(image, conf_threshold=0.35, iou_threshold=0.55):
@@ -198,66 +215,210 @@ def calculate_iou(box1, box2):
 
 
 
+# async def process_single_dog(image):
+#     top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
+#     if top1_prob < 0.15:
+#         initial_state = {
+#             "explanation": "The image is unclear or the breed is not in the dataset. Please upload a clearer image of a dog.",
+#             "buttons": [],
+#             "show_back": False,
+#             "image": None,
+#             "is_multi_dog": False
+#         }
+#         return initial_state["explanation"], None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
+
+#     breed = topk_breeds[0]
+#     description = get_dog_description(breed)
+
+#     if top1_prob >= 0.45:
+#         formatted_description = format_description(description, breed)
+#         initial_state = {
+#             "explanation": formatted_description,
+#             "buttons": [],
+#             "show_back": False,
+#             "image": image,
+#             "is_multi_dog": False
+#         }
+#         return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
+#     else:
+#         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."
+#         )
+#         buttons = [
+#             gr.update(visible=True, value=f"More about {topk_breeds[0]}"),
+#             gr.update(visible=True, value=f"More about {topk_breeds[1]}"),
+#             gr.update(visible=True, value=f"More about {topk_breeds[2]}")
+#         ]
+#         initial_state = {
+#             "explanation": explanation,
+#             "buttons": buttons,
+#             "show_back": True,
+#             "image": image,
+#             "is_multi_dog": False
+#         }
+#         return explanation, image, buttons[0], buttons[1], buttons[2], gr.update(visible=True), initial_state
+
+
+# async def predict(image):
+#     if image is None:
+#         return "Please upload an image to start.", None, gr.update(visible=False, choices=[]), None
+
+#     try:
+#         if isinstance(image, np.ndarray):
+#             image = Image.fromarray(image)
+
+#         dogs = await detect_multiple_dogs(image)
+        
+#         color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
+#         buttons = []
+#         annotated_image = image.copy()
+#         draw = ImageDraw.Draw(annotated_image)
+#         font = ImageFont.load_default()
+
+#         dogs_info = ""
+
+#         for i, (cropped_image, detection_confidence, box) in enumerate(dogs):
+#             buttons_html = ""  
+#             top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
+#             color = color_list[i % len(color_list)]
+#             draw.rectangle(box, outline=color, width=3)
+#             draw.text((box[0] + 5, box[1] + 5), f"Dog {i+1}", fill=color, font=font)
+        
+#             combined_confidence = detection_confidence * top1_prob
+#             dogs_info += f'<div class="dog-info" style="border-left: 5px solid {color}; margin-bottom: 20px; padding: 15px;">'
+#             dogs_info += f'<h2>Dog {i+1}</h2>'
+        
+#             if top1_prob >= 0.45:
+#                 breed = topk_breeds[0]
+#                 description = get_dog_description(breed)
+#                 dogs_info += format_description_html(description, breed)
+        
+#             elif combined_confidence >= 0.15:
+#                 dogs_info += f"<p>Top 3 possible breeds:</p><ul>"
+#                 for j, (breed, prob) in enumerate(zip(topk_breeds[:3], topk_probs_percent[:3])):
+#                     prob = float(prob.replace('%', ''))
+#                     dogs_info += f"<li><strong>{breed}</strong> ({prob:.2f}% confidence)</li>"
+#                 dogs_info += "</ul>"
+        
+#                 for breed in topk_breeds[:3]:
+#                     button_id = f"Dog {i+1}: More about {breed}"
+#                     buttons_html += f'<button class="breed-button" onclick="handle_button_click(\'{button_id}\')">{breed}</button>'
+#                     buttons.append(button_id)
+        
+#             else:
+#                 dogs_info += "<p>The image is unclear or the breed is not in the dataset. Please upload a clearer image.</p>"
+        
+#             dogs_info += '</div>'    
+
+        
+#         buttons_html = ""  
+      
+#         html_output = f"""
+#         <style>
+#         .dog-info {{ border: 1px solid #ddd; margin-bottom: 20px; padding: 15px; border-radius: 5px; box-shadow: 0 2px 5px rgba(0,0,0,0.1); }}
+#         .dog-info h2 {{ background-color: #f0f0f0; padding: 10px; margin: -15px -15px 15px -15px; border-radius: 5px 5px 0 0; }}
+#         .breed-buttons {{ margin-top: 10px; }}
+#         .breed-button {{ margin-right: 10px; margin-bottom: 10px; padding: 5px 10px; background-color: #4CAF50; color: white; border: none; border-radius: 3px; cursor: pointer; }}
+#         </style>
+#         {dogs_info}
+#         """     
+
+#         if buttons:
+#             html_output += """
+#             <script>
+#             function handle_button_click(button_id) {
+#                 const radio = document.querySelector('input[type=radio][value="' + button_id + '"]');
+#                 if (radio) {
+#                     radio.click();
+#                 } else {
+#                     console.error("Radio button not found:", button_id);
+#                 }
+#             }
+#             </script>
+#             """
+#             initial_state = {
+#                 "dogs_info": dogs_info,
+#                 "buttons": buttons,
+#                 "show_back": True,
+#                 "image": annotated_image,
+#                 "is_multi_dog": len(dogs) > 1,
+#                 "html_output": html_output  
+#             }
+#             return html_output, annotated_image, gr.update(visible=True, choices=buttons), initial_state
+#         else:
+#             initial_state = {
+#                 "dogs_info": dogs_info,
+#                 "buttons": [],
+#                 "show_back": False,
+#                 "image": annotated_image,
+#                 "is_multi_dog": len(dogs) > 1,
+#                 "html_output": html_output  
+#             }
+#             return html_output, annotated_image, gr.update(visible=False, choices=[]), initial_state
+
+
+#     except Exception as e:
+#         error_msg = f"An error occurred: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
+#         print(error_msg)
+#         return error_msg, None, gr.update(visible=False, choices=[]), None
+
+
+
 async def process_single_dog(image):
-    top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(image)
+    top1_prob, topk_breeds, relative_probs = await predict_single_dog(image)
+    
+    # Case 1: Low confidence - unclear image or breed not in dataset
     if top1_prob < 0.15:
         initial_state = {
             "explanation": "The image is unclear or the breed is not in the dataset. Please upload a clearer image of a dog.",
-            "buttons": [],
-            "show_back": False,
             "image": None,
             "is_multi_dog": False
         }
-        return initial_state["explanation"], None, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
+        return initial_state["explanation"], None, initial_state
 
     breed = topk_breeds[0]
-    description = get_dog_description(breed)
-
+    
+    # Case 2: High confidence - single breed result
     if top1_prob >= 0.45:
+        description = get_dog_description(breed)
         formatted_description = format_description(description, breed)
         initial_state = {
             "explanation": formatted_description,
-            "buttons": [],
-            "show_back": False,
             "image": image,
             "is_multi_dog": False
         }
-        return formatted_description, image, gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), gr.update(visible=False), initial_state
+        return formatted_description, image, initial_state
+        
+    # Case 3: Medium confidence - show top 3 breeds with relative probabilities
     else:
-        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."
-        )
-        buttons = [
-            gr.update(visible=True, value=f"More about {topk_breeds[0]}"),
-            gr.update(visible=True, value=f"More about {topk_breeds[1]}"),
-            gr.update(visible=True, value=f"More about {topk_breeds[2]}")
-        ]
+        breeds_info = ""
+        for i, (breed, prob) in enumerate(zip(topk_breeds, relative_probs)):
+            description = get_dog_description(breed)
+            formatted_description = format_description(description, breed)
+            breeds_info += f"\n\nBreed {i+1}: **{breed}** (Confidence: {prob})\n{formatted_description}"
+
         initial_state = {
-            "explanation": explanation,
-            "buttons": buttons,
-            "show_back": True,
+            "explanation": breeds_info,
             "image": image,
             "is_multi_dog": False
         }
-        return explanation, image, buttons[0], buttons[1], buttons[2], gr.update(visible=True), initial_state
-
+        return breeds_info, image, initial_state
+        
 
 async def predict(image):
     if image is None:
-        return "Please upload an image to start.", None, gr.update(visible=False, choices=[]), None
+        return "Please upload an image to start.", None, None
 
     try:
         if isinstance(image, np.ndarray):
             image = Image.fromarray(image)
 
         dogs = await detect_multiple_dogs(image)
-        
         color_list = ['#FF0000', '#00FF00', '#0000FF', '#FFFF00', '#00FFFF', '#FF00FF', '#800080', '#FFA500']
-        buttons = []
         annotated_image = image.copy()
         draw = ImageDraw.Draw(annotated_image)
         font = ImageFont.load_default()
@@ -265,89 +426,73 @@ async def predict(image):
         dogs_info = ""
 
         for i, (cropped_image, detection_confidence, box) in enumerate(dogs):
-            buttons_html = ""  
-            top1_prob, topk_breeds, topk_probs_percent = await predict_single_dog(cropped_image)
             color = color_list[i % len(color_list)]
             draw.rectangle(box, outline=color, width=3)
             draw.text((box[0] + 5, box[1] + 5), f"Dog {i+1}", fill=color, font=font)
         
+            top1_prob, topk_breeds, relative_probs = await predict_single_dog(cropped_image)
             combined_confidence = detection_confidence * top1_prob
+            
             dogs_info += f'<div class="dog-info" style="border-left: 5px solid {color}; margin-bottom: 20px; padding: 15px;">'
             dogs_info += f'<h2>Dog {i+1}</h2>'
-        
-            if top1_prob >= 0.45:
+            
+            if combined_confidence < 0.15:
+                dogs_info += "<p>The image is unclear or the breed is not in the dataset. Please upload a clearer image.</p>"
+                
+            elif top1_prob >= 0.45:
                 breed = topk_breeds[0]
                 description = get_dog_description(breed)
                 dogs_info += format_description_html(description, breed)
-        
-            elif combined_confidence >= 0.15:
-                dogs_info += f"<p>Top 3 possible breeds:</p><ul>"
-                for j, (breed, prob) in enumerate(zip(topk_breeds[:3], topk_probs_percent[:3])):
-                    prob = float(prob.replace('%', ''))
-                    dogs_info += f"<li><strong>{breed}</strong> ({prob:.2f}% confidence)</li>"
-                dogs_info += "</ul>"
-        
-                for breed in topk_breeds[:3]:
-                    button_id = f"Dog {i+1}: More about {breed}"
-                    buttons_html += f'<button class="breed-button" onclick="handle_button_click(\'{button_id}\')">{breed}</button>'
-                    buttons.append(button_id)
-        
+                
             else:
-                dogs_info += "<p>The image is unclear or the breed is not in the dataset. Please upload a clearer image.</p>"
-        
-            dogs_info += '</div>'    
+                dogs_info += "<h3>Top 3 possible breeds:</h3>"
+                for breed, prob in zip(topk_breeds, relative_probs):
+                    description = get_dog_description(breed)
+                    dogs_info += f"<div class='breed-section'>"
+                    dogs_info += f"<h4>{breed} (Confidence: {prob})</h4>"
+                    dogs_info += format_description_html(description, breed)
+                    dogs_info += "</div>"
+            
+            dogs_info += '</div>'
 
-        
-        buttons_html = ""  
-      
         html_output = f"""
         <style>
-        .dog-info {{ border: 1px solid #ddd; margin-bottom: 20px; padding: 15px; border-radius: 5px; box-shadow: 0 2px 5px rgba(0,0,0,0.1); }}
-        .dog-info h2 {{ background-color: #f0f0f0; padding: 10px; margin: -15px -15px 15px -15px; border-radius: 5px 5px 0 0; }}
-        .breed-buttons {{ margin-top: 10px; }}
-        .breed-button {{ margin-right: 10px; margin-bottom: 10px; padding: 5px 10px; background-color: #4CAF50; color: white; border: none; border-radius: 3px; cursor: pointer; }}
+        .dog-info {{ 
+            border: 1px solid #ddd; 
+            margin-bottom: 20px; 
+            padding: 15px; 
+            border-radius: 5px; 
+            box-shadow: 0 2px 5px rgba(0,0,0,0.1); 
+        }}
+        .dog-info h2 {{ 
+            background-color: #f0f0f0; 
+            padding: 10px; 
+            margin: -15px -15px 15px -15px; 
+            border-radius: 5px 5px 0 0; 
+        }}
+        .breed-section {{
+            margin-bottom: 20px;
+            padding: 10px;
+            background-color: #f8f8f8;
+            border-radius: 5px;
+        }}
         </style>
         {dogs_info}
-        """     
-
-        if buttons:
-            html_output += """
-            <script>
-            function handle_button_click(button_id) {
-                const radio = document.querySelector('input[type=radio][value="' + button_id + '"]');
-                if (radio) {
-                    radio.click();
-                } else {
-                    console.error("Radio button not found:", button_id);
-                }
-            }
-            </script>
-            """
-            initial_state = {
-                "dogs_info": dogs_info,
-                "buttons": buttons,
-                "show_back": True,
-                "image": annotated_image,
-                "is_multi_dog": len(dogs) > 1,
-                "html_output": html_output  
-            }
-            return html_output, annotated_image, gr.update(visible=True, choices=buttons), initial_state
-        else:
-            initial_state = {
-                "dogs_info": dogs_info,
-                "buttons": [],
-                "show_back": False,
-                "image": annotated_image,
-                "is_multi_dog": len(dogs) > 1,
-                "html_output": html_output  
-            }
-            return html_output, annotated_image, gr.update(visible=False, choices=[]), initial_state
+        """
 
+        initial_state = {
+            "dogs_info": dogs_info,
+            "image": annotated_image,
+            "is_multi_dog": len(dogs) > 1,
+            "html_output": html_output
+        }
+        
+        return html_output, annotated_image, initial_state
 
     except Exception as e:
         error_msg = f"An error occurred: {str(e)}\n\nTraceback:\n{traceback.format_exc()}"
         print(error_msg)
-        return error_msg, None, gr.update(visible=False, choices=[]), None
+        return error_msg, None, None
 
 
 def show_details_html(choice, previous_output, initial_state):
@@ -402,47 +547,76 @@ def go_back(state):
     )
 
 
+# with gr.Blocks() as iface:
+#     gr.HTML("<h1 style='text-align: center;'>🐶 Dog Breed Classifier 🔍</h1>")
+#     gr.HTML("<p style='text-align: center;'>Upload a picture of a dog, and the model will predict its breed, provide detailed information, and include an extra information link!</p>")
+    
+#     with gr.Row():
+#         input_image = gr.Image(label="Upload a dog image", type="pil")
+#         output_image = gr.Image(label="Annotated Image")
+    
+#     output = gr.HTML(label="Prediction Results")  
+    
+#     breed_buttons = gr.Radio(choices=[], label="More Information", visible=False)
+    
+#     back_button = gr.Button("Back", visible=False)
+    
+#     initial_state = gr.State()
+    
+#     input_image.change(
+#         predict,
+#         inputs=input_image,
+#         outputs=[output, output_image, breed_buttons, initial_state]
+#     )
+
+#     breed_buttons.change(
+#         show_details_html,
+#         inputs=[breed_buttons, output, initial_state],
+#         outputs=[output, back_button, initial_state]
+#     )
+
+#     back_button.click(
+#         go_back,
+#         inputs=[initial_state],
+#         outputs=[output, output_image, breed_buttons, back_button, initial_state]
+#     )
+    
+#     gr.Examples(
+#         examples=['Border_Collie.jpg', 'Golden_Retriever.jpeg', 'Saint_Bernard.jpeg', 'French_Bulldog.jpeg', 'Samoyed.jpg'],
+#         inputs=input_image
+#     )
+
+#     gr.HTML('For more details on this project and other work, feel free to visit my GitHub <a href="https://github.com/Eric-Chung-0511/Learning-Record/tree/main/Data%20Science%20Projects/Dog_Breed_Classifier">Dog Breed Classifier</a>')
+
+
+# if __name__ == "__main__":
+#     iface.launch()
+
+
+
 with gr.Blocks() as iface:
     gr.HTML("<h1 style='text-align: center;'>🐶 Dog Breed Classifier 🔍</h1>")
-    gr.HTML("<p style='text-align: center;'>Upload a picture of a dog, and the model will predict its breed, provide detailed information, and include an extra information link!</p>")
+    gr.HTML("<p style='text-align: center;'>Upload a picture of a dog, and the model will predict its breed and provide detailed information!</p>")
     
     with gr.Row():
         input_image = gr.Image(label="Upload a dog image", type="pil")
         output_image = gr.Image(label="Annotated Image")
     
-    output = gr.HTML(label="Prediction Results")  
-    
-    breed_buttons = gr.Radio(choices=[], label="More Information", visible=False)
-    
-    back_button = gr.Button("Back", visible=False)
-    
+    output = gr.HTML(label="Prediction Results")
     initial_state = gr.State()
     
     input_image.change(
         predict,
         inputs=input_image,
-        outputs=[output, output_image, breed_buttons, initial_state]
-    )
-
-    breed_buttons.change(
-        show_details_html,
-        inputs=[breed_buttons, output, initial_state],
-        outputs=[output, back_button, initial_state]
-    )
-
-    back_button.click(
-        go_back,
-        inputs=[initial_state],
-        outputs=[output, output_image, breed_buttons, back_button, initial_state]
+        outputs=[output, output_image, initial_state]
     )
     
     gr.Examples(
         examples=['Border_Collie.jpg', 'Golden_Retriever.jpeg', 'Saint_Bernard.jpeg', 'French_Bulldog.jpeg', 'Samoyed.jpg'],
         inputs=input_image
     )
-
+    
     gr.HTML('For more details on this project and other work, feel free to visit my GitHub <a href="https://github.com/Eric-Chung-0511/Learning-Record/tree/main/Data%20Science%20Projects/Dog_Breed_Classifier">Dog Breed Classifier</a>')
 
-
 if __name__ == "__main__":
     iface.launch()