jnhg_door/main.py

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())