Update app.py

app.py

@@ -5,6 +5,9 @@ import folium
 from fastai.vision.all import *
 from groq import Groq
 from PIL import Image
+import time
+import json
+from functools import lru_cache
 
 # Load the trained model
 learn = load_learner('export.pkl')
@@ -15,6 +18,9 @@ client = Groq(
     api_key=os.environ.get("GROQ_API_KEY"),
 )
 
+# Cache directory for API responses
+os.makedirs("cache", exist_ok=True)
+
 # Language translations
 translations = {
     "en": {
@@ -34,7 +40,11 @@ translations = {
         "answer_title": "Answer:",
         "habitat_map_title": "Natural Habitat Map for",
         "detailed_info_title": "Detailed Information",
-        "language_label": "Language / Lugha"
+        "language_label": "Language / Lugha",
+        "loading": "Processing...",
+        "low_confidence": "Low confidence prediction. Results may not be accurate.",
+        "not_a_bird": "The image may not contain a bird. Please upload a clear image of a bird.",
+        "other_message": "This bird is not in our trained dataset or the image may not be of a bird. Please try uploading a different image."
     },
     "sw": {
         "app_title": "Mtafiti wa Ndege: Utambuzi wa Kiotomatiki kwa Watafiti",
@@ -53,38 +63,117 @@ translations = {
         "answer_title": "Jibu:",
         "habitat_map_title": "Ramani ya Makazi Asilia ya",
         "detailed_info_title": "Taarifa za Kina",
-        "language_label": "Language / Lugha"
+        "language_label": "Language / Lugha",
+        "loading": "Inachakata...",
+        "low_confidence": "Utabiri wa uhakika mdogo. Matokeo yanaweza kuwa si sahihi.",
+        "not_a_bird": "Picha inaweza isiwe ya ndege. Tafadhali pakia picha wazi ya ndege.",
+        "other_message": "Ndege huyu haipatikani katika hifadhidata yetu au picha inaweza isiwe ya ndege. Tafadhali jaribu kupakia picha nyingine."
     }
 }
 
 def clean_bird_name(name):
     """Clean bird name by removing numbers and special characters, and fix formatting"""
+    # Remove numbers and dots at the beginning
     cleaned = re.sub(r'^\d+\.', '', name)
+    # Replace underscores with spaces
     cleaned = cleaned.replace('_', ' ')
+    # Remove any remaining special characters
     cleaned = re.sub(r'[^\w\s]', '', cleaned)
+    # Fix spacing
     cleaned = ' '.join(cleaned.split())
     return cleaned
 
+def get_cache_path(function_name, key):
+    """Generate a cache file path"""
+    safe_key = re.sub(r'[^\w]', '_', key)
+    return f"cache/{function_name}_{safe_key}.json"
+
+def save_to_cache(function_name, key, data):
+    """Save API response to cache"""
+    try:
+        cache_path = get_cache_path(function_name, key)
+        with open(cache_path, 'w') as f:
+            json.dump({"data": data, "timestamp": time.time()}, f)
+    except Exception as e:
+        print(f"Error saving to cache: {e}")
+
+def load_from_cache(function_name, key, max_age=86400):  # Default max age: 1 day
+    """Load API response from cache if it exists and is not too old"""
+    try:
+        cache_path = get_cache_path(function_name, key)
+        if os.path.exists(cache_path):
+            with open(cache_path, 'r') as f:
+                cached = json.load(f)
+                if time.time() - cached["timestamp"] < max_age:
+                    return cached["data"]
+    except Exception as e:
+        print(f"Error loading from cache: {e}")
+    return None
+
+def is_likely_bird_image(img):
+    """Basic check to see if the image might contain a bird"""
+    try:
+        # Convert to numpy array for analysis
+        img_array = np.array(img)
+        
+        # Simple checks that might indicate a bird isn't present:
+        # 1. Check if image is too dark or too bright overall
+        mean_brightness = np.mean(img_array)
+        if mean_brightness < 20 or mean_brightness > 235:
+            return False
+            
+        # 2. Check if image has very little color variation (might be a solid background)
+        std_dev = np.std(img_array)
+        if std_dev < 15:
+            return False
+            
+        # 3. If image is very small, it might not be a useful bird photo
+        if img_array.shape[0] < 100 or img_array.shape[1] < 100:
+            return False
+            
+        return True
+    except:
+        # If any error occurs during the check, assume it might be a bird
+        return True
+
 def get_bird_habitat_map(bird_name, check_tanzania=True):
-    """Get habitat map locations for the bird using Groq API"""
+    """Get habitat map locations for the bird using Groq API with caching"""
     clean_name = clean_bird_name(bird_name)
     
-    if check_tanzania:
-        tanzania_check_prompt = f"""
-        Is the {clean_name} bird native to or commonly found in Tanzania? 
-        Answer with ONLY "yes" or "no".
-        """
-        try:
-            tanzania_check = client.chat.completions.create(
-                messages=[{"role": "user", "content": tanzania_check_prompt}],
-                model="llama-3.3-70b-versatile",
-            )
-            is_in_tanzania = "yes" in tanzania_check.choices[0].message.content.lower()
-        except:
-            is_in_tanzania = True
+    # Check cache for Tanzania check
+    tanzania_cache_key = f"{clean_name}_tanzania"
+    cached_tanzania = load_from_cache("tanzania_check", tanzania_cache_key)
+    if cached_tanzania is not None:
+        is_in_tanzania = cached_tanzania
     else:
-        is_in_tanzania = True
+        # First check if the bird is endemic to Tanzania
+        if check_tanzania:
+            tanzania_check_prompt = f"""
+            Is the {clean_name} bird native to or commonly found in Tanzania? 
+            Answer with ONLY "yes" or "no".
+            """
+            
+            try:
+                tanzania_check = client.chat.completions.create(
+                    messages=[{"role": "user", "content": tanzania_check_prompt}],
+                    model="llama-3.3-70b-versatile",
+                )
+                is_in_tanzania = "yes" in tanzania_check.choices[0].message.content.lower()
+                # Cache result
+                save_to_cache("tanzania_check", tanzania_cache_key, is_in_tanzania)
+            except:
+                # Default to showing Tanzania if we can't determine
+                is_in_tanzania = True
+        else:
+            is_in_tanzania = True
     
+    # Check cache for habitat locations
+    habitat_cache_key = f"{clean_name}_habitat"
+    cached_habitat = load_from_cache("habitat", habitat_cache_key)
+    if cached_habitat is not None:
+        return cached_habitat, is_in_tanzania
+    
+    # Now get the habitat locations
     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:
@@ -104,78 +193,120 @@ def get_bird_habitat_map(bird_name, check_tanzania=True):
     
     try:
         chat_completion = client.chat.completions.create(
-            messages=[{"role": "user", "content": prompt}],
+            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": "Primary habitat region", "lat": 0, "lon": 0, 
                  "description": "Could not retrieve specific habitat information for this bird."}
             ]
+        
+        # Cache the result
+        save_to_cache("habitat", habitat_cache_key, locations)
+        
         return locations, is_in_tanzania
+        
     except Exception as e:
         return [{"name": "Error retrieving data", "lat": 0, "lon": 0, 
                 "description": "Please try again or check your connection."}], False
 
 def create_habitat_map(habitat_locations):
     """Create a folium map with the habitat locations"""
+    # Find center point based on valid coordinates
     valid_coords = [(loc.get("lat", 0), loc.get("lon", 0)) 
                     for loc in habitat_locations 
                     if loc.get("lat", 0) != 0 or loc.get("lon", 0) != 0]
     
     if valid_coords:
+        # Calculate the average of the coordinates
         avg_lat = sum(lat for lat, _ in valid_coords) / len(valid_coords)
         avg_lon = sum(lon for _, lon in valid_coords) / len(valid_coords)
+        # Create map centered on the average coordinates
         m = folium.Map(location=[avg_lat, avg_lon], zoom_start=3)
     else:
+        # Default world map if no valid coordinates
         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 format_bird_info(raw_info, language="en"):
     """Improve the formatting of bird information"""
+    # Add proper line breaks between sections and ensure consistent heading levels
     formatted = raw_info
+    
+    # Get translation of warning text based on language
     warning_text = "NOT TYPICALLY FOUND IN TANZANIA"
     warning_translation = "HAPATIKANI SANA TANZANIA" if language == "sw" else warning_text
     
+    # Fix heading levels (make all main sections h3)
     formatted = re.sub(r'#+\s+' + warning_text, 
-                       f'<div class="alert alert-warning"><strong>⚠️ {warning_translation}</strong></div>', 
-                       formatted)
+                        f'<div class="alert alert-warning"><strong>⚠️ {warning_translation}</strong></div>', 
+                        formatted)
     
