Update app.py

app.py

@@ -1,6 +1,7 @@
 import os
 import gradio as gr
 import re
+import folium
 from fastai.vision.all import *
 from groq import Groq
 from PIL import Image
@@ -26,6 +27,87 @@ def clean_bird_name(name):
     cleaned = ' '.join(cleaned.split())
     return cleaned
 
+def get_bird_habitat_map(bird_name):
+    """Get habitat map locations for the bird using Groq API"""
+    clean_name = clean_bird_name(bird_name)
+    
+    prompt = f"""
+    Provide a JSON array of the main habitat locations for the {clean_name} bird in the world. 
+    Return ONLY a JSON array with 3-5 entries, each containing:
+    1. "name": Location name
+    2. "lat": Latitude (numeric value)
+    3. "lon": Longitude (numeric value)
+    4. "description": Brief description of why this is a key habitat (2-3 sentences)
+    
+    Example format:
+    [
+      {{"name": "Example Location", "lat": 12.34, "lon": 56.78, "description": "Brief description"}},
+      ...
+    ]
+    
+    If this bird is not commonly found in Tanzania, include at least one location in Tanzania with a note about its unusual presence there.
+    """
+    
+    try:
+        chat_completion = client.chat.completions.create(
+            messages=[
+                {
+                    "role": "user",
+                    "content": prompt,
+                }
+            ],
+            model="llama-3.3-70b-versatile",
+        )
+        response = chat_completion.choices[0].message.content
+        
+        # Extract JSON from response (in case there's additional text)
+        import json
+        import re
+        
+        # Find JSON pattern in response
+        json_match = re.search(r'\[.*\]', response, re.DOTALL)
+        if json_match:
+            locations = json.loads(json_match.group())
+        else:
+            # Fallback if JSON extraction fails
+            locations = [
+                {"name": "Habitat information unavailable", "lat": 0, "lon": 0, 
+                 "description": "Could not retrieve habitat information for this bird."}
+            ]
+        
+        return locations
+        
+    except Exception as e:
+        return [{"name": f"Error retrieving habitat data: {str(e)}", "lat": 0, "lon": 0, 
+                 "description": "Please try again or check your connection."}]
+
+def create_habitat_map(habitat_locations):
+    """Create a folium map with the habitat locations"""
+    # Create a world-centered map
+    m = folium.Map(location=[20, 0], zoom_start=2)
+    
+    # Add markers for each habitat location
+    for location in habitat_locations:
+        name = location.get("name", "Unknown")
+        lat = location.get("lat", 0)
+        lon = location.get("lon", 0)
+        description = location.get("description", "No description available")
+        
+        # Skip invalid coordinates
+        if lat == 0 and lon == 0:
+            continue
+            
+        # Add marker
+        folium.Marker(
+            location=[lat, lon],
+            popup=folium.Popup(f"<b>{name}</b><br>{description}", max_width=300),
+            tooltip=name
+        ).add_to(m)
+    
+    # Save map to HTML
+    map_html = m._repr_html_()
+    return map_html
+
 def get_bird_info(bird_name):
     """Get detailed information about a bird using Groq API"""
     clean_name = clean_bird_name(bird_name)
@@ -79,10 +161,25 @@ def predict_and_get_info(img):
     # Also keep a clean version for display
     clean_top_bird = clean_bird_name(top_bird)
     
+    # Get habitat locations and create map
+    habitat_locations = get_bird_habitat_map(top_bird)
+    habitat_map_html = create_habitat_map(habitat_locations)
+    
     # Get detailed information about the top predicted bird
     bird_info = get_bird_info(top_bird)
     
-    return prediction_results, bird_info, clean_top_bird
+    # Create combined info with map at the top
+    combined_info = f"""
+## Natural Habitat Map for {clean_top_bird}
+<div style="height: 400px; width: 100%; border: 1px solid #ddd; border-radius: 8px; overflow: hidden; margin-bottom: 20px;">
+{habitat_map_html}
+</div>
+
+## Detailed Information
+{bird_info}
+"""
+    
+    return prediction_results, combined_info, clean_top_bird
 
 def follow_up_question(question, bird_name):
     """Allow researchers to ask follow-up questions about the identified bird"""
@@ -132,7 +229,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
         
         with gr.Column(scale=2):
             prediction_output = gr.Label(label="Top 5 Predictions", num_top_classes=5)
-            bird_info_output = gr.Markdown(label="Bird Information")
+            bird_info_output = gr.HTML(label="Bird Information")
     
     # Clear divider
     gr.Markdown("---")