dschandra commited on
Commit
94fd512
·
verified ·
1 Parent(s): f5561a2

Update detector.py

Browse files
Files changed (1) hide show
  1. detector.py +5 -6
detector.py CHANGED
@@ -1,15 +1,14 @@
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):
13
  results = self.model.predict(source=frame, conf=0.3, save=False, verbose=False)
14
  detections = results[0].boxes.data.cpu().numpy() # x1, y1, x2, y2, conf, class
15
- return detections, results[0].names
 
1
  from ultralytics import YOLO
 
 
2
  import cv2
3
+ import torch
4
+ import torch.serialization
5
 
6
  class LBWDetector:
7
  def __init__(self, model_path='best.pt'):
8
+ with torch.serialization.safe_globals([torch.nn.modules.container.Sequential]):
9
+ self.model = YOLO(model_path)
10
 
11
  def detect_objects(self, frame):
12
  results = self.model.predict(source=frame, conf=0.3, save=False, verbose=False)
13
  detections = results[0].boxes.data.cpu().numpy() # x1, y1, x2, y2, conf, class
14
+ return detections, results[0].names