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, pad_size = letterbox(image, new_shape=[640, 640]) # 对图像进行缩放与填充 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, pad_size 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: ''' bwarn = False warn_text = "" # 是否有检测区域，有先绘制检测区域 由于在该函数生成了polygon对象，所有需要在检测区域前调用。 if check_area == 1: self.draw_polygon(image, polygon, (255, 0, 0)) if outputs: output = outputs[0] # 只放了一张图片 -- #是否能批量验证？ # 后处理 -- boxout 是 tensor-list: [tensor([[],[].[]])] --[x1,y1,x2,y2,置信度,coco_index] # 利用非极大值抑制处理模型输出，conf_thres 为置信度阈值，iou_thres 为iou阈值 output_torch = torch.tensor(output) boxout = nms(output_torch, conf_thres=conf_threshold, iou_thres=iou_thres) del output_torch pred_all = boxout[0].numpy() # 转换为numpy数组 -- [[],[],[]] --[x1,y1,x2,y2,置信度,coco_index] # pred_all[:, :4] 取所有行的前4列，pred_all[:,1]--第一列 scale_coords([640, 640], pred_all[:, :4], image.shape, ratio_pad=(scale_ratio, pad_size)) # 将推理结果缩放到原始图片大小 # 过滤掉不是目标标签的数据 -- 序号0-- person self.labels_dict --这个考虑下是否可以放到nms前面#? filtered_pred_all = pred_all[pred_all[:, 5] == 0] # 绘制检测结果 --- 也需要封装在类里， for pred in filtered_pred_all: x1, y1, x2, y2 = int(pred[0]), int(pred[1]), int(pred[2]), int(pred[3]) # # 绘制目标识别的锚框 --已经在draw_bbox里处理 # cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2) 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!" draw_bbox(filtered_pred_all, image, (0, 255, 0), 2, self.labels_dict) # 画出检测框、类别、概率 # 清理内存 del outputs, output del boxout del pred_all, filtered_pred_all #图片绘制是在原图生效 image 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!!!")