tuoheng_alg_move_20240715/test/demo/demo1.py

import copy
import json
import os
import time
from concurrent.futures import ThreadPoolExecutor
from multiprocessing import Queue, Process

from loguru import logger
import subprocess as sp

import cv2
import numpy as np
from aip import AipImageClassify
from enums.BaiduSdkEnum import BAIDUERRORDATA, VehicleEnumVALUE
from enums.ExceptionEnum import ExceptionType
from enums.ModelTypeEnum import ModelType
from exception.CustomerException import ServiceException


def get_recording_video_info(url):
    try:
        video_info = 'ffprobe -show_format -show_streams -of json %s' % url
        p = sp.Popen(video_info, stdout=sp.PIPE, stderr=sp.PIPE, shell=True)
        out, err = p.communicate(timeout=17)
        if p.returncode != 0:
            raise Exception("未获取视频信息！！！！！")
        probe = json.loads(out.decode('utf-8'))
        if probe is None or probe.get("streams") is None:
            raise Exception("未获取视频信息！！！！！:")
        video_stream = next((stream for stream in probe['streams'] if stream.get('codec_type') == 'video'), None)
        if video_stream is None:
            raise Exception("未获取视频信息！！！！！")
        width = video_stream.get('width')
        height = video_stream.get('height')
        nb_frames = video_stream.get('nb_frames')
        fps = video_stream.get('r_frame_rate')
        up, down = str(fps).split('/')
        fps = int(eval(up) / eval(down))
        return (width, height, nb_frames, fps)
    except Exception as e:
        raise e

client = AipImageClassify(str(31096670), 'Dam3O4tgPRN3qh4OYE82dbg7', '1PGZ9LAXRR5zcT5MN9rHcW8kLBIS5DAa')
def vehicleDetect(client, iamge, options={}):
    reply_num = 0
    reply_value = None
    while True:
        try:
            options["show"] = "true"
            res_image = client.vehicleDetect(iamge,options)
            error_code = res_image.get("error_code")
            if error_code:
                enum = BAIDUERRORDATA.get(error_code)
                # 如果异常编码未知, 返回空值
                if enum is None:
                    logger.error("百度云车辆检测异常！error_code:{}", error_code)
                    return None
                # 重试指定次数后，还是异常，输出统一内部异常
                if enum.value[3] == 0:
                    if reply_value is None:
                        reply_value = 10
                    logger.error("百度云车辆检测异常！error_code:{}, error_msg:{}, reply_num:{}", enum.value[0], enum.value[2], reply_num)
                    raise Exception()
                # 重试指定次数后，还是异常，输出对应的异常
                if enum.value[3] == 1:
                    if reply_value is None:
                        reply_value = 10
                    raise ServiceException(str(enum.value[0]), enum.value[2])
                # 重试指定次数后，还是异常，输出空
                if enum.value[3] == 2:
                    if reply_value is None:
                        reply_value = 10
                    if reply_num >= reply_value:
                        return None
            return res_image
        except ServiceException as s:
            time.sleep(0.2)
            reply_num += 1
            if reply_num > reply_value:
                logger.exception("车辆检测识别失败: {}", s.msg)
                raise ServiceException(e.code, e.msg)
        except Exception as e:
            logger.exception("车辆检测失败: {}, 当前重试次数：{}", e, reply_num)
            time.sleep(0.2)
            reply_num += 1
            if reply_num > reply_value:
                logger.exception("车辆检测识别失败: {}", e)
                raise ServiceException(ExceptionType.SERVICE_INNER_EXCEPTION.value[0],
                                       ExceptionType.SERVICE_INNER_EXCEPTION.value[1])

def mark(content, info, img, color):
    score = info.get("probability")
    if score is None:
        score = info.get("location").get("score")
    text = "%s: %.2f]" % (content, score)
    text_xy = (info.get("location").get("left"), info.get("location").get("top") - 25)
    img_lu = (info.get("location").get("left"), info.get("location").get("top"))
    img_rd = (info.get("location").get("left") + info.get("location").get("width"),
              info.get("location").get("top") + info.get("location").get("height"))
    cv2.putText(img, text, text_xy, cv2.FONT_HERSHEY_SIMPLEX, 1.0, color, 2, cv2.LINE_AA)
    count = 1
    if img.shape[1] > 1600:
        count = 2
    cv2.rectangle(img, img_lu, img_rd, color, count)
    return img

