Spaces:
Sleeping
Sleeping
Update detector.py
Browse files- detector.py +3 -0
detector.py
CHANGED
|
@@ -1,8 +1,11 @@
|
|
| 1 |
from ultralytics import YOLO
|
| 2 |
import cv2
|
|
|
|
|
|
|
| 3 |
|
| 4 |
class LBWDetector:
|
| 5 |
def __init__(self, model_path='best.pt'):
|
|
|
|
| 6 |
self.model = YOLO(model_path)
|
| 7 |
|
| 8 |
def detect_objects(self, frame):
|
|
|
|
| 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):
|