+    # Replace markdown headings with HTML headings for better control
     formatted = re.sub(r'#+\s+(.*)', r'<h3>\1</h3>', formatted)
+    
+    # Add paragraph tags for better spacing
     formatted = re.sub(r'\n\*\s+(.*)', r'<p>• \1</p>', formatted)
     formatted = re.sub(r'\n([^<\n].*)', r'<p>\1</p>', formatted)
     
+    # Remove any duplicate paragraph tags
     formatted = formatted.replace('<p><p>', '<p>')
     formatted = formatted.replace('</p></p>', '</p>')
+    
     return formatted
 
 def get_bird_info(bird_name, language="en"):
-    """Get detailed information about a bird using Groq API"""
+    """Get detailed information about a bird using Groq API with caching"""
     clean_name = clean_bird_name(bird_name)
     
-    lang_instruction = "" if language == "en" else " Provide your response in Swahili language."
+    # Check cache first
+    cache_key = f"{clean_name}_{language}"
+    cached_info = load_from_cache("bird_info", cache_key)
+    if cached_info is not None:
+        return cached_info
+    
+    # Adjust language for the prompt
+    lang_instruction = ""
+    if language == "sw":
+        lang_instruction = " Provide your response in Swahili language."
     
     prompt = f"""
     Provide detailed information about the {clean_name} bird, including:
@@ -192,130 +323,184 @@ def get_bird_info(bird_name, language="en"):
     
     try:
         chat_completion = client.chat.completions.create(
-            messages=[{"role": "user", "content": prompt}],
+            messages=[
+                {
+                    "role": "user",
+                    "content": prompt,
+                }
+            ],
             model="llama-3.3-70b-versatile",
         )
-        return chat_completion.choices[0].message.content
+        response = chat_completion.choices[0].message.content
+        # Cache the result
+        save_to_cache("bird_info", cache_key, response)
+        return response
     except Exception as e:
         error_msg = "Hitilafu katika kupata taarifa" if language == "sw" else "Error fetching information"
         return f"{error_msg}: {str(e)}"
 
+def create_message_html(message, icon="🔍", language="en"):
+    """Create a styled message container for notifications"""
+    custom_css = """
+    <style>
+    .message-container {
+        font-family: Arial, sans-serif;
+        padding:
+        20px;
+        background-color: #f8f9fa;
+        border-radius: 8px;
+        text-align: center;
+        margin: 20px 0;
+    }
+    .message-icon {
+        font-size: 48px;
+        margin-bottom: 15px;
+    }
+    .message-text {
+        font-size: 18px;
+        color: #495057;
+    }
+    </style>
+    """
+    
+    html = f"""
+    {custom_css}
+    <div class="message-container">
+        <div class="message-icon">{icon}</div>
+        <div class="message-text">{message}</div>
+    </div>
+    """
+    return html
+
 def predict_and_get_info(img, language="en"):
     """Predict bird species and get detailed information"""
+    # Get translations
     t = translations[language]
     
-    img = PILImage.create(img)
-    pred, pred_idx, probs = learn.predict(img)
+    # Check if an image was provided
+    if img is None:
+        message = t['upload_prompt']
+        return None, create_message_html(message, "📷", language), "", ""
     
-    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]
+    # Basic check if the image might contain a bird
+    if not is_likely_bird_image(img):
+        message = t['not_a_bird']
+        return None, create_message_html(message, "⚠️", language), "", ""
     
-    prediction_results = {clean_bird_name(top_labels[i]): float(top_probs[i]) for i in range(num_classes)}
-    
-    top_bird = str(pred)
-    clean_top_bird = clean_bird_name(top_bird)
-    
-    # Check if top prediction is "Other" with high confidence (>80%)
-    if top_bird.lower() == "other" and float(probs[pred_idx]) > 0.8:
-        message = (
-            "This image does not appear to match any bird species in our trained dataset. "
-            "Please upload a clear image of a bird for accurate identification."
-        ) if language == "en" else (
-            "Picha hii haionekani kulingana na spishi yoyote ya ndege katika seti yetu ya mafunzo. "
-            "Tafadhali pakia picha ya wazi ya ndege kwa utambuzi sahihi."
-        )
+    try:
+        # Process the image
+        img = PILImage.create(img)
+        
+        # Get prediction
+        pred, pred_idx, probs = learn.predict(img)
+        
+        # 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 with cleaned names for display
+        prediction_results = {clean_bird_name(top_labels[i]): float(top_probs[i]) for i in range(num_classes)}
+        
+        # Get top prediction (original format for info retrieval)
+        top_bird = str(pred)
+        # Also keep a clean version for display
+        clean_top_bird = clean_bird_name(top_bird)
+        
+        # Check if the model's confidence is low
+        if float(top_probs[0]) < 0.4:
+            low_confidence_warning = t['low_confidence']
+        else:
+            low_confidence_warning = ""
+        
+        # Check if the top prediction is "Other" and has high confidence
+        if "other" in clean_top_bird.lower():
+            # Create a message informing the user that the bird wasn't recognized
+            other_message = t['other_message']
+            combined_info = create_message_html(other_message, "🔍", language)
+            return prediction_results, combined_info, clean_top_bird, ""
+        
+        # Get habitat locations and create map
+        habitat_locations, is_in_tanzania = 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, language)
+        formatted_info = format_bird_info(bird_info, language)
+        
+        # Create combined info with map at the top and properly formatted information
         custom_css = """
         <style>
         .bird-container {
             font-family: Arial, sans-serif;
             padding: 10px;
         }
