import gradio as gr
import torch
from torchvision import transforms
from PIL import Image
from transformers import AutoModelForImageClassification, AutoProcessor

# Load model
model_name = "ttheland/insect-identification"
processor = AutoProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)

# Define inference function
def classify_insect(image):
    image = Image.fromarray(image)
    inputs = processor(images=image, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model(**inputs)
    
    predicted_class = outputs.logits.argmax().item()
    
    # Suggested commercial products for elimination
    insect_products = {
        "mosquito": "Talstar P Professional Insecticide",
        "ant": "Advion Ant Gel Bait",
        "roach": "Combat Max Roach Killing Gel",
        "termite": "Termidor SC Termiticide",
        "fly": "Noble Chemical Eliminate Flying Insect Killer"
    }
    
    product = insect_products.get(predicted_class, "No specific product found.")
    
    return f"Detected Insect Species: {predicted_class}\nRecommended Product: {product}"

# Gradio interface
iface = gr.Interface(
    fn=classify_insect,
    inputs="image",
    outputs="text",
    title="Insect Identification",
    description="Upload an insect image to identify its species and get elimination strategies."
)

if __name__ == "__main__":
    iface.launch()