Update app.py

app.py

@@ -153,8 +153,8 @@ def format_description(description, breed):
 
     return formatted_description
 
-async def predict_single_dog(image):
-    return await asyncio.to_thread(_predict_single_dog, image)
+# async def predict_single_dog(image):
+#     return await asyncio.to_thread(_predict_single_dog, image)
 
 # def _predict_single_dog(image):
 #     image_tensor = preprocess_image(image)
@@ -182,14 +182,16 @@ async def predict_single_dog(image):
 
 
 async def predict_single_dog(image):
-    image_tensor = preprocess(image).unsqueeze(0)
+    image_tensor = preprocess_image(image)
     with torch.no_grad():
         output = model(image_tensor)
-        probabilities = torch.nn.functional.softmax(output[0], dim=0)
-        top3_prob, top3_catid = torch.topk(probabilities, 3)
-        top3_breeds = [dog_breeds[idx.item()] for idx in top3_catid]
-        top3_probs = [f"{prob.item()*100:.2f}%" for prob in top3_prob]
-    return top3_prob[0].item(), top3_breeds, top3_probs
+        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 detect_multiple_dogs(image, conf_threshold=0.3, iou_threshold=0.45):
     results = model_yolo(image, conf=conf_threshold, iou=iou_threshold)[0]
@@ -402,7 +404,7 @@ async def process_single_dog(image):
 
 async def predict(image):
     if image is None:
-        return "Please upload an image to start.", None, [], gr.update(visible=False), None
+        return "Please upload an image to start.", None, gr.update(visible=False, choices=[]), None
 
     try:
         if isinstance(image, np.ndarray):
@@ -460,13 +462,12 @@ async def predict(image):
                 "is_multi_dog": len(dogs) > 1,
                 "dogs_info": explanations
             }
-            return final_explanation, annotated_image, gr.update(visible=False), initial_state
+            return final_explanation, annotated_image, gr.update(visible=False, choices=[]), initial_state
 
     except Exception as e:
-        error_msg = f"An error occurred: {str(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), None
-        
+        return error_msg, None, gr.update(visible=False, choices=[]), None
 
 def show_details(choice, previous_output, initial_state):
     if not choice:
@@ -483,7 +484,7 @@ def show_details(choice, previous_output, initial_state):
         return formatted_description, gr.update(visible=True), initial_state
     except Exception as e:
         error_msg = f"An error occurred while showing details: {e}"
-        logger.error(error_msg)
+        print(error_msg)
         return error_msg, gr.update(visible=True), initial_state
 
 def go_back(state):
@@ -495,7 +496,6 @@ def go_back(state):
         gr.update(visible=False),
         state
     )
-    
 
 with gr.Blocks() as iface:
     gr.HTML("<h1 style='text-align: center;'>🐶 Dog Breed Classifier 🔍</h1>")
@@ -530,7 +530,7 @@ with gr.Blocks() as iface:
         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
@@ -539,6 +539,4 @@ with gr.Blocks() as iface:
     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()
-
-    
+    iface.launch()