from model.plugins.ModelBase import ModelBase
from myutils.ConfigManager import myCongif
from shapely.geometry import Point, Polygon
import torch
import cv2
import av
import numpy as np
class Model(ModelBase):
    def __init__(self,path):
        super().__init__()
        self.name = "人员入侵"
        self.version = "V1.0"
        self.model_type = 2
        #实例化模型
        self.model = torch.hub.load(myCongif.get_data("yolov5_path"), 'custom', path=path, source='local')

    def verify(self,image,data):
        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