Update app.py

app.py

@@ -1,44 +1,29 @@
+from ultralytics import YOLOvv8
 import gradio as gr
-import torch
-from ultralytics.nn.tasks import DetectionModel
-from torch.nn.modules.container import Sequential
-from ultralytics.nn.modules import Conv
-from torch.nn.modules.conv import Conv2d  # Import Conv2d
 
-# Whitelist the globals to bypass the pickle errors (only do this if you trust the source!)
-torch.serialization.add_safe_globals([DetectionModel, Sequential, Conv, Conv2d])
+# Load the pretrained model directly using YOLOvv8.from_pretrained
+model = YOLOvv8.from_pretrained("foduucom/plant-leaf-detection-and-classification")
 
-from ultralyticsplus import YOLO, render_result
-from PIL import Image
-
-# Load the YOLOv8s Leaf Detection and Classification model from Hugging Face
-model = YOLO('foduucom/plant-leaf-detection-and-classification')
-
-# Set recommended model parameters as per the model card
-model.overrides['conf'] = 0.25       # NMS confidence threshold
-model.overrides['iou'] = 0.45        # NMS IoU threshold
-model.overrides['agnostic_nms'] = False  # NMS class-agnostic setting
-model.overrides['max_det'] = 1000    # Maximum detections per image
-
-def count_leaves(image):
-    # Convert the input to a PIL image (ensuring RGB)
-    image = Image.open(image).convert("RGB")
-    
-    # Perform inference with the model
-    results = model.predict(image)
+def predict_leaves(image_path):
+    """
+    Given an image file path, run prediction on it using the YOLOvv8 model.
+    The model will automatically save the result if configured.
+    """
+    # Run the prediction; you can pass additional kwargs as needed (like save=True)
+    results = model.predict(source=image_path, save=True)
     
-    # Count the detected leaves using the bounding boxes from the first result
-    num_leaves = len(results[0].boxes)
+    # Optionally, count the detected leaves using the first result
+    count = len(results[0].boxes)
     
-    return f"Number of leaves detected: {num_leaves}"
+    return f"Detected leaves: {count}"
 
-# Build a Gradio interface for the leaf counter
+# Build a Gradio Interface for the leaf detection app
 iface = gr.Interface(
-    fn=count_leaves,
-    inputs=gr.Image(type="filepath"),
+    fn=predict_leaves,
+    inputs=gr.Image(type="filepath"),  # Users can upload an image file
     outputs="text",
-    title="Leaf Counter",
-    description="Upload an image of a plant and the model will detect and count the number of leaves."
+    title="Leaf Detection & Classification",
+    description="Upload an image to detect and count leaves using the YOLOvv8 model."
 )
 
 if __name__ == "__main__":