You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
91 lines
4.5 KiB
91 lines
4.5 KiB
import os.path
|
|
from model.plugins.ModelBase import ModelBase
|
|
from myutils.ConfigManager import myCongif
|
|
from model.base_model.ascnedcl.det_utils import get_labels_from_txt, letterbox, scale_coords, nms, draw_bbox # 模型前后处理相关函数
|
|
import cv2
|
|
import numpy as np
|
|
import torch # 深度学习运算框架，此处主要用来处理数据
|
|
|
|
class Model(ModelBase):
|
|
def __init__(self,path,threshold=0.5):
|
|
# 找pt模型路径 -- 一个约束py文件和模型文件的路径关系需要固定， -- 上传模型时，要解压好路径
|
|
dirpath, filename = os.path.split(path)
|
|
self.model_file = os.path.join(dirpath, "yolov5s_bs1.om") # 目前约束模型文件和py文件在同一目录
|
|
self.coco_file = os.path.join(dirpath, "coco_names.txt")
|
|
super().__init__(self.model_file) #acl环境初始化基类负责类的实例化
|
|
self.model_id = None # 模型 id
|
|
self.input_dataset = None # 输入数据结构
|
|
self.output_dataset = None # 输出数据结构
|
|
self.model_desc = None # 模型描述信息
|
|
self._input_num = 0 # 输入数据个数
|
|
self._output_num = 0 # 输出数据个数
|
|
self._output_info = [] # 输出信息列表
|
|
self._is_released = False # 资源是否被释放
|
|
self.name = "人员入侵"
|
|
self.version = "V1.0"
|
|
self.model_type = 2
|
|
|
|
self.neth = 640 # 缩放的目标高度, 也即模型的输入高度
|
|
self.netw = 640 # 缩放的目标宽度, 也即模型的输入宽度
|
|
self.conf_threshold = threshold # 置信度阈值
|
|
|
|
self._init_acl()
|
|
self.init_ok = True
|
|
if not self._init_resource():
|
|
print("加载模型文件失败！")
|
|
self.init_ok = False
|
|
|
|
|
|
def __del__(self):
|
|
if self.init_ok:
|
|
self.release()
|
|
self._del_acl()
|
|
|
|
|
|
def verify(self,image,data,isdraw=1):
|
|
labels_dict = get_labels_from_txt('/mnt/zfbox/model/plugins/RYRQ_ACL/coco_names.txt') # 得到类别信息，返回序号与类别对应的字典
|
|
# 数据前处理
|
|
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 # 转换为内存连续存储的数组
|
|
|
|
# 模型推理, 得到模型输出
|
|
output = self.execute([img,])[0]
|
|
|
|
# 后处理 -- boxout 是 tensor-list: [tensor([[],[].[]])] --[x1,y1,x2,y2,置信度,coco_index]
|
|
boxout = nms(torch.tensor(output), conf_thres=0.4,
|
|
iou_thres=0.5) # 利用非极大值抑制处理模型输出，conf_thres 为置信度阈值，iou_thres 为iou阈值
|
|
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)) # 将推理结果缩放到原始图片大小
|
|
|
|
bwarn = False
|
|
warn_text = ""
|
|
#是否有检测区域，有先绘制检测区域 由于在该函数生成了polygon对象，所有需要在检测区域前调用。
|
|
if data[1] == 1:
|
|
self.draw_polygon(image,data[2],(255,0,0))
|
|
#过滤掉不是目标标签的数据
|
|
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 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 = "Intruder detected!"
|
|
img_dw = draw_bbox(filtered_pred_all, image, (0, 255, 0), 2, labels_dict) # 画出检测框、类别、概率
|
|
cv2.imwrite('img_res.png', img_dw)
|
|
return filtered_pred_all, False, ""
|
|
|
|
|
|
def testRun(self):
|
|
print("1111")
|