import google.generativeai as genai
import gradio as gr
from PIL import Image
import os

# Set up Gemini API key
GEMINI_API_KEY = os.getenv('GEMINI_API_KEY')
genai.configure(api_key="GEMINI_API_KEY")  # Replace with your API key
def predict_rat(image):
    """Predict if the uploaded image contains a rat using Google Gemini Pro Vision API."""
    
    model = genai.GenerativeModel("learnlm-2.0-flash-experimental")

    # Open image using PIL
    img = Image.open(image)

    # Generate prediction
    response = model.generate_content(
        [img, "Is this an image of a rat? Answer with 'Yes' or 'No' and provide a brief explanation."]
    )
    
    # Extract prediction and explanation
    result = response.text
    return result

# Gradio UI
iface = gr.Interface(
    fn=predict_rat,
    inputs=gr.Image(type="filepath"),
    outputs="text",
    title="Rodent Detection App",
    description="Upload an image to check if it contains a rat."
)

# Run the app
if __name__ == "__main__":
    iface.launch()