Update app.py

app.py

@@ -1,11 +1,45 @@
 import gradio as gr
-from transformers import pipeline
+import torch
+from torchvision import transforms
+from PIL import Image
+from transformers import AutoModelForImageClassification, AutoProcessor
 
-model_name = "Genius-Society/insecta"
-detection_pipeline = pipeline("object-detection", model=model_name)
+# Load model
+model_name = "ttheland/insect-identification"
+processor = AutoProcessor.from_pretrained(model_name)
+model = AutoModelForImageClassification.from_pretrained(model_name)
 
-def detect_objects(image):
-    return detection_pipeline(image)
+# 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}"
 
-iface = gr.Interface(fn=detect_objects, inputs="image", outputs="json")
-iface.launch()
+# 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()