Update app.py

app.py

@@ -7,18 +7,26 @@ from PIL import Image
 import sys
 import os
 
+import sys
 import torch
 import torch.serialization
+from pathlib import Path
+
+# Add yolov10 repo to Python path
+sys.path.append(str(Path(__file__).resolve().parent / "yolov10"))
+
+# Now import YOLOv10 classes
 from ultralytics import YOLO
 from ultralytics.nn.tasks import YOLOv10DetectionModel
 
-# Register the custom class for safe loading
+# Register YOLOv10DetectionModel to allow safe unpickling
 torch.serialization.add_safe_globals({'YOLOv10DetectionModel': YOLOv10DetectionModel})
 
-# Load the model inside a safe context
+# Load the model safely
 with torch.serialization.safe_globals({'YOLOv10DetectionModel': YOLOv10DetectionModel}):
-    model = YOLO('/home/user/app/best.pt')  # Load your YOLOv10 model
+    model = YOLO('/home/user/app/best.pt')  # Path to your trained YOLOv10n model
 
+    
 # Frame processing function
 def process_frame(frame):
     img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)