Spaces:
Sleeping
Sleeping
from ultralytics import YOLO | |
import cv2 | |
import torch | |
import torch.serialization | |
class LBWDetector: | |
def __init__(self, model_path='best.pt'): | |
with torch.serialization.safe_globals([torch.nn.modules.container.Sequential]): | |
self.model = YOLO(model_path) | |
def detect_objects(self, frame): | |
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 |