Update app.py

app.py

@@ -1,17 +1,19 @@
+
 import os
 import gradio as gr
 import cv2
 import numpy as np
 from ultralytics import YOLO
 import tempfile
+import torch
 
 # Set Ultralytics config path
 os.environ['YOLO_CONFIG_DIR'] = '/tmp/Ultralytics'
 
-# Custom function to load model with weights_only=False
+# Custom function to load model with trusted weights
 def load_model_with_trusted_weights(model_path):
-    import torch
-    with torch.serialization.safe_globals():
+    import ultralytics.nn.tasks
+    with torch.serialization.safe_globals([ultralytics.nn.tasks.DetectionModel]):
         return YOLO(model_path)
 
 # Load both YOLO models