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.
49 lines
2.3 KiB
49 lines
2.3 KiB
import os.path
|
|
|
|
from model.plugins.ModelBase import ModelBase
|
|
from myutils.ConfigManager import myCongif
|
|
import torch
|
|
import cv2
|
|
class Model(ModelBase):
|
|
def __init__(self,path):
|
|
super().__init__()
|
|
self.name = "人员入侵"
|
|
self.version = "V1.0"
|
|
self.model_type = 2
|
|
#找pt模型路径 -- 一个约束py文件和模型文件的路径关系需要固定， -- 上传模型时，要解压好路径
|
|
dirpath,filename = os.path.split(path)
|
|
model_file = os.path.join(dirpath,"yolov5s.pt") #目前约束模型文件和py文件在同一目录
|
|
yolov5_path = myCongif.get_data("yolov5_path")
|
|
print(f"************{model_file},{yolov5_path}")
|
|
#实例化模型
|
|
self.model = torch.hub.load(yolov5_path, 'custom', path=model_file, source='local')
|
|
|
|
def verify(self,image,data,isdraw=1):
|
|
results = self.model(image) # 进行模型检测 --- 需要统一接口
|
|
detections = results.pandas().xyxy[0].to_dict(orient="records")
|
|
bwarn = False
|
|
warn_text = ""
|
|
#是否有检测区域，有先绘制检测区域 由于在该函数生成了polygon对象，所有需要在检测区域前调用。
|
|
if data[1] == 1:
|
|
self.draw_polygon(image,data[2],(0, 255, 0))
|
|
# 绘制检测结果 --- 也需要封装在类里，
|
|
for det in detections:
|
|
if det['name'] == 'person': #标签是人
|
|
x1, y1, x2, y2 = int(det['xmin']), int(det['ymin']), int(det['xmax']), int(det['ymax'])
|
|
# 绘制目标识别的锚框
|
|
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!"
|
|
return detections,bwarn,warn_text
|
|
|
|
def testRun(self):
|
|
print("1111")
|