jnhg_door/main.py

import asyncio
import aiohttp
import time
import cv2
from stream import StreamCapture
from infer import DoorInference
from logger import logger
import time
from datetime import datetime

def is_time_in_period(start_time_str, end_time_str):
    # 解析开始和结束时间字符串为time对象
    start_time = datetime.strptime(start_time_str, "%H:%M").time()
    end_time = datetime.strptime(end_time_str, "%H:%M").time()
    current_time = datetime.now().time()
    
    # 判断时间段是否跨天
    if start_time <= end_time:
        # 时间段在同一天
        return start_time <= current_time <= end_time
    else:
        # 时间段跨天（如22:00到次日02:00）
        return current_time >= start_time or current_time <= end_time

async def upload_image(session, url, payload, files):
    try:
        form_data = aiohttp.FormData()
        # 添加普通表单数据
        for key, value in payload.items():
            form_data.add_field(key, value)

        # 添加文件数据
        for key, (filename, content, content_type) in files:
            form_data.add_field(key, content, filename=filename, content_type=content_type)
        
        # 发起 POST 请求
        async with session.post(url, data=form_data) 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/yolo11m.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)
    
    start_det_time = "06:00"
    end_det_time = "18:00"

    posttime = time.time() - 30
    async with aiohttp.ClientSession() as session:
        while True:
            if not is_time_in_period(start_det_time, end_det_time):
                time.sleep(10)
                continue
            for frame, ret in stream():
                if not is_time_in_period(start_det_time, end_det_time):
                    logger.info(f"当前时间不在时间段 {start_det_time} ~ {end_det_time}")
                    stream.close()
                    break
                if not ret: 
                    continue
                if time.time() - posttime < 30:
                    continue
                image = frame.copy()
                result = instance(image)
                if len(result) > 0:
                    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('imencode image error')
                            continue
                        content = encoded_image.tobytes()
                        successori, encoded_imageori = cv2.imencode('.jpg', frame)
                        if not successori:
                            logger.error('imencode original image error')
                            continue
                        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())