Update app.py

app.py

@@ -3,12 +3,15 @@ import torch
 import cv2
 from ultralytics import YOLO
 
-# Add necessary safe globals to avoid the UnpicklingError
-torch.serialization.add_safe_globals([torch, "ultralytics.nn.tasks.DetectionModel"])
+# Safe load method to handle custom YOLO class during deserialization
+def safe_load_yolo_model(path):
+    # Add necessary safe globals to allow the detection model class during loading
+    torch.serialization.add_safe_globals([torch, 'ultralytics.nn.tasks.DetectionModel'])
+    return YOLO(path)
 
 # Load YOLO models
-model_yolo11 = YOLO('./data/yolo11n.pt')
-model_best = YOLO('./data/best.pt')
+model_yolo11 = safe_load_yolo_model('./data/yolo11n.pt')
+model_best = safe_load_yolo_model('./data/best.pt')
 
 def process_video(video):
     # Read video input