Update app.py

app.py

@@ -50,8 +50,14 @@ def predict_and_get_info(img):
     # Get prediction
     pred, pred_idx, probs = learn.predict(img)
     
-    # Format prediction results
-    prediction_results = {labels[i]: float(probs[i]) for i in range(len(labels))}
+    # Get top 5 predictions (or all if less than 5)
+    num_classes = min(5, len(labels))
+    top_indices = probs.argsort(descending=True)[:num_classes]
+    top_probs = probs[top_indices]
+    top_labels = [labels[i] for i in top_indices]
+    
+    # Format as dictionary for display
+    prediction_results = {top_labels[i]: float(top_probs[i]) for i in range(num_classes)}
     
     # Get top prediction
     top_bird = str(pred)
@@ -59,10 +65,13 @@ def predict_and_get_info(img):
     # Get detailed information about the top predicted bird
     bird_info = get_bird_info(top_bird)
     
-    return prediction_results, bird_info
+    return prediction_results, bird_info, top_bird
 
 def follow_up_question(question, bird_name):
     """Allow researchers to ask follow-up questions about the identified bird"""
+    if not question.strip() or not bird_name:
+        return "Please identify a bird first and ask a specific question about it."
+    
     prompt = f"""
     The researcher is asking about the {bird_name} bird: "{question}"
     
@@ -90,47 +99,89 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
     gr.Markdown("# Bird Species Identification for Researchers")
     gr.Markdown("Upload an image to identify bird species and get detailed information relevant to research in Tanzania and climate change studies.")
     
+    # Store the current bird for context
+    current_bird = gr.State("")
+    
+    # Main identification section
     with gr.Row():
         with gr.Column(scale=1):
             input_image = gr.Image(type="pil", label="Upload Bird Image")
             submit_btn = gr.Button("Identify Bird", variant="primary")
         
         with gr.Column(scale=2):
-            with gr.Row():
-                prediction_output = gr.Label(label="Prediction Results")
-            with gr.Row():
-                bird_info_output = gr.Markdown(label="Bird Information")
+            prediction_output = gr.Label(label="Top 5 Predictions", num_top_classes=5)
+            bird_info_output = gr.Markdown(label="Bird Information")
+    
+    # Clear divider
+    gr.Markdown("---")
+    
+    # Follow-up question section with improved UI
+    gr.Markdown("## Research Questions")
+    
+    conversation_history = gr.Markdown("")
     
-    # For follow-up questions
     with gr.Row():
-        gr.Markdown("## Ask Follow-up Questions")
+        follow_up_input = gr.Textbox(
+            label="Ask a question about this bird", 
+            placeholder="Example: How has climate change affected this bird's migration pattern?",
+            lines=2
+        )
     
     with gr.Row():
-        with gr.Column(scale=1):
-            current_bird = gr.State("")
-            follow_up_input = gr.Textbox(label="Your Question", placeholder="Ask more about this bird species...")
-            follow_up_btn = gr.Button("Submit Question")
-        
-        with gr.Column(scale=2):
-            follow_up_output = gr.Markdown(label="Answer")
+        follow_up_btn = gr.Button("Submit Question", variant="primary")
+        clear_btn = gr.Button("Clear Conversation")
     
     # Set up event handlers
     def process_image(img):
-        pred_results, info = predict_and_get_info(img)
-        # Extract bird name from top prediction
-        bird_name = max(pred_results.items(), key=lambda x: x[1])[0]
-        return pred_results, info, bird_name
+        if img is None:
+            return None, "Please upload an image", "", ""
+        
+        try:
+            pred_results, info, bird_name = predict_and_get_info(img)
+            return pred_results, info, bird_name, ""
+        except Exception as e:
+            return None, f"Error processing image: {str(e)}", "", ""
+    
+    def update_conversation(question, bird_name, history):
+        if not question.strip():
+            return history
+        
+        answer = follow_up_question(question, bird_name)
+        
+        # Format the conversation with clear separation
+        new_exchange = f"""
+### Question:
+{question}
+
+### Answer:
+{answer}
+
+---
+"""
+        updated_history = new_exchange + history
+        return updated_history
+    
+    def clear_conversation_history():
+        return ""
     
     submit_btn.click(
         process_image,
         inputs=[input_image],
-        outputs=[prediction_output, bird_info_output, current_bird]
+        outputs=[prediction_output, bird_info_output, current_bird, conversation_history]
     )
     
     follow_up_btn.click(
-        follow_up_question,
-        inputs=[follow_up_input, current_bird],
-        outputs=[follow_up_output]
+        update_conversation,
+        inputs=[follow_up_input, current_bird, conversation_history],
+        outputs=[conversation_history]
+    ).then(
+        lambda: "", 
+        outputs=follow_up_input
+    )
+    
+    clear_btn.click(
+        clear_conversation_history,
+        outputs=[conversation_history]
     )
 
 # Launch the app