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

# 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"),
)

# Language translations
translations = {
    "en": {
        "app_title": "Bird Species Explorer: AI-Powered Identification for Researchers",
        "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"
    },
    "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"
    }
}

def clean_bird_name(name):
    """Clean bird name by removing numbers and special characters, and fix formatting"""
    cleaned = re.sub(r'^\d+\.', '', name)
    cleaned = cleaned.replace('_', ' ')
    cleaned = re.sub(r'[^\w\s]', '', cleaned)
    cleaned = ' '.join(cleaned.split())
    return cleaned

def get_bird_habitat_map(bird_name, check_tanzania=True):
    """Get habitat map locations for the bird using Groq API"""
    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
    else:
        is_in_tanzania = True
    
    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
        
        json_match = re.search(r'\[.*\]', response, re.DOTALL)
        if json_match:
            locations = json.loads(json_match.group())
        else:
            locations = [
                {"name": "Primary habitat region", "lat": 0, "lon": 0, 
                 "description": "Could not retrieve specific habitat information for this bird."}
            ]
        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"""
    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:
        avg_lat = sum(lat for lat, _ in valid_coords) / len(valid_coords)
        avg_lon = sum(lon for _, lon in valid_coords) / len(valid_coords)
        m = folium.Map(location=[avg_lat, avg_lon], zoom_start=3)
    else:
        m = folium.Map(location=[20, 0], zoom_start=2)
    
    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")
        
        if lat == 0 and lon == 0:
            continue
            
        folium.Marker(
            location=[lat, lon],
            popup=folium.Popup(f"<b>{name}</b><br>{description}", max_width=300),
            tooltip=name
        ).add_to(m)
    
    map_html = m._repr_html_()
    return map_html

def format_bird_info(raw_info, language="en"):
    """Improve the formatting of bird information"""
    formatted = raw_info
    warning_text = "NOT TYPICALLY FOUND IN TANZANIA"
    warning_translation = "HAPATIKANI SANA TANZANIA" if language == "sw" else warning_text
    
    formatted = re.sub(r'#+\s+' + warning_text, 
                       f'<div class="alert alert-warning"><strong>⚠️ {warning_translation}</strong></div>', 
                       formatted)
    
    formatted = re.sub(r'#+\s+(.*)', r'<h3>\1</h3>', formatted)
    formatted = re.sub(r'\n\*\s+(.*)', r'<p>• \1</p>', formatted)
    formatted = re.sub(r'\n([^<\n].*)', r'<p>\1</p>', formatted)
    
    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"""
    clean_name = clean_bird_name(bird_name)
    
    lang_instruction = "" if language == "en" else " 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",
        )
        return chat_completion.choices[0].message.content
    except Exception as e:
        error_msg = "Hitilafu katika kupata taarifa" if language == "sw" else "Error fetching information"
        return f"{error_msg}: {str(e)}"

def predict_and_get_info(img, language="en"):
    """Predict bird species and get detailed information"""
    t = translations[language]
    
    img = PILImage.create(img)
    pred, pred_idx, probs = learn.predict(img)
    
    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]
    
    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."
        )
        
        custom_css = """
        <style>
        .bird-container {
            font-family: Arial, sans-serif;
            padding: 10px;
        }
        .alert-info {
            background-color: #d9edf7;
            border: 1px solid #bce8f1;
            color: #31708f;
            padding: 10px;
            margin-bottom: 15px;
            border-radius: 4px;
        }
        </style>
        """
        
        combined_info = f"""
        {custom_css}
        <div class="bird-container">
            <div class="alert-info">
                <strong>ℹ️ {message}</strong>
            </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

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."
    
    lang_instruction = "" if language == "en" else " 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",
        )
        return chat_completion.choices[0].message.content
    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_lang = gr.State("en")
    current_bird = gr.State("")
    
    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"
            )
    
    description_md = gr.Markdown(f"{translations['en']['app_description']}")
    
    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'])
    
    gr.Markdown("---")
    
    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'])
    
    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)}", "", ""
    
    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)
        
        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):
        lang = "sw" if choice == "Kiswahili" else "en"
        t = translations[lang]
        
        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']
        )
    
    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
        ]
    )
    
    submit_btn.click(
        process_image,
        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)