zfbox/model/plugins/ModelBase.py

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

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.system = platform.system() #获取系统平台
        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(ast.literal_eval(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



    @abstractmethod
    def verify(self,image,data,isdraw=1):
        '''
        :param image:  需要验证的图片
        :param data: select t1.model_id,t1.check_area,t1.polygon ,t2.duration_time,t2.proportion,t2.model_path
        :param isdraw: 是否需要绘制线框：0-不绘制，1-绘制
        :return: detections,bwarn,warntext  bwarn:0-没有识别到符合要求的目标，1-没有识别到符合要求的目标。
        '''
        pass