Update app.py

app.py

@@ -1,31 +1,56 @@
 import gradio as gr
-from transformers import pipeline
+from ultralytics import YOLO
 from PIL import Image
+import numpy as np
+import cv2
 
-# Load the Hugging Face image classification pipeline with EfficientNetB0
-# This model is a general-purpose model for plant disease classification
-classifier = pipeline("image-classification", model="nateraw/efficientnet-b0")
+# Load the pre-trained YOLOv5 model
+# Use 'yolov5s' or any available YOLO model from ultralytics
+model = YOLO('yolov5s')  # Change to a leaf disease detection model if available
 
 def identify_disease(image):
-    # Use the classifier to predict the disease
-    predictions = classifier(image)
+    # Convert the image to RGB if it's not
+    if image.mode != 'RGB':
+        image = image.convert('RGB')
     
-    # Format the output to include disease name and confidence score
-    results = [{"Disease": pred["label"], "Confidence": f"{pred['score'] * 100:.2f}%"} for pred in predictions]
+    # Perform inference
+    results = model(image)
     
-    # Return the uploaded image along with the results
-    return image, results
+    # Extract predictions
+    predictions = results[0]
+    boxes = predictions.boxes
+    labels = boxes.cls.cpu().numpy()
+    scores = boxes.conf.cpu().numpy()
+    class_names = model.names
+    
+    # Annotate image with bounding boxes and labels
+    annotated_image = np.array(image)
+    for box, label, score in zip(boxes.xyxy.cpu().numpy(), labels, scores):
+        x1, y1, x2, y2 = map(int, box)
+        class_name = class_names[int(label)]
+        confidence = f"{score * 100:.2f}%"
+        annotated_image = cv2.rectangle(annotated_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
+        annotated_image = cv2.putText(annotated_image, f"{class_name} {confidence}", (x1, y1 - 10),
+                                      cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
+    
+    # Convert annotated image back to PIL format
+    annotated_image = Image.fromarray(annotated_image)
+    
+    # Prepare results for display
+    results_list = [{"Disease": class_names[int(label)], "Confidence": f"{score * 100:.2f}%"} for label, score in zip(labels, scores)]
+    
+    return annotated_image, results_list
 
 # Define Gradio interface
 interface = gr.Interface(
-    fn=identify_disease, 
+    fn=identify_disease,
     inputs=gr.inputs.Image(type="pil"),
     outputs=[
-        gr.outputs.Image(type="pil", label="Uploaded Image"),
-        gr.outputs.Dataframe(type="pandas", label="Predictions")
+        gr.outputs.Image(type="pil", label="Annotated Image"),
+        gr.outputs.Dataframe(headers=["Disease", "Confidence"], label="Predictions")
     ],
-    title="Plant Disease Identifier",
-    description="Upload an image of a plant leaf, and this tool will identify the disease along with the confidence score."
+    title="Leaf Disease Identifier with YOLOv5",
+    description="Upload an image of a leaf, and this tool will identify the disease with confidence scores."
 )
 
 # Launch the app