Update app.py

app.py

@@ -1,28 +1,29 @@
 import torch
 from ultralytics.nn.tasks import DetectionModel
 from torch.nn.modules.container import Sequential
+from ultralytics.nn.modules import Conv  # Import Conv
 
-# Whitelist safe globals for unpickling (only do this if you trust the model checkpoint)
-torch.serialization.add_safe_globals([DetectionModel, Sequential])
+# Whitelist safe globals for unpickling (only if you trust the model checkpoint!)
+torch.serialization.add_safe_globals([DetectionModel, Sequential, Conv])
 
 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
+# Download the weights file (ensure the filename is correct)
 weights_path = hf_hub_download(
     repo_id="foduucom/plant-leaf-detection-and-classification",
     filename="best.pt"
 )
 
-# Load the YOLO model using the local weights file
+# Load the model using the local weights file
 model = YOLO(weights_path)
 
 def predict_leaves(image_path):
-    # Run prediction on the given image
+    # Run prediction on the provided image
     results = model.predict(source=image_path, save=True)
-    # Count the number of detected leaves using the first result
+    # Count the number of detected leaves (using the first result)
     count = len(results[0].boxes)
     return f"Detected leaves: {count}"