-        .alert-info {
-            background-color: #d9edf7;
-            border: 1px solid #bce8f1;
-            color: #31708f;
+        .map-container {
+            height: 400px;
+            width: 100%;
+            border: 1px solid #ddd;
+            border-radius: 8px;
+            overflow: hidden;
+            margin-bottom: 20px;
+        }
+        .info-container {
+            line-height: 1.6;
+        }
+        .info-container h3 {
+            margin-top: 20px;
+            margin-bottom: 10px;
+            color: #2c3e50;
+            border-bottom: 1px solid #eee;
+            padding-bottom: 5px;
+        }
+        .info-container p {
+            margin-bottom: 10px;
+        }
+        .alert {
             padding: 10px;
             margin-bottom: 15px;
             border-radius: 4px;
         }
+        .alert-warning {
+            background-color: #fcf8e3;
+            border: 1px solid #faebcc;
+            color: #8a6d3b;
+        }
+        .confidence-warning {
+            background-color: #fff3cd;
+            color: #856404;
+            padding: 8px;
+            border-radius: 4px;
+            margin-bottom: 15px;
+            font-weight: bold;
+        }
         </style>
         """
         
+        # Add low confidence warning if needed
+        confidence_warning_html = f'<div class="confidence-warning">{low_confidence_warning}</div>' if low_confidence_warning else ''
+        
         combined_info = f"""
         {custom_css}
         <div class="bird-container">
-            <div class="alert-info">
-                <strong>ℹ️ {message}</strong>
+            {confidence_warning_html}
+            <h2>{t['habitat_map_title']} {clean_top_bird}</h2>
+            <div class="map-container">
+                {habitat_map_html}
+            </div>
+            
+            <div class="info-container">
+                <h2>{t['detailed_info_title']}</h2>
+                {formatted_info}
             </div>
         </div>
         """
-        return prediction_results, combined_info, ""
-    
-    habitat_locations, is_in_tanzania = get_bird_habitat_map(top_bird)
-    habitat_map_html = create_habitat_map(habitat_locations)
-    
-    bird_info = get_bird_info(top_bird, language)
-    formatted_info = format_bird_info(bird_info, language)
-    
-    custom_css = """
-    <style>
-    .bird-container {
-        font-family: Arial, sans-serif;
-        padding: 10px;
-    }
-    .map-container {
-        height: 400px;
-        width: 100%;
-        border: 1px solid #ddd;
-        border-radius: 8px;
-        overflow: hidden;
-        margin-bottom: 20px;
-    }
-    .info-container {
-        line-height: 1.6;
-    }
-    .info-container h3 {
-        margin-top: 20px;
-        margin-bottom: 10px;
-        color: #2c3e50;
-        border-bottom: 1px solid #eee;
-        padding-bottom: 5px;
-    }
-    .info-container p {
-        margin-bottom: 10px;
-    }
-    .alert {
-        padding: 10px;
-        margin-bottom: 15px;
-        border-radius: 4px;
-    }
-    .alert-warning {
-        background-color: #fcf8e3;
-        border: 1px solid #faebcc;
-        color: #8a6d3b;
-    }
-    </style>
-    """
-    
-    combined_info = f"""
-    {custom_css}
-    <div class="bird-container">
-        <h2>{t['habitat_map_title']} {clean_top_bird}</h2>
-        <div class="map-container">
-            {habitat_map_html}
-        </div>
         
-        <div class="info-container">
-            <h2>{t['detailed_info_title']}</h2>
-            {formatted_info}
-        </div>
-    </div>
-    """
-    
-    return prediction_results, combined_info, clean_top_bird
+        return prediction_results, combined_info, clean_top_bird, ""
+    except Exception as e:
+        error_msg = "Hitilafu katika kuchakata picha" if language == "sw" else "Error processing image"
+        return None, create_message_html(f"{error_msg}: {str(e)}", "⚠️", language), "", ""
 
 def follow_up_question(question, bird_name, language="en"):
     """Allow researchers to ask follow-up questions about the identified bird"""
@@ -324,7 +509,16 @@ def follow_up_question(question, bird_name, language="en"):
     if not question.strip() or not bird_name:
         return "Please identify a bird first and ask a specific question about it." if language == "en" else "Tafadhali tambua ndege kwanza na uulize swali maalum kuhusu ndege huyo."
     
-    lang_instruction = "" if language == "en" else " Provide your response in Swahili language."
+    # Check cache first
+    cache_key = f"{bird_name}_{question}_{language}".replace(" ", "_")[:100]  # Limit key length
+    cached_answer = load_from_cache("follow_up", cache_key)
+    if cached_answer is not None:
+        return cached_answer
+    
+    # Adjust language for the prompt
+    lang_instruction = ""
+    if language == "sw":
+        lang_instruction = " Provide your response in Swahili language."
     
     prompt = f"""
     The researcher is asking about the {bird_name} bird: "{question}"
@@ -341,19 +535,29 @@ def follow_up_question(question, bird_name, language="en"):
     
     try:
         chat_completion = client.chat.completions.create(
-            messages=[{"role": "user", "content": prompt}],
+            messages=[
+                {
+                    "role": "user",
+                    "content": prompt,
+                }
+            ],
             model="llama-3.3-70b-versatile",
         )
-        return chat_completion.choices[0].message.content
+        response = chat_completion.choices[0].message.content
+        # Cache the result
+        save_to_cache("follow_up", cache_key, response)
+        return response
     except Exception as e:
         error_msg = "Hitilafu katika kupata taarifa" if language == "sw" else "Error fetching information"
         return f"{error_msg}: {str(e)}"
 
 # Create the Gradio interface
 with gr.Blocks(theme=gr.themes.Soft()) as app:
+    # Current language and bird state
     current_lang = gr.State("en")
     current_bird = gr.State("")
     
+    # Header with language switcher
     with gr.Row():
         with gr.Column(scale=3):
             title_md = gr.Markdown(f"# {translations['en']['app_title']}")
@@ -364,8 +568,10 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
                 value="English"
             )
     
