app.py

import os
import gradio as gr
import re
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')
labels = learn.dls.vocab

# Initialize Groq client
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": {
        "app_title": "AvianEye Tanzania",
        "app_description": "🔍 Upload a bird photo to instantly identify species and access comprehensive data on habitats, behaviors, and climate change impacts. A powerful tool for Tanzania-based ornithological research.",
        "upload_label": "Upload Bird Image",
        "identify_button": "Identify Bird",
        "predictions_label": "Top 5 Predictions",
        "bird_info_label": "Bird Information",
        "research_questions": "Research Questions",
        "question_placeholder": "Example: How has climate change affected this bird's migration pattern?",
        "question_label": "Ask a question about this bird",
        "submit_question": "Submit Question",
        "clear_conversation": "Clear Conversation",
        "upload_prompt": "Please upload an image",
        "question_title": "Question:",
        "answer_title": "Answer:",
        "habitat_map_title": "Natural Habitat Map for",
        "detailed_info_title": "Detailed Information",
        "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",
        "app_description": "🔍 Pakia picha ya ndege ili kutambua spishi mara moja na kupata data kamili kuhusu makazi, tabia, na athari za mabadiliko ya tabianchi. Zana yenye nguvu kwa utafiti wa ndege nchini Tanzania.",
        "upload_label": "Pakia Picha ya Ndege",
        "identify_button": "Tambua Ndege",
        "predictions_label": "Utabiri Bora 5",
        "bird_info_label": "Taarifa za Ndege",
        "research_questions": "Maswali ya Utafiti",
        "question_placeholder": "Mfano: Je, mabadiliko ya tabianchi yameathiri vipi mfumo wa uhamiaji wa ndege huyu?",
        "question_label": "Uliza swali kuhusu ndege huyu",
        "submit_question": "Wasilisha Swali",
        "clear_conversation": "Futa Mazungumzo",
        "upload_prompt": "Tafadhali pakia picha",
        "question_title": "Swali:",
        "answer_title": "Jibu:",
        "habitat_map_title": "Ramani ya Makazi Asilia ya",
        "detailed_info_title": "Taarifa za Kina",
        "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 with caching"""
    clean_name = clean_bird_name(bird_name)
    
    # 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:
        # 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:
    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 is_in_tanzania else 'DO NOT include any locations in Tanzania as this bird is not native to or commonly found 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": "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)
    
    # 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 with caching"""
    clean_name = clean_bird_name(bird_name)
    
    # 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:
    1. Physical characteristics and appearance
    2. Habitat and distribution
    3. Diet and behavior
    4. Migration patterns (emphasize if this pattern has changed in recent years due to climate change)
    5. Conservation status
    
    If this bird is not commonly found in Tanzania, explicitly flag that this bird is "NOT TYPICALLY FOUND IN TANZANIA" at the beginning of your response and explain why its presence might be unusual.
    
    Format your response in markdown for better readability.{lang_instruction}
    """
    
    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
        # 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]
    
    # Check if an image was provided
    if img is None:
        message = t['upload_prompt']
        return None, create_message_html(message, "📷", language), "", ""
    
    # 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), "", ""
    
    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;
        }
        .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">
            {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, 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"""
    t = translations[language]
    
    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."
    
    # 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}"
    
    Provide a detailed, scientific answer focusing on accurate ornithological information. 
    If the question relates to Tanzania or climate change impacts, emphasize those aspects in your response.
    
    IMPORTANT: Do not repeat basic introductory information about the bird that would have already been provided in a general description.
    Do not start your answer with phrases like "Introduction to the {bird_name}" or similar repetitive headers.
    Directly answer the specific question asked.
    
    Format your response in markdown for better readability.{lang_instruction}
    """
    
    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
        # 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']}")
        with gr.Column(scale=1):
            language_selector = gr.Radio(
                choices=["English", "Kiswahili"], 
                label=translations['en']['language_label'], 
                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'])
            submit_btn = gr.Button(translations['en']['identify_button'], variant="primary")
        
        with gr.Column(scale=2):
            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("")
    
    with gr.Row():
        follow_up_input = gr.Textbox(
            label=translations['en']['question_label'], 
            placeholder=translations['en']['question_placeholder'],
            lines=2
        )
    
    with gr.Row():
        follow_up_btn = gr.Button(translations['en']['submit_question'], variant="primary")
        clear_btn = gr.Button(translations['en']['clear_conversation'])
    
    # Functions for event handlers
    def update_conversation(question, bird_name, history, lang):
        t = translations[lang]
        
        if not question.strip():
            return history
        
        answer = follow_up_question(question, bird_name, lang)
        
        # Format the conversation with clear separation
        new_exchange = f"""
### {t['question_title']}
{question}
### {t['answer_title']}
{answer}
---
"""
        updated_history = new_exchange + history
        return updated_history
    
    def clear_conversation_history():
        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']}", 
            f"{t['app_description']}", 
            t['upload_label'],
            t['identify_button'],
            t['predictions_label'],
            t['bird_info_label'],
            f"## {t['research_questions']}",
            t['question_label'],
            t['question_placeholder'],
            t['submit_question'],
            t['clear_conversation']
        )
    
    # Set up event handlers
    language_selector.change(
        update_language,
        inputs=[language_selector],
        outputs=[
            current_lang,
            title_md,
            description_md,
            input_image,
            submit_btn,
            prediction_output,
            bird_info_output,
            questions_header,
            follow_up_input,
            follow_up_input,
            follow_up_btn,
            clear_btn
        ]
    )
    
    # Add loading state for better UX
    submit_btn.click(
        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]
    )
    
    follow_up_btn.click(
        update_conversation,
        inputs=[follow_up_input, current_bird, conversation_history, current_lang],
        outputs=[conversation_history]
    ).then(
        lambda: "", 
        outputs=follow_up_input
    )
    
    clear_btn.click(
        clear_conversation_history,
        outputs=[conversation_history]
    )

# Launch the app
app.launch(share=True)