Update app.py

app.py

@@ -2,33 +2,42 @@ 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  # Import Conv
+from ultralytics.nn.modules import Conv
 
-# Whitelist the globals to bypass the pickle errors (only do this if you trust the model source!)
+# Whitelist safe globals (only do this if you trust the source of the model)
 torch.serialization.add_safe_globals([DetectionModel, Sequential, Conv])
 
-from ultralyticsplus import YOLO
+from ultralyticsplus import YOLO, render_result
 from PIL import Image
 
-# Load your custom YOLOv8 leaf detection model
+# Load the YOLOv8s Leaf Detection and Classification model from Hugging Face
 model = YOLO('foduucom/plant-leaf-detection-and-classification')
 
+# Set recommended model parameters
+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 image to a PIL Image and ensure it's in RGB
+    # Convert the input to a PIL image (ensuring RGB)
     image = Image.open(image).convert("RGB")
-    # Run inference
+    
+    # Perform inference with the model
     results = model.predict(image)
-    # Count the number of detected leaves
+    
+    # Count the detected leaves using the bounding boxes from the first result
     num_leaves = len(results[0].boxes)
+    
     return f"Number of leaves detected: {num_leaves}"
 
-# Gradio UI
+# Build a Gradio interface for the leaf counter
 iface = gr.Interface(
     fn=count_leaves,
     inputs=gr.Image(type="filepath"),
     outputs="text",
     title="Leaf Counter",
-    description="Upload an image of a plant, and the model will detect and count the number of leaves."
+    description="Upload an image of a plant and the model will detect and count the number of leaves."
 )
 
 if __name__ == "__main__":