python --fastapi推流AI推理

python 复制代码

<!DOCTYPE html>
<html>
<body>
    <img id="videoStream" style="width:100%; max-width:800px;">
    <script>
        const params = new URLSearchParams({
            url: 'rtsp://192.168.8.119:555/rtp/34020000001320000206_34020000001320000206',
            box_rgb: '1',
            algo: '1',
            threshold: '0.50',
            regions: JSON.stringify([])
        });
        document.getElementById('videoStream').src =
            'http://192.168.8.119:49351/instant_ai?' + params.toString();
    </script>
</body>
</html>

python 复制代码

@app.get('/instant_ai') # 实时AI
async def instant_ai(request: Request) -> dict: return await instant_ai_views(request)

python 复制代码

async def instant_ai_views(request: Request):
    '''
    实时AI
    @params  url            -->  摄像头地址;
    @params  box_rgb        -->  检测框颜色(1红 2绿  3蓝 4白 5黑 6紫);
    @params  algo           -->  算法(1: 人员越线; 2: 安全帽; 3: 明火; 4:未穿反光衣; 5: 未穿工服);
    @params  threshold      -->  阈值;
    @params  regions        -->  区域;
    '''
    params = request.query_params  # 请求参数
    rtx = {k: params.get(k) for k in ('url', 'box_rgb', 'algo', 'threshold', 'regions')}
    if rtx['threshold'] in ('0', 0, None, ''):
        rtx['threshold'] = 0.5

    if rtx['box_rgb'] in ('0', 0, None, ''):
        rtx['box_rgb'] = 1

    return StreamingResponse(
        reid_event_generator(rtx),
        media_type='multipart/x-mixed-replace; boundary=frame')

python 复制代码

# coding: utf-8
import ast
import asyncio

import cv2
from loguru import logger
from yolo import utils


def get_func(rtx: dict):
    '''
    根据枚举获取算法
    算法(1: 人员越线; 2: 安全帽; 3: 明火; 4:未穿反光衣; 5: 未穿工服);
    '''
    algo: str = rtx['algo']
    if algo in (1, '1'):
        func =  utils.yuexian_yolo.yuexian_predict
    elif algo in (2, '2'):
        func = utils.anquanmao_yolo.aqm_predict
    elif algo in (3, '3'):
        func = utils.minghuo_yolo.flame_predict
    elif algo in (4, '4'):
        func = utils.fanguangyi_yolo.fanguangyi_predict
    else:
        func = utils.gongfu_yolo.gongfu_predict
    return func

def get_color(rtx: dict):
    '''
    获取颜色
    红(255, 0, 0) 绿(0, 255, 0) 蓝(0, 0, 255) 白(255, 255, 255) 黑(0, 0, 0) 紫(255, 0, 255)
    '''
    box_rgb = int(rtx['box_rgb'])
    color = {
        1: (255, 0, 0), # 红
        2: (0, 255, 0), # 绿
        3: (0, 0, 255), # 蓝
        4: (255, 255, 255), # 白
        5: (0, 0, 0), # 黑
        6: (255, 0, 255) # 紫
    }
    return color[box_rgb]


async def reid_event_generator(rtx: dict):
    '''响应层'''
    func = get_func(rtx)
    color = get_color(rtx)
    threshold = float(rtx['threshold'])
    regions = ast.literal_eval(rtx['regions'])
    cap = cv2.VideoCapture(rtx['url'])
    cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)  # 缓冲区只留1帧

    try:
        while True:
            for _ in range(3): # 跳帧
                cap.grab()
            success, frame = cap.retrieve()
            if not success:
                await asyncio.sleep(0.01)
                continue

            _, frame = func(frame, regions, threshold, color) # 自己写的AI推理 返回检测的图片
            ret, buffer = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 70])
            jpg_bytes = buffer.tobytes()
            yield (b'--frame\r\n'
                   b'Content-Type: image/jpeg\r\n\r\n' + jpg_bytes + b'\r\n')
            await asyncio.sleep(0.01)
    finally:
        cap.release()
        logger.success(f'浏览器已被关闭,退出拉流,摄像头资源已释放!')

python --fastapi推流AI推理

前端demo

服务端