zfbox/core/ModelNode.py

import threading
import importlib.util
import time
from myutils.MyDeque import MyDeque
from myutils.ConfigManager import myCongif
from myutils.MyLogger_logger import LogHandler
from core.ACLModelManager import ACLModeManger
from core.DataStruct import ModelinData,ModeloutData
from threading import Lock

class ModelNode:
    def __init__(self,device,model_path):
        self.device = device
        self.model_path = model_path
        self.model = None       #模型对象
        self.model_th = None    #模型线程句柄
        self.brun = True        #模型控制标识
        self.model_th_status = 0    #模型线程运行状态 0--初始状态，1-线程执行成功，2-线程退出
        self.in_mq = MyDeque(50)    #
        self.channel_list = {}   #channel_id  out_mq --需要线程安全
        self.clist_Lock = Lock()    #channel_list的维护锁
        self.ch_count = 0       #关联启动的通道数量
        self.count_Lock = Lock()    #count的维护锁
        self.model_platform = myCongif.get_data("model_platform")
        self.logger = LogHandler().get_logger("ModelNode")



    def __del__(self):
        pass

    def _reset(self): #重置数据
        #self.model_th_status = 0  # 模型线程运行状态 0--初始状态，1-线程执行成功，2-线程退出
        self.in_mq.myclear()

    def _import_model(self,model_path,threshold=0.5,iou_thres=0.5):
        '''
        根据路径，动态导入模块
        :param model_path: 模块路径
        :param threshold:   置信阈值
        :param iou_thres:   iou阈值
        :return:
        '''
        try:
            module_path = model_path.replace("/", ".").rsplit(".", 1)[0]
            print(module_path)
            # 动态导入模块
            module = importlib.import_module(module_path)
            # 从模块中获取指定的类
            Model = getattr(module, "Model")
            # 使用 Model 类
            model_instance = Model(model_path,threshold,iou_thres)
            return model_instance
        except ModuleNotFoundError as e:
            print(f"Module not found: {e}")
            return None
        except AttributeError as e:
            print(f"Class not found in module: {e}")
            return None
        except Exception as e:
            print(f"An unexpected error occurred: {e}")
            return None

    def _model_th(self):
        # 加载自定义模型文件
        self.model = self._import_model(self.model_path)  # 动态加载模型处理文件py
        if not self.model:
            self.logger.error("自定义模型文件加载失败，退出model线程")
            self.model_th_status = 2
            return
        # 初始化模型运行资源
        context = None
        if self.model_platform == "acl":  # ACL线程中初始化内容
            context = ACLModeManger.th_inti_acl(self.device)  # 创建context
            # 初始化模型资源 -- 加载模型文件
            ret = self.model.init_acl_resource()  # 加载和初始化离线模型文件--om文件
            if not ret:
                print("初始化模型资源出错，退出线程！")
                self.model_th_status = 2
                return
        #执行工作
        self.model_th_status = 1
        while self.brun:
            inData = self.in_mq.mypopleft() #空时，返回None  #(self,channel_id,img,image,scale_ratio, pad_size):
            if inData:
                outputs = self.model.execute([inData.img,])#创建input，执行模型，返回结果 --失败返回None
                outdata = ModeloutData(inData.image,inData.scale_ratio,inData.pad_size,outputs)
                del inData.img
                with self.clist_Lock:
                    if inData.channel_id in self.channel_list:
                        self.channel_list[inData.channel_id].myappend(outdata)
            else:
                time.sleep(0.05)

        #结束线程，释放资源
        self.model_th_status = 0
        self._reset()
        # 反初始化
        if self.model_platform == "acl":
            try:
                self.model.release()  # 释放模型资源资源
                # 删除模型对象
                del self.model
                # 释放context
                if context:  # ACL线程中反初始化内容  -- 若线程异常退出，这些资源就不能正常释放了
                    # 再释放context
                    ACLModeManger.th_del_acl(context)
            except Exception as e:
                print(e)

    def start_model_th(self,channel_id,out_mq):
        with self.count_Lock:
            with self.clist_Lock:
                if channel_id in self.channel_list:
                    return  #这个可以删除老的，新增新的
                self.channel_list[channel_id] = out_mq
            if self.ch_count == 0:  #第一次启动线程
                self.brun = True
                self.model_th = threading.Thread(target=self._model_th)
                self.model_th.start()
            self.ch_count += 1      #有通道调用一次就加一

    def stop_model_th(self,channel_id):
        with self.count_Lock:
            with self.clist_Lock:
                if channel_id in self.channel_list:
                    del self.channel_list[channel_id]
            self.ch_count -= 1
            if self.ch_count == 0:   #所有通道结束
                self.brun = False
                self.model_th.join()
                self.model_th = None