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