Update app.py

app.py

@@ -1,24 +1,30 @@
+import torch
+from ultralytics.nn.tasks import DetectionModel
+
+# Allow the DetectionModel global for safe unpickling (only if you trust the source!)
+torch.serialization.add_safe_globals([DetectionModel])
+
 from huggingface_hub import hf_hub_download
 from ultralytics import YOLO
 import gradio as gr
 
 # Download the weights file (e.g., "best.pt") from the Hugging Face Hub
 weights_path = hf_hub_download(
-    repo_id="foduucom/plant-leaf-detection-and-classification", 
+    repo_id="foduucom/plant-leaf-detection-and-classification",
     filename="best.pt"
 )
 
-# Load the model using the local weights file
+# Load the YOLO model using the local weights file
 model = YOLO(weights_path)
 
 def predict_leaves(image_path):
-    # Run the model prediction on the provided image path
+    # Run the prediction on the provided image path
     results = model.predict(source=image_path, save=True)
-    # Count the number of detected leaves (using the first result)
+    # Count the number of detected leaves (assuming results[0].boxes is available)
     count = len(results[0].boxes)
     return f"Detected leaves: {count}"
 
-# Create a Gradio interface
+# Create a Gradio interface for your app
 iface = gr.Interface(
     fn=predict_leaves,
     inputs=gr.Image(type="filepath"),