from abc import ABC, abstractmethod
from ultralytics import YOLO
import yaml


class Predictor(ABC):
    @abstractmethod
    def predict(self, frame):
        pass


class OnnxPredictor(Predictor):
    def __init__(self):
        self.model = YOLO('azaion.onnx')
        self.model.task = 'detect'
        with open('data.yaml', 'r') as f:
            data_yaml = yaml.safe_load(f)
        class_names = data_yaml['names']

        names = self.model.names

    def predict(self, frame):
        results = self.model.track(frame, persist=True, tracker='bytetrack.yaml')
        return results[0].plot()


class YoloPredictor(Predictor):
    def __init__(self):
        self.model = YOLO('azaion.pt')

    def predict(self, frame):
        results = self.model.track(frame, persist=True, tracker='bytetrack.yaml')
        return results[0].plot()