dschandra commited on
Commit
83382ca
·
verified ·
1 Parent(s): 1b47d0b

Update detector.py

Browse files
Files changed (1) hide show
  1. detector.py +3 -2
detector.py CHANGED
@@ -1,11 +1,12 @@
1
  from ultralytics import YOLO
2
- import cv2
3
  from torch.serialization import add_safe_globals
4
  from ultralytics.nn.tasks import DetectionModel
 
 
 
5
 
6
  class LBWDetector:
7
  def __init__(self, model_path='best.pt'):
8
- add_safe_globals([DetectionModel])
9
  self.model = YOLO(model_path)
10
 
11
  def detect_objects(self, frame):
 
1
  from ultralytics import YOLO
 
2
  from torch.serialization import add_safe_globals
3
  from ultralytics.nn.tasks import DetectionModel
4
+ import cv2
5
+
6
+ add_safe_globals([DetectionModel]) # Allow loading YOLOv8 models with full class
7
 
8
  class LBWDetector:
9
  def __init__(self, model_path='best.pt'):
 
10
  self.model = YOLO(model_path)
11
 
12
  def detect_objects(self, frame):