def pull_stream(url, queue, nb_frames):
    command = ['ffmpeg -re -y -i ' + url +' -f rawvideo -pix_fmt bgr24 -an -']
    pull_p = sp.Popen(command, stdout=sp.PIPE, shell=True)
    aa = 0
    try:
        while True:
            if queue.qsize() == 200:
                time.sleep(1)
                continue
            in_bytes = pull_p.stdout.read(width*height*3)
            if in_bytes is not None and len(in_bytes) > 0:
                img = np.frombuffer(in_bytes, np.uint8).reshape([height, width, 3])
                queue.put({"status": "1", "img": img})
                aa+=1
            else:
                if aa -10 > nb_frames:
                    queue.put({"status": "2"})
                    pull_p.terminate()
                    pull_p.wait()
                    break;
    except Exception as e:
        logger.exception("拉流异常: {}", e)
    finally:
        pull_p.terminate()
        pull_p.wait()

def getQueue(queue):
    eBody = None
    try:
        eBody = queue.get(block=False)
        return eBody
    except Exception as e:
        pass
    return eBody

def buildFrame(queue, senlin_mod, client, width, height, nb_frames, fps):
    frames = []
    status = None
    for i in range(queue.qsize()):
        frame_result = getQueue(queue)
        if frame_result is None:
            time.sleep(0.01)
            continue
        if frame_result.get("status") == '1':
            frames.append((frame_result.get("img"),  senlin_mod, client, width, height, nb_frames, fps))
        else:
            status = frame_result.get("status")
    return frames, status

def process(frame):
    try:
        p_result, timeOut = frame[1].process(copy.deepcopy(frame[0]), frame[3])
        or_result, or_image = cv2.imencode(".jpg", frame[0])
        result = vehicleDetect(frame[2], or_image)
        if result is not None:
            vehicleInfo = result.get("vehicle_info")
            if vehicleInfo is not None and len(vehicleInfo) > 0:
                for i, info in enumerate(vehicleInfo):
                    value = VehicleEnumVALUE.get(info.get("type"))
                    if value is None:
                        logger.error("车辆识别出现未支持的目标类型！type:{}", info.get("type"))
                        raise ServiceException(ExceptionType.SERVICE_INNER_EXCEPTION.value[0],
                                               ExceptionType.SERVICE_INNER_EXCEPTION.value[1])
                    p_result[1] = mark(value.value[1], info, p_result[1], (255, 0, 255))
        frame_merge = np.hstack((frame[0], p_result[1]))
        return frame_merge
    except Exception as e:
        logger.exception("模型分析异常: {}", e)
        return None

queue = Queue(200)
url ='/home/th/tuo_heng/dev/11.mp4'
width, height, nb_frames, fps = get_recording_video_info(url)
my_process = Process(target = pull_stream, args=(url, queue, nb_frames))
#启动子进程
my_process.start()



current_path = os.path.abspath(os.path.dirname(__file__))
import GPUtil
senlin_mod = Model(str(GPUtil.getAvailable()[0]), [2,3,4], logger, "11111", ModelType.FOREST_FARM_MODEL)
or_video_file = cv2.VideoWriter("aaa.mp4", cv2.VideoWriter_fourcc(*'mp4v'), fps,
                                (int(width) * 2, int(height)))


with ThreadPoolExecutor(max_workers=3) as t:
    task_frame = None
    while True:
        frames = []
        status = None
        if task_frame is not None:
            frames, status = task_frame.result()
        task_frame = t.submit(buildFrame, queue, senlin_mod, client, width, height, nb_frames, fps)
        if len(frames) == 0 and status is None:
            time.sleep(0.02)
            continue
        if frames is not None and len(frames) > 0:
            for result in t.map(process, frames):
                if result is not None:
                    or_video_file.write(result)
        if status is None:
            continue
        if status.get("status") == "2":
            t.shutdown(wait=False)
            or_video_file.release()

t.shutdown(wait=False)
or_video_file.release()