使用协程改造

main.py

@@ -0,0 +1,81 @@
+import asyncio
+import aiohttp
+import time
+import cv2
+from logger import logger
+from stream import StreamCapture
+from infer import DoorInference
+from logger import logger
+
+async def upload_image(session, url, payload, files):
+    try:
+        async with session.post(url, data=payload, files=files) as response:
+            result = await response.text()
+            logger.info(result)
+    except Exception as error:
+        logger.error(f'Error: {str(error)}')
+
+async def process_stream():
+    logger.info("====== Start Server =======")
+    human_model_path = "models/work_clo_person_head_hat.pt"
+    door_model_path = "models/door_classify.pt"
+    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 = time.time() - 30
+    async with aiohttp.ClientSession() as session:
+        for frame, ret in stream():
+            if not ret: 
+                continue
+            
+            image = frame.copy()
+            result = instance(image)
+            if len(result) > 0 and time.time() - posttime > 30:
+                try:
+                    posttime = time.time()
+                    videoTime = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
+                    fileTime = time.strftime('%Y-%m-%d-%H:%M:%S', time.localtime())
+                    filename = fileTime + ".jpg"
+                    filenameori = fileTime + "det.jpg"
+                    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)
+                    if not success:
+                        logger.error(f'imencode image error')
+                    content = encoded_image.tobytes()
+                    successori, encoded_imageori = cv2.imencode('.jpg', frame)
+                    if not successori:
+                        logger.error(f'imencode original image error')
+                    contentori = encoded_imageori.tobytes()
+                    
+                    payload = {
+                        'channel': '45',
+                        'classIndex': '8',
+                        'ip': '172.19.152.231',
+                        'videoTime': videoTime,
+                        'videoUrl': stream.stream_url
+                    }
+                    
+                    files = [
+                        ('file', (filename, content, 'image/jpeg')),
+                        ('oldFile', (filenameori, contentori, 'image/jpeg'))
+                    ]
+                    
+                    # 使用协程上传图像
+                    await upload_image(session, 'http://172.19.152.231/open/api/operate/upload', payload, files)
+                except Exception as error:
+                    logger.error(f'Error: {str(error)}')
+
+    logger.info("======= EXIT =======")
+
+if __name__ == "__main__":
+    asyncio.run(process_stream())

test.py

@@ -1,274 +1,80 @@
-'''
-@File    :   infer.py
-@Time    :   2024/10/17 11:18:00
-@Author  :   leon 
-@Version :   1.0
-@Desc    :   None
-'''
-import cv2
-import torch
+import asyncio
+import aiohttp
 import time
-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
-
+import cv2
+from logger import logger
+from stream import StreamCapture
+from infer import DoorInference
 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})
+async def upload_image(session, url, payload, files):
+    try:
+        async with session.post(url, data=payload, files=files) as response:
+            result = await response.text()
+            logger.info(result)
+    except Exception as error:
+        logger.error(f'Error: {str(error)}')
 
-        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__":
+async def process_stream():
     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) 这个区域是为了测试，画大了的
+    door_model_path = "models/door_classify.pt"
     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 = time.time() - 30
-
-    frame = cv2.imread("inference/test.jpg")
-
-    image = frame.copy()
-    result = instance(image)
-    if len(result) > 0 and time.time() - posttime > 30:
-        posttime  = time.time()
-        videoTime = time.strftime('%Y-%m-%d %H:%M:%S',time.localtime())
-        fileTime  = time.strftime('%Y-%m-%d-%H:%M:%S',time.localtime())
-        filename  = fileTime + ".jpg"
-        filenameori = fileTime + "det.jpg"
-        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',
+    async with aiohttp.ClientSession() as session:
+        frame = cv2.imread("inference/test.jpg")
+
+        image = frame.copy()
+            
+        result = instance(image)
+        if len(result) > 0 and time.time() - posttime > 30:
+            try:
+                posttime = time.time()
+                videoTime = time.strftime('%Y-%m-%d %H:%M:%S', time.localtime())
+                fileTime = time.strftime('%Y-%m-%d-%H:%M:%S', time.localtime())
+                filename = fileTime + ".jpg"
+                filenameori = fileTime + "det.jpg"
+                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)
+                if not success:
+                    logger.error(f'imencode image error')
+                content = encoded_image.tobytes()
+                successori, encoded_imageori = cv2.imencode('.jpg', frame)
+                if not successori:
+                    logger.error(f'imencode original image error')
+                contentori = encoded_imageori.tobytes()
+                
+                payload = {
+                    'channel': '45',
                     'classIndex': '8',
                     'ip': '172.19.152.231',
                     'videoTime': videoTime,
-                    'videoUrl': stream.stream_url}
-        files = [
-                        ('file', (filename, content, 'image/jpeg')),
-                        ('oldFile', (filenameori, contentori, 'image/jpeg')),
-                    ]
-    
-        result = requests.post('http://172.19.152.231/open/api/operate/upload', data=payload, files=files)
-        logger.info(result)
+                    'videoUrl': stream.stream_url
+                }
+                
+                files = [
+                    ('file', (filename, content, 'image/jpeg')),
+                    ('oldFile', (filenameori, contentori, 'image/jpeg'))
+                ]
+                
+                # 使用协程上传图像
+                await upload_image(session, 'http://172.19.152.231/open/api/operate/upload', payload, files)
+            except Exception as error:
+                logger.error(f'Error: {str(error)}')
+
+    logger.info("======= EXIT =======")
 
-    logger.info("=======  EXIT  =======")
-    #cv2.imwrite("result/result.jpg", image)
+if __name__ == "__main__":
+    asyncio.run(process_stream())