jnhg_door/infer.py

'''
@File    :   infer.py
@Time    :   2024/10/17 11:18:00
@Author  :   leon 
@Version :   1.0
@Desc    :   None
'''
import cv2
import torch
import numpy as np
from ultralytics import YOLO
from shapely.geometry.polygon import Polygon

from collections import namedtuple
from typing import List, Tuple

Box = namedtuple("Box", ["left", "top", "right", "bottom", "confidence", "label"])

class PerspectiveMatrix:
    def __init__(self, matrix: np.ndarray, target: Tuple[int, int]) -> None:
        self.matrix = matrix
        self.target = target
    
    def __repr__(self):
        matrix_str = np.array2string(self.matrix, formatter={'float_kind': lambda x: f"{x:.2f}"})
        return f"PerspectiveMatrix(matrix={matrix_str}, target={self.target})"
    
    @staticmethod
    def perspective_matrix(src: List[Tuple[int, int]], dst: List[Tuple[int, int]], target : Tuple[int, int]) -> "PerspectiveMatrix":
        """
        计算透视变换矩阵。
        
        参数:
        - src: 源图像上的 4 个点 [(x1, y1), (x2, y2), (x3, y3), (x4, y4)]
        - dst: 目标图像上的 4 个点 [(u1, v1), (u2, v2), (u3, v3), (u4, v4)]
        
        返回:
        - PerspectiveMatrix: 透视变换矩阵
        """
        src_pts = np.array(src, dtype=np.float32)
        dst_pts = np.array(dst, dtype=np.float32)

        # 计算透视变换矩阵
        # matrix = cv2.getPerspectiveTransform(src_pts, dst_pts)
        # 单应性矩阵
        matrix, _ = cv2.findHomography(src_pts, dst_pts)

        return PerspectiveMatrix(matrix=matrix, target=target)

class ImageTransformer:
    @staticmethod
    def PerspectiveTransform(
        image: np.ndarray, 
        perspective_matrix: PerspectiveMatrix,
        flags=cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_CONSTANT,
        borderValue=(114, 114, 114)
    ) -> np.ndarray:
        if image is None or perspective_matrix is None:
            raise ValueError("Input image and warpaffine_matrix cannot be None.")

        transformed_image = cv2.warpPerspective(
            image, 
            perspective_matrix.matrix, 
            perspective_matrix.target, 
            flags=flags,
            borderMode=borderMode,
            borderValue=borderValue)
        return transformed_image

class AreaManage(object):
    def __init__(self):
        self.door_areas = [
            # [(934, 143), (958, 130), (961, 165), (936, 181)],
            [(564, 399), (682, 432), (684, 528), (574, 493)]]
        
        self.door_perspective_matrix = []

        self.dst_door_points = [(0,0), (224, 0), (224, 224), (0, 224)]

        for p in self.door_areas:
            perspective_matrix = PerspectiveMatrix.perspective_matrix(p, self.dst_door_points, (224, 224))
            self.door_perspective_matrix.append(perspective_matrix)

        # 预设值的三个门对应的人的活动区域
        self.person_areas = [
            # [(900, 152),  (914, 84), (637, 20), (604, 73)],
            [(860, 200), (958, 226), (687, 432), (586, 395)]]
        
        self.person_areas_polygon = [Polygon(area) for area in self.person_areas]
        

    def update(self, person_area):
        areas = []
        person_area_polygon = Polygon(person_area)
        for pap in self.person_areas_polygon:
            areas.append(person_area_polygon.intersection(pap).area)
        max_area = max(areas)
        if max_area > 0:
            max_idx = areas.index(max_area)
            self.person_areas_polygon[max_idx] = person_area_polygon
        
    def inner_area(self, person_polygon):
        for i in range(len(self.person_areas_polygon)):
            overlap_ratop = person_polygon.intersection(self.person_areas_polygon[i]).area / person_polygon.area
            if overlap_ratop > 0.4:
                return i
        return -1


        

