from algo.model_manager import AlgoModelExec
from model_handler.base_model_handler import BaseModelHandler


class SafeEngineModelHandler(BaseModelHandler):

    def __init__(self, model: AlgoModelExec):
        super().__init__(model)
        self.model_names = {
            0: '人',
            1: '头',
            2: '安全帽',
            3: '摔倒',
            4: '烟头'
        }
        self.model_ids = list(self.model_names.keys())