zfbox/model/plugins/RYRQ_ACL/RYRQ_Model_ACL.py

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")