+    # App description
     description_md = gr.Markdown(f"{translations['en']['app_description']}")
     
+    # Main identification section
     with gr.Row():
         with gr.Column(scale=1):
             input_image = gr.Image(type="pil", label=translations['en']['upload_label'])
@@ -375,8 +581,10 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
             prediction_output = gr.Label(label=translations['en']['predictions_label'], num_top_classes=5)
             bird_info_output = gr.HTML(label=translations['en']['bird_info_label'])
     
+    # Clear divider
     gr.Markdown("---")
     
+    # Follow-up question section with improved UI
     questions_header = gr.Markdown(f"## {translations['en']['research_questions']}")
     
     conversation_history = gr.Markdown("")
@@ -392,17 +600,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
         follow_up_btn = gr.Button(translations['en']['submit_question'], variant="primary")
         clear_btn = gr.Button(translations['en']['clear_conversation'])
     
-    def process_image(img, lang):
-        if img is None:
-            return None, translations[lang]['upload_prompt'], "", ""
-        
-        try:
-            pred_results, info, clean_bird_name = predict_and_get_info(img, lang)
-            return pred_results, info, clean_bird_name, ""
-        except Exception as e:
-            error_msg = "Hitilafu katika kuchakata picha" if lang == "sw" else "Error processing image"
-            return None, f"{error_msg}: {str(e)}", "", ""
-    
+    # Functions for event handlers
     def update_conversation(question, bird_name, history, lang):
         t = translations[lang]
         
@@ -411,6 +609,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
         
         answer = follow_up_question(question, bird_name, lang)
         
+        # Format the conversation with clear separation
         new_exchange = f"""
 ### {t['question_title']}
 {question}
@@ -425,9 +624,11 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
         return ""
     
     def update_language(choice):
+        # Convert selection to language code
         lang = "sw" if choice == "Kiswahili" else "en"
         t = translations[lang]
         
+        # Return updated UI components based on selected language
         return (
             lang, 
             f"# {t['app_title']}", 
@@ -443,6 +644,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
             t['clear_conversation']
         )
     
+    # Set up event handlers
     language_selector.change(
         update_language,
         inputs=[language_selector],
@@ -462,8 +664,13 @@ with gr.Blocks(theme=gr.themes.Soft()) as app:
         ]
     )
     
+    # Add loading state for better UX
     submit_btn.click(
-        process_image,
+        lambda x, y: (None, create_message_html(translations[y]['loading'], "⏳", y), "", ""),
+        inputs=[input_image, current_lang],
+        outputs=[prediction_output, bird_info_output, current_bird, conversation_history]
+    ).then(
+        predict_and_get_info,
         inputs=[input_image, current_lang],
         outputs=[prediction_output, bird_info_output, current_bird, conversation_history]
     )