app.py

import os
import gradio as gr
import re
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"),
)

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_bird_info(bird_name):
    """Get detailed information about a bird using Groq API"""
    clean_name = clean_bird_name(bird_name)
    
    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.
    """
    
    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:
        return f"Error fetching information: {str(e)}"

def predict_and_get_info(img):
    """Predict bird species and get detailed information"""
    # 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)
    
    # Get detailed information about the top predicted bird
    bird_info = get_bird_info(top_bird)
    
    return prediction_results, bird_info, clean_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}"
    
    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.
    """
    
    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:
        return f"Error fetching information: {str(e)}"

# Create the Gradio interface
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):
            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("")
    
    with gr.Row():
        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():
        follow_up_btn = gr.Button("Submit Question", variant="primary")
        clear_btn = gr.Button("Clear Conversation")
    
    # Set up event handlers
    def process_image(img):
        if img is None:
            return None, "Please upload an image", "", ""
        
        try:
            pred_results, info, clean_bird_name = predict_and_get_info(img)
            return pred_results, info, clean_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, conversation_history]
    )
    
    follow_up_btn.click(
        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
app.launch(share=True)