Update app.py

app.py

@@ -1,12 +1,14 @@
 import torch
 from ultralytics.nn.tasks import DetectionModel
+from torch.nn.modules.container import Sequential
 
-# Allow the DetectionModel global for safe unpickling (only if you trust the source!)
-torch.serialization.add_safe_globals([DetectionModel])
+# Whitelist safe globals for unpickling (only do this if you trust the model checkpoint)
+torch.serialization.add_safe_globals([DetectionModel, Sequential])
 
 from huggingface_hub import hf_hub_download
 from ultralytics import YOLO
 import gradio as gr
+from PIL import Image
 
 # Download the weights file (e.g., "best.pt") from the Hugging Face Hub
 weights_path = hf_hub_download(
@@ -18,9 +20,9 @@ weights_path = hf_hub_download(
 model = YOLO(weights_path)
 
 def predict_leaves(image_path):
-    # Run the prediction on the provided image path
+    # Run prediction on the given image
     results = model.predict(source=image_path, save=True)
-    # Count the number of detected leaves (assuming results[0].boxes is available)
+    # Count the number of detected leaves using the first result
     count = len(results[0].boxes)
     return f"Detected leaves: {count}"