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
import requests

from logger import logger
from stream import StreamCapture 

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)):
            if person_polygon.intersection(self.person_areas_polygon[i]).area / person_polygon.area > 0.5:
                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=[4], 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()
        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
    
    
if __name__ == "__main__":
    logger.info("Start Server")
    human_model_path = "models/work_clo_person_head_hat.pt"
    door_model_path  = "models/door_classify.pt"
    #image_path = "images/camera1_9ce1aca78db14c5807a271cb154fed5b.jpg"
    #image = cv2.imread(image_path)
    # logger.info(video)
    # (222, 59), (432, 3), (528, 96), (318, 198) 这个区域是为了测试，画大了的
    test_area = [[(222, 59), (432, 3), (528, 96), (318, 198)]]
    instance = DoorInference(human_model_path, door_model_path, person_areas=None)
    # 返回需要门和人
    ip = '172.19.152.231'
    channel = '45'
    stream = StreamCapture(ip, channel)
    posttime = 0
    for frame, ret in stream():
        if not ret:continue
        image = frame.copy()
        result = instance(image)
        if len(result)>0 and time.time()-posttime>30:
            videoTime = time.strftime('%Y-%m-%d %H:%M:%S',time.localtime())
            logger.info(videoTime)
            logger.info(result)
            for res in result:
                cv2.rectangle(image, tuple(map(int, (res.left, res.top))), tuple(map(int, (res.right, res.bottom))), (255,0, 0), 4)
            success, encoded_image = cv2.imencode('.jpg', image)
            content = encoded_image.tobytes()
            successori, encoded_imageori = cv2.imencode('.jpg',frame)
            contentori = encoded_imageori.tobytes()
            payload = {'channel': '45',
                        'classIndex': '8',
                        'ip': '172.19.152.231',
                        'videoTime': videoTime,
                        'videoUrl': video}
            files = [
                            ('file', (filename, content, 'image/jpeg')),
                            ('oldFile', (filenameori, contentori, 'image/jpeg')),
                        ]
            try:
                result = requests.post('http://172.19.145.197/open/api/operate/upload', data=payload, files=files)
                logger.info(result)
            except Exception:
                logger.info('posterror')
    logger.info("EXIT")
    #cv2.imwrite("result/result.jpg", image)
    

"""
1. 检测工服、安全帽等
2. 检测人的模型
3. 安全帽的分类


逻辑
1. 检测人
2. 检测安全帽、工服、头  检测到人没有工服报警
3. 安全帽头再分类
"""