"""
1. 人员检测
2. 检查区域内是否有人
3. 如果区域内有人，分类该区域的门是否关闭
4. 如果关闭，上报违章
"""
class DoorInference(object):
    """
    human_model_path : 检测人的模型地址
    door_model_path  : 门分类的模型地址
    person_areas     : 电子围栏区域 [[(x,y),(x,y),(x,y),(x,y),...],[(x,y),(x,y),(x,y),(x,y),...],...]
    device_id        : 显卡id, -1表示使用cpu
    confidence_threshold : 检测人的模型的阈值
    """
    def __init__(self, human_model_path, door_model_path, person_areas, device_id=0, confidence_threshold= 0.5) -> "DoorInference":
        self.device = torch.device(f"cuda:{device_id}" if torch.cuda.is_available() and device_id !=-1 else "cpu")
        self.confidence_threshold = confidence_threshold
        self.human_model = YOLO(human_model_path).to(self.device)
        self.door_model  = YOLO(door_model_path).to(self.device)

        self.door_names  = {0: 'block', 1: 'close', 2: 'open'}

        self.am = AreaManage()

        if person_areas:
            for person_area in person_areas:
                self.am.update(person_area)

    def __repr__(self):
        infer_str = f"DoorInference(device = {self.device})"
        return infer_str
    
    def rect(self, point_list):
        """
        根据给定的点计算矩形框
        :param point_list: [(x1, y1), (x2, y2), ..., (xn, yn)] 多个点
        :return: 左上角和右下角的坐标，表示矩形的框
        """
        # 获取所有点的 x 和 y 坐标的最小和最大值
        min_x = min(point[0] for point in point_list)
        max_x = max(point[0] for point in point_list)
        min_y = min(point[1] for point in point_list)
        max_y = max(point[1] for point in point_list)

        # 计算矩形框的左上角和右下角
        left, top = min_x, min_y
        right, bottom = max_x, max_y

        return left, top ,right, bottom

    """
    返回所有检测到的人
    """
    def person_detect(self, image):
        objs  = []
        confs = []
        results = self.human_model(image, stream=False, classes=[0], conf=self.confidence_threshold, iou=0.3, imgsz=640)
        for result in results:
            boxes = result.boxes.cpu()
            for box in boxes:
                conf = box.conf.cpu().numpy().tolist()[0]
                left, top, right, bottom = box.xyxy.tolist()[0]
                objs.append([(left, top), (right, top), (right, bottom), (left, bottom)])
                confs.append(conf)
        return objs, confs

    """
    判断门是否关闭
    """
    def is_door_close(self, image, index) -> bool:
        # 图片透视变换, 只需要门的区域
        # 比只使用矩形框更大程度的只保留门的区域
        door_image  = ImageTransformer.PerspectiveTransform(image, self.am.door_perspective_matrix[index])
        # cv2.imshow('0',door_image)
        # cv2.waitKey(30000)
        res = self.door_model(door_image)
        class_idx = res[0].probs.top1
        class_conf = res[0].probs.top1conf.cpu().numpy().item()
        return class_idx == 1, class_conf

    """ 
    image 输入待识别图片
    返回需要画框的门和人的坐标
    """
    def __call__(self, image):
        inner_person = {0:[], 1:[], 2:[]}
        person_boxes, person_confs = self.person_detect(image)
        for person_box, person_conf in zip(person_boxes, person_confs):
            person_polygon = Polygon(person_box)
            idx = self.am.inner_area(person_polygon)
            if idx == -1: continue
            inner_person[idx].append({"box" : person_box, "conf" : person_conf})

        result = []
        for i, persons in inner_person.items():
            if len(persons) == 0:
                continue
            close, class_conf = self.is_door_close(image, i)
            if close:
                left, top ,right, bottom = self.rect(self.am.door_areas[i])
                result.append(Box(left=left, top=top, right=right, bottom=bottom, confidence=class_conf, label="door_close"))
                for person in persons:
                    left, top ,right, bottom = self.rect(person["box"])
                    conf = person["conf"]
                    result.append(Box(left=left, top=top, right=right, bottom=bottom, confidence=conf, label="person"))
        return result