zfbox/model/plugins/RYRQ_Model_ACL/RYRQ_Model_ACL.py

import os.path
from model.plugins.ModelBase import ModelBase
#from model.base_model.ascnedcl.det_utils import get_labels_from_txt, letterbox, scale_coords, nms, draw_bbox  # 模型前后处理相关函数
from model.base_model.ascnedcl.det_utils_v10 import draw_box,draw_bbox_old, letterbox,non_max_suppression_v10
import cv2
import numpy as np
import torch  # 深度学习运算框架，此处主要用来处理数据

class Model(ModelBase):
    def __init__(self,path,threshold=0.5,iou_thres=0.5):
        # 找pt模型路径  -- 一个约束py文件和模型文件的路径关系需要固定， -- 上传模型时，要解压好路径
        dirpath, filename = os.path.split(path)
        #self.model_file = os.path.join(dirpath, "yolov5s_bs1.om")  # 目前约束模型文件和py文件在同一目录
        self.model_file = os.path.join(dirpath, "yolov10m_310B4.om")  # 目前约束模型文件和py文件在同一目录
        self.coco_file = os.path.join(dirpath, "coco_names.txt")
        #self.labels_dict = get_labels_from_txt(self.coco_file)  # 得到类别信息，返回序号与类别对应的字典

        super().__init__(self.model_file)  # acl环境初始化基类负责类的实例化

        self.name = "人员入侵-yolov10"
        self.version = "V1.0"
        self.model_type = 2

        self.neth = 640  # 缩放的目标高度, 也即模型的输入高度
        self.netw = 640  # 缩放的目标宽度, 也即模型的输入宽度

        self.conf_threshold = threshold  # 置信度阈值
        self.iou_thres = iou_thres      #IOU阈值

    def prework(self,image):
        '''模型输入图片数据前处理 --- 针对每个模型特有的预处理内容 -'''
        img, scale_ratio, dw,dh = letterbox(image, new_shape=[self.netw, self.neth])  # 对图像进行缩放与填充
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, HWC to CHW     #图片在输入时已经做了转换
        img = np.ascontiguousarray(img, dtype=np.float32) / 255.0  # 转换为内存连续存储的数组
        return img,scale_ratio, (dw,dh)

    def postwork(self,image,outputs,scale_ratio,pad_size,check_area,polygon,conf_threshold,iou_thres):
        '''
        针对每个模型特有的后处理内容
        :param image:
        :param outputs:
        :param scale_ratio:
        :param pad_size:
        :param check_area:
        :param polygon:
        :param conf_threshold:
        :param iou_thres:
        :return:
        '''
        filtered_pred_all = None
        bwarn = False
        warn_text = ""
        # 是否有检测区域，有先绘制检测区域 由于在该函数生成了polygon对象，所有需要在检测区域前调用。
        if check_area == 1:
            self.draw_polygon(image, polygon, (255, 0, 0))

        if outputs:
            output = np.squeeze(outputs[0])  # 移除张量为1的维度 --暂时不明白其具体意义
            dw,dh = pad_size
            pred_all = non_max_suppression_v10(output, self.conf_threshold, scale_ratio, dw, dh)
            for xmin, ymin, xmax, ymax, confidence, label in pred_all:
                # # 绘制目标识别的锚框  --已经在draw_bbox里处理
                # cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
                draw_box(image, [xmin, ymin, xmax, ymax], confidence, label)  # 画出检测框、类别、概率
                if label == 0:  # person
                    # 判断是否产生告警
                    x1, y1, x2, y2 = int(xmin), int(ymin), int(xmax), int(ymax)
                    if check_area == 1:  # 指定了检测区域
                        x_center = (x1 + x2) / 2
                        y_center = (y1 + y2) / 2
                        # 绘制中心点？
                        cv2.circle(image, (int(x_center), int(y_center)), 5, (0, 0, 255), -1)
                        # 判断是否区域点
                        if not self.is_point_in_region((x_center, y_center)):
                            continue  # 没产生报警-继续
                    # 产生报警 -- 有一个符合即可
                    bwarn = True
                    warn_text = "People Intruder detected!"
            # 清理内存
            del outputs, output
            del pred_all, filtered_pred_all
        # cv2.imwrite('img_res.png', img_dw)
        return bwarn, warn_text

    def verify(self,image,data,isdraw=1):
        # 数据前处理
        img, scale_ratio, dw,dh = letterbox(image, new_shape=[self.netw, self.neth])  # 对图像进行缩放与填充
        img = img[:, :, ::-1].transpose(2, 0, 1)  # BGR to RGB, HWC to CHW     #图片在输入时已经做了转换
        img = np.ascontiguousarray(img, dtype=np.float32) / 255.0  # 转换为内存连续存储的数组

        # 模型推理, 得到模型输出
        outputs = self.execute([img,])#创建input，执行模型，返回结果 --失败返回None

        filtered_pred_all = None
        bwarn = False
        warn_text = ""
        # 是否有检测区域，有先绘制检测区域 由于在该函数生成了polygon对象，所有需要在检测区域前调用。
        if data[1] == 1:
            self.draw_polygon(image, data[2], (255, 0, 0))

        if outputs:
            output = np.squeeze(outputs[0])     #移除张量为1的维度 --暂时不明白其具体意义
            pred_all = non_max_suppression_v10(output,self.conf_threshold,scale_ratio,dw,dh)

            for xmin, ymin, xmax, ymax, confidence, label in pred_all:
                # # 绘制目标识别的锚框  --已经在draw_bbox里处理
                # cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
                draw_box(image, [xmin, ymin, xmax, ymax], confidence, label)  # 画出检测框、类别、概率
                if label == 0: # person
                    #判断是否产生告警
                    x1, y1, x2, y2 = int(xmin), int(ymin), int(xmax), int(ymax)
                    if data[1] == 1:  # 指定了检测区域
                        x_center = (x1 + x2) / 2
                        y_center = (y1 + y2) / 2
                        # 绘制中心点？
                        cv2.circle(image, (int(x_center), int(y_center)), 5, (0, 0, 255), -1)
                        # 判断是否区域点
                        if not self.is_point_in_region((x_center, y_center)):
                            continue  # 没产生报警-继续
                    # 产生报警 -- 有一个符合即可
                    bwarn = True
                    warn_text = "People Intruder detected!"
            # 清理内存
            del outputs, output
            del pred_all,filtered_pred_all
        #cv2.imwrite('img_res.png', img_dw)
        return bwarn, warn_text

    def testRun(self):
        print("I am  RYRQ-Model-ACL!!!")