jnhg_door/test.py

import asyncio
import aiohttp
import time
import io
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:
        form_data = aiohttp.FormData()
        # 添加普通表单数据
        for key, value in payload.items():
            form_data.add_field(key, value)

        # 添加文件数据
        for key, (filename, content, content_type) in files:
            content_io = io.BytesIO(content)
            content_io.seek(0)
            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/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:
        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('imencode image error')
                content = encoded_image.tobytes()
                successori, encoded_imageori = cv2.imencode('.jpg', frame)
                if not successori:
                    logger.error('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())