Update detector.py

detector.py

@@ -1,17 +1,20 @@
 from ultralytics import YOLO
 import cv2
 import torch
-import torch.serialization
-import ultralytics  # Required to access ultralytics-specific classes
 
 class LBWDetector:
     def __init__(self, model_path='best.pt'):
-        """Initialize YOLO model with safe globals for PyTorch 2.6+."""
-        with torch.serialization.safe_globals([torch.nn.modules.container.Sequential, ultralytics.nn.tasks.DetectionModel, ultralytics.nn.modules.conv.Conv, torch.nn.modules.conv.Conv2d]):
-            self.model = YOLO(model_path)
+        # Temporarily override torch.load to use weights_only=False
+        original_load = torch.load
+        def custom_load(*args, **kwargs):
+            kwargs['weights_only'] = False
+            return original_load(*args, **kwargs)
+        torch.load = custom_load
+        self.model = YOLO(model_path)
+        # Restore original torch.load
+        torch.load = original_load
 
     def detect_objects(self, frame):
-        """Detect objects in a frame and return bounding boxes and class names."""
         results = self.model.predict(source=frame, conf=0.3, save=False, verbose=False)
         detections = results[0].boxes.data.cpu().numpy()  # x1, y1, x2, y2, conf, class
         return detections, results[0].names