Update app.py

app.py

@@ -1,21 +1,23 @@
 import gradio as gr
 import torch
-from torch.nn.modules.container import Sequential  # import Sequential
-# Allow safe globals for the Sequential class
-torch.serialization.add_safe_globals([Sequential])
+from ultralytics.nn.tasks import DetectionModel
+from torch.nn.modules.container import Sequential  # Allow Sequential if needed
+
+# Whitelist the globals to bypass the pickle error (only do this if you trust the model source!)
+torch.serialization.add_safe_globals([DetectionModel, Sequential])
 
 from ultralyticsplus import YOLO
 from PIL import Image
 
-# Load your custom YOLOv8 leaf detection model.
+# Load your custom YOLOv8 leaf detection model
 model = YOLO('foduucom/plant-leaf-detection-and-classification')
 
 def count_leaves(image):
-    # Ensure the image is a PIL Image and convert to RGB
+    # Convert image to a PIL Image and ensure it's in RGB
     image = Image.open(image).convert("RGB")
     # Run inference
     results = model.predict(image)
-    # Count detected leaves
+    # Count the number of detected leaves
     num_leaves = len(results[0].boxes)
     return f"Number of leaves detected: {num_leaves}"