zfbox/model/plugins/ModelBase.py

from abc import abstractmethod,ABC
from shapely.geometry import Point, Polygon
import numpy as np
import cv2

class ModelBase(ABC):
    def __init__(self):
        self.name = None        #基于name来查询，用户对模型的配置参数，代表着模型名称需要唯一 2024-6-18 -逻辑还需要完善和验证
        self.version = None
        self.model_type = None  # 模型类型 1-图像分类，2-目标检测（yolov5)，3-分割模型，4-关键点
        self.do_map = {  # 定义插件的入口函数 --
            # POCType.POC: self.do_verify,
            # POCType.SNIFFER: self.do_sniffer,
            # POCType.BRUTE: self.do_brute
        }

    def __del__(self):
        print("资源释放")

    def draw_polygon(self, img, polygon_points,color=(0, 255, 0)):
        self.polygon = Polygon(polygon_points)

        points = np.array([self.polygon.exterior.coords], dtype=np.int32)
        cv2.polylines(img, points, isClosed=True, color=color, thickness=2)

    def is_point_in_region(self, point):
        '''判断点是否在区域内，需要先执行draw_polygon'''
        if self.polygon:
            return self.polygon.contains(Point(point))
        else:
            return False

    def save_warn(self):
        '''保存报警信息'''
        pass

    def send_warn(self):
        '''发送报警信息'''
        pass

    @abstractmethod
    def verify(self,image,data):
        '''
        :param image:  需要验证的图片
        :param data: select t1.model_id,t1.check_area,t1.polygon ,t2.duration_time,t2.proportion,t2.model_path
        :return: detections,bwarn,warntext
        '''
        pass