from ultralytics import YOLO


class YoloModel:

    def __init__(self, pt_url=None):
        if pt_url:
            self.model = YOLO(pt_url)
        else:
            self.model = YOLO('yolo11n.pt')

    def predict(self, source):
        """
        对rtsp视频流进行预测
        """
        results = self.model.predict(source=source, conf=0.6, device='cpu')
        return results