Spaces:
Runtime error
Runtime error
| 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() |