Update app.py

app.py

@@ -1,39 +1,34 @@
-import torch
-import torchvision.transforms as transforms
-from PIL import Image
+from ultralytics import YOLO
 import gradio as gr
+from PIL import Image
+import torch
 
-# Load model
-model = torch.load("best.pt", map_location=torch.device('cpu'))
-model.eval()
+# Load YOLOv5 model
+model = YOLO("best.pt")  # assumes model is in the same directory
 
-# Define class names (update these based on model's training)
-classes = [
-    "Apple___Apple_scab", "Apple___Black_rot", "Apple___Cedar_apple_rust",
-    "Apple___healthy", "Blueberry___healthy", "Cherry_(including_sour)___Powdery_mildew",
-    "Cherry_(including_sour)___healthy", "Corn_(maize)___Cercospora_leaf_spot Gray_leaf_spot",
-    # ...continue all 39 classes...
-]
+# Prediction function
+def predict(image):
+    results = model(image)
+    boxes = results[0].boxes
+    names = model.names
+    output = ""
 
-# Image transforms
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-])
+    if len(boxes) == 0:
+        output = "✅ Healthy: No disease detected!"
+    else:
+        for box in boxes:
+            class_id = int(box.cls[0])
+            class_name = names[class_id]
+            confidence = float(box.conf[0])
+            output += f"🚨 Detected: {class_name} ({confidence*100:.2f}%)\n"
 
-# Prediction function
-def predict(img):
-    img = transform(img).unsqueeze(0)
-    with torch.no_grad():
-        outputs = model(img)
-        _, predicted = torch.max(outputs, 1)
-        label = classes[predicted.item()]
-    return f"Prediction: {label}"
+    return output
 
-# Gradio Interface
-gr.Interface(fn=predict,
-             inputs=gr.Image(type="pil"),
-             outputs="text",
-             title="🌿 AgroScan: Plant Disease Detector",
-             description="Upload a plant leaf image to detect disease."
+# Gradio UI
+gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs="text",
+    title="🌿 AgroScan: Plant Disease Detector",
+    description="Upload a plant leaf image to detect disease using a YOLOv5 model.",
 ).launch()