Update app.py

app.py

@@ -8,23 +8,23 @@ import sys
 import os
 
 import sys
-import torch
-import torch.serialization
 from pathlib import Path
+import torch
+from torch.serialization import add_safe_globals, safe_globals
 
-# Add yolov10 repo to Python path
+# Add yolov10 directory to Python path
 sys.path.append(str(Path(__file__).resolve().parent / "yolov10"))
 
-# Now import YOLOv10 classes
+# Now import from yolov10
 from ultralytics import YOLO
 from ultralytics.nn.tasks import YOLOv10DetectionModel
 
-# Register YOLOv10DetectionModel to allow safe unpickling
-torch.serialization.add_safe_globals({'YOLOv10DetectionModel': YOLOv10DetectionModel})
+# Register the actual YOLOv10DetectionModel class for safe loading
+add_safe_globals([YOLOv10DetectionModel])
 
-# Load the model safely
-with torch.serialization.safe_globals({'YOLOv10DetectionModel': YOLOv10DetectionModel}):
-    model = YOLO('/home/user/app/best.pt')  # Path to your trained YOLOv10n model
+# Load the model safely using context manager
+with safe_globals([YOLOv10DetectionModel]):
+    model = YOLO('/home/user/app/best.pt')  # Make sure this file exists
 
     
 # Frame processing function