#ne -*- coding: utf-8 -*- from concurrent.futures import ThreadPoolExecutor from multiprocessing import Process from os import getpid import os from os.path import join from time import time, sleep from traceback import format_exc import cv2 import numpy as np import psutil from loguru import logger from enums.ExceptionEnum import ExceptionType from enums.ModelTypeEnum import ModelType from exception.CustomerException import ServiceException from util import ImageUtils from util.Cv2Utils import video_conjuncing, write_or_video, write_ai_video, push_video_stream, close_all_p from util.ImageUtils import url2Array, add_water_pic from util.LogUtils import init_log from util.PlotsUtils import draw_painting_joint, filterBox, xywh2xyxy2, draw_name_joint from util.QueUtil import get_no_block_queue, put_queue, clear_queue class PushStreamProcess(Process): __slots__ = ("_msg", "_push_queue", "_image_queue", '_push_ex_queue', '_hb_queue', "_context") def __init__(self, *args): super().__init__() # 传参 self._msg, self._push_queue, self._image_queue, self._push_ex_queue, self._hb_queue, self._context = args self._algStatus = False # 默认关闭 self._algSwitch = self._context['service']['algSwitch'] #0521: default_enabled = str(self._msg.get("defaultEnabled", "True")).lower() == "true" if default_enabled: print("执行默认程序（defaultEnabled=True）") self._algSwitch = True # 这里放默认逻辑的代码 else: print("执行替代程序（defaultEnabled=False）") # 这里放非默认逻辑的代码 self._algSwitch = False print("---line53 :PushVideoStreamProcess.py---",self._algSwitch) def build_logo_url(self): logo = None if self._context["video"]["video_add_water"]: logo = self._msg.get("logo_url") if logo: logo = url2Array(logo, enable_ex=False) if logo is None: logo = cv2.imread(join(self._context['base_dir'], "image/logo.png"), -1) self._context["logo"] = logo @staticmethod def handle_image(det_xywh, det, frame_score, copy_frame, draw_config, code_list): code, det_result = det # 每个单独模型处理 # 模型编号、100帧的所有问题, 检测目标、颜色、文字图片 if len(det_result) > 0: font_config, allowedList = draw_config["font_config"], draw_config[code]["allowedList"] rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"] for qs in det_result: box, score, cls = xywh2xyxy2(qs) if cls not in allowedList or score < frame_score: continue label_array, color = label_arrays[cls], rainbows[cls] draw_painting_joint(box, copy_frame, label_array, score, color, font_config) if det_xywh.get(code) is None: det_xywh[code], code_list[code] = {}, {} cd = det_xywh[code].get(cls) if cd is None: code_list[code][cls] = 1 det_xywh[code][cls] = [[cls, box, score, label_array, color]] else: code_list[code][cls] += 1 det_xywh[code][cls].append([cls, box, score, label_array, color]) class OnPushStreamProcess(PushStreamProcess): __slots__ = () def run(self): self.build_logo_url() msg, context = self._msg, self._context base_dir, env, orFilePath, aiFilePath, logo, service_timeout, frame_score = context["base_dir"], \ context['env'], context["orFilePath"], context["aiFilePath"], context["logo"], \ int(context["service"]["timeout"]) + 120, context["service"]["filter"]["frame_score"] request_id, push_url = msg["request_id"], msg["push_url"] push_queue, image_queue, push_ex_queue, hb_queue = self._push_queue, self._image_queue, self._push_ex_queue, \ self._hb_queue or_video_file, ai_video_file, push_p, ex = None, None, None, None ex_status = True # 图片相似度开关 picture_similarity = bool(context["service"]["filter"]["picture_similarity"]) qs_np_tmp = None pix_dis = 60 try: init_log(base_dir, env) logger.info("开始实时启动推流进程！requestId:{},pid:{}, ppid:{}", request_id,os.getpid(),os.getppid()) with ThreadPoolExecutor(max_workers=2) as t: # 定义三种推流、写原视频流、写ai视频流策略 # 第一个参数时间, 第二个参数重试次数 p_push_status, or_write_status, ai_write_status = [0, 0], [0, 0], [0, 0] start_time = time() while True: # 检测推流执行超时时间, 1.防止任务运行超时 2.主进程挂了,子进程运行超时 if time() - start_time > service_timeout: logger.error("推流超时, requestId: {}", request_id) raise ServiceException(ExceptionType.TASK_EXCUTE_TIMEOUT.value[0], ExceptionType.TASK_EXCUTE_TIMEOUT.value[1]) # 系统由于各种问题可能会杀死内存使用多的进程, 自己杀掉自己 if psutil.Process(getpid()).ppid() == 1: logger.info("推流进程检测到父进程异常停止, 自动停止推流进程, requestId: {}", request_id) ex_status = False for q in [push_queue, image_queue, push_ex_queue, hb_queue]: clear_queue(q) break # 获取推流的视频帧 push_r = get_no_block_queue(push_queue) if push_r is not None: if push_r[0] == 1: frame_list, frame_index_list, all_frames, draw_config, push_objs = push_r[1] for i, frame in enumerate(frame_list): pix_dis = int((frame.shape[0]//10)*1.2) # 复制帧用来画图 copy_frame = frame.copy() det_xywh, thread_p = {}, [] det_xywh2 = {} # 所有问题的矩阵集合 qs_np = None qs_reurn = [] for det in push_objs[i]: code, det_result = det # 每个单独模型处理 # 模型编号、100帧的所有问题, 检测目标、颜色、文字图片 if len(det_result) > 0: font_config, allowedList = draw_config["font_config"], draw_config[code]["allowedList"] rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"] for qs in det_result: try: # 应对NaN情况 box, score, cls = xywh2xyxy2(qs) except: continue if cls not in allowedList or score < frame_score: continue label_array, color = label_arrays[cls], rainbows[cls] if ModelType.CHANNEL2_MODEL.value[1] == str(code) and cls == 2: rr = t.submit(draw_name_joint, box, copy_frame, draw_config[code]["label_dict"], score, color, font_config, qs[6]) else: rr = t.submit(draw_painting_joint, box, copy_frame, label_array, score, color, font_config) thread_p.append(rr) if det_xywh.get(code) is None: det_xywh[code] = {} cd = det_xywh[code].get(cls) if not (ModelType.CHANNEL2_MODEL.value[1] == str(code) and cls == 2): if cd is None: det_xywh[code][cls] = [[cls, box, score, label_array, color]] else: det_xywh[code][cls].append([cls, box, score, label_array, color]) if qs_np is None: qs_np = np.array([box[0][0], box[0][1], box[1][0], box[1][1], box[2][0], box[2][1], box[3][0], box[3][1], score, cls, code],dtype=np.float32) else: result_li = np.array([box[0][0], box[0][1], box[1][0], box[1][1], box[2][0], box[2][1], box[3][0], box[3][1], score, cls, code],dtype=np.float32) qs_np = np.row_stack((qs_np, result_li)) if logo: frame = add_water_pic(frame, logo, request_id) copy_frame = add_water_pic(copy_frame, logo, request_id) if len(thread_p) > 0: for r in thread_p: r.result() #print('----line173:',self._algSwitch,self._algStatus) if self._algSwitch and (not self._algStatus): # frame_merge = video_conjuncing(frame, frame.copy()) frame_merge = frame.copy() else: # frame_merge = video_conjuncing(frame, copy_frame) frame_merge = copy_frame # 写原视频到本地 write_or_video_result = t.submit(write_or_video, frame, orFilePath, or_video_file, or_write_status, request_id) # 写识别视频到本地 write_ai_video_result = t.submit(write_ai_video, frame_merge, aiFilePath, ai_video_file, ai_write_status, request_id) push_stream_result = t.submit(push_video_stream, frame_merge, push_p, push_url, p_push_status, request_id) # 如果有问题, 走下面的逻辑 if qs_np is not None: if len(qs_np.shape) == 1: qs_np = qs_np[np.newaxis,...] qs_np_id = qs_np.copy() b = np.ones(qs_np_id.shape[0]) qs_np_id = np.column_stack((qs_np_id,b)) if qs_np_tmp is None: if picture_similarity: qs_np_tmp = qs_np_id.copy() b = np.zeros(qs_np.shape[0]) qs_reurn = np.column_stack((qs_np,b)) else: qs_reurn = filterBox(qs_np, qs_np_tmp, pix_dis) if picture_similarity: qs_np_tmp = np.append(qs_np_tmp,qs_np_id,axis=0) qs_np_tmp[:, 11] += 1 qs_np_tmp = np.delete(qs_np_tmp, np.where((qs_np_tmp[:, 11] >= 75))[0], axis=0) has = False new_lab = [] for j in qs_reurn: if j[11] == 1: has = True new_lab.append(j[9]) if has: for q in qs_reurn: if q[11] >= 1: cls = int(q[9]) if not (cls in new_lab): continue # 为了防止其他类别被带出 code = str(int(q[10])).zfill(3) if det_xywh2.get(code) is None: det_xywh2[code] = {} cd = det_xywh2[code].get(cls) score = q[8] rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"] label_array, color = label_arrays[cls], rainbows[cls] box = [(int(q[0]), int(q[1])), (int(q[2]), int(q[3])), (int(q[4]), int(q[5])), (int(q[6]), int(q[7]))] is_new = False if q[11] == 1: is_new = True if cd is None: det_xywh2[code][cls] = [[cls, box, score, label_array, color, is_new]] else: det_xywh2[code][cls].append([cls, box, score, label_array, color, is_new]) if len(det_xywh2) > 0: put_queue(image_queue, (1, [det_xywh2, frame, frame_index_list[i], all_frames, draw_config["font_config"]])) push_p = push_stream_result.result(timeout=60) ai_video_file = write_ai_video_result.result(timeout=60) or_video_file = write_or_video_result.result(timeout=60) # 接收停止指令 if push_r[0] == 2: logger.info("拉流进程收到控制命令为:{}, requestId: {}",push_r[1] ,request_id) if 'algStart' == push_r[1]: self._algStatus = True;logger.info("算法识别开启, requestId: {}", request_id) if 'algStop' == push_r[1]: self._algStatus = False;logger.info("算法识别关闭, requestId: {}", request_id) if 'stop' == push_r[1]: logger.info("停止推流进程, requestId: {}", request_id) break if 'stop_ex' == push_r[1]: ex_status = False logger.info("停止推流进程, requestId: {}", request_id) break del push_r else: sleep(1) except ServiceException as s: logger.error("推流进程异常:{}, requestId:{}", s.msg, request_id) ex = s.code, s.msg except Exception: logger.error("推流进程异常:{}, requestId:{}", format_exc(), request_id) ex = ExceptionType.SERVICE_INNER_EXCEPTION.value[0], ExceptionType.SERVICE_INNER_EXCEPTION.value[1] finally: # 关闭推流管, 原视频写对象, 分析视频写对象 close_all_p(push_p, or_video_file, ai_video_file, request_id) if ex: code, msg = ex put_queue(push_ex_queue, (1, code, msg), timeout=2) else: if ex_status: # 关闭推流的时候, 等待1分钟图片队列处理完，如果1分钟内没有处理完, 清空图片队列, 丢弃没有上传的图片 c_time = time() while time() - c_time < 60: if image_queue.qsize() == 0 or image_queue.empty(): break sleep(2) for q in [push_queue, image_queue, hb_queue]: clear_queue(q) logger.info("推流进程停止完成！图片队列大小: {}, requestId:{}", image_queue.qsize(), request_id) class OffPushStreamProcess(PushStreamProcess): __slots__ = () def run(self): self.build_logo_url() msg, context = self._msg, self._context request_id = msg["request_id"] base_dir, env = context["base_dir"], context['env'] push_queue, image_queue, push_ex_queue, hb_queue = self._push_queue, self._image_queue, self._push_ex_queue, \ self._hb_queue aiFilePath, logo = context["aiFilePath"], context["logo"] ai_video_file, push_p, push_url = None, None, msg["push_url"] service_timeout = int(context["service"]["timeout"]) + 120 frame_score = context["service"]["filter"]["frame_score"] ex = None ex_status = True # 图片相似度开关 picture_similarity = bool(context["service"]["filter"]["picture_similarity"]) qs_np_tmp = None pix_dis = 60 if msg['taskType']==0: self._algStatus = False else: self._algStatus = True try: init_log(base_dir, env) logger.info("开始启动离线推流进程！requestId:{}", request_id) with ThreadPoolExecutor(max_workers=2) as t: # 定义三种推流、写原视频流、写ai视频流策略 # 第一个参数时间, 第二个参数重试次数 p_push_status, ai_write_status = [0, 0], [0, 0] start_time = time() while True: # 检测推流执行超时时间 if time() - start_time > service_timeout: logger.error("离线推流超时, requestId: {}", request_id) raise ServiceException(ExceptionType.TASK_EXCUTE_TIMEOUT.value[0], ExceptionType.TASK_EXCUTE_TIMEOUT.value[1]) # 系统由于各种问题可能会杀死内存使用多的进程, 自己杀掉自己 if psutil.Process(getpid()).ppid() == 1: logger.info("离线推流进程检测到父进程异常停止, 自动停止推流进程, requestId: {}", request_id) ex_status = False for q in [push_queue, image_queue, push_ex_queue, hb_queue]: clear_queue(q) break # 获取推流的视频帧 push_r = get_no_block_queue(push_queue) if push_r is not None: # [(1, ...] 视频帧操作 # [(2, 操作指令)] 指令操作 if push_r[0] == 1: frame_list, frame_index_list, all_frames, draw_config, push_objs = push_r[1] # 处理每一帧图片 for i, frame in enumerate(frame_list): pix_dis = int((frame.shape[0]//10)*1.2) if frame_index_list[i] % 300 == 0 and frame_index_list[i] <= all_frames: task_process = "%.2f" % (float(frame_index_list[i]) / float(all_frames)) put_queue(hb_queue, {"hb_value": task_process}, timeout=2) # 复制帧用来画图 copy_frame = frame.copy() # 所有问题记录字典 det_xywh, thread_p = {}, [] det_xywh2 = {} # 所有问题的矩阵集合 qs_np = None qs_reurn = [] for det in push_objs[i]: code, det_result = det # 每个单独模型处理 # 模型编号、100帧的所有问题, 检测目标、颜色、文字图片 if len(det_result) > 0: font_config, allowedList = draw_config["font_config"], draw_config[code]["allowedList"] rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"] for qs in det_result: box, score, cls = xywh2xyxy2(qs) if cls not in allowedList or score < frame_score: continue label_array, color = label_arrays[cls], rainbows[cls] if ModelType.CHANNEL2_MODEL.value[1] == str(code) and cls == 2: rr = t.submit(draw_name_joint, box, copy_frame, draw_config[code]["label_dict"], score, color, font_config, qs[6]) else: rr = t.submit(draw_painting_joint, box, copy_frame, label_array, score, color, font_config) thread_p.append(rr) if det_xywh.get(code) is None: det_xywh[code] = {} cd = det_xywh[code].get(cls) if not (ModelType.CHANNEL2_MODEL.value[1] == str(code) and cls == 2): if cd is None: det_xywh[code][cls] = [[cls, box, score, label_array, color]] else: det_xywh[code][cls].append([cls, box, score, label_array, color]) if qs_np is None: qs_np = np.array([box[0][0], box[0][1], box[1][0], box[1][1], box[2][0], box[2][1], box[3][0], box[3][1], score, cls, code],dtype=np.float32) else: result_li = np.array([box[0][0], box[0][1], box[1][0], box[1][1], box[2][0], box[2][1], box[3][0], box[3][1], score, cls, code],dtype=np.float32) qs_np = np.row_stack((qs_np, result_li)) if logo: frame = add_water_pic(frame, logo, request_id) copy_frame = add_water_pic(copy_frame, logo, request_id) if len(thread_p) > 0: for r in thread_p: r.result() if self._algSwitch and (not self._algStatus): # frame_merge = video_conjuncing(frame, frame.copy()) frame_merge = frame.copy() else: # frame_merge = video_conjuncing(frame, copy_frame) frame_merge = copy_frame # 写识别视频到本地 write_ai_video_result = t.submit(write_ai_video, frame_merge, aiFilePath, ai_video_file, ai_write_status, request_id) push_stream_result = t.submit(push_video_stream, frame_merge, push_p, push_url, p_push_status, request_id) if qs_np is not None: if len(qs_np.shape) == 1: qs_np = qs_np[np.newaxis,...] qs_np_id = qs_np.copy() b = np.ones(qs_np_id.shape[0]) qs_np_id = np.column_stack((qs_np_id,b)) if qs_np_tmp is None: if picture_similarity: qs_np_tmp = qs_np_id.copy() b = np.zeros(qs_np.shape[0]) qs_reurn = np.column_stack((qs_np,b)) else: qs_reurn = filterBox(qs_np, qs_np_tmp, pix_dis) if picture_similarity: qs_np_tmp = np.append(qs_np_tmp,qs_np_id,axis=0) qs_np_tmp[:, 11] += 1 qs_np_tmp = np.delete(qs_np_tmp, np.where((qs_np_tmp[:, 11] >= 75))[0], axis=0) has = False new_lab = [] for j in qs_reurn: if j[11] == 1: has = True new_lab.append(j[9]) if has: for q in qs_reurn: if q[11] >= 1: cls = int(q[9]) if not (cls in new_lab): continue # 为了防止其他类别被带出 code = str(int(q[10])).zfill(3) if det_xywh2.get(code) is None: det_xywh2[code] = {} cd = det_xywh2[code].get(cls) score = q[8] rainbows, label_arrays = draw_config[code]["rainbows"], draw_config[code]["label_arrays"] label_array, color = label_arrays[cls], rainbows[cls] box = [(int(q[0]), int(q[1])), (int(q[2]), int(q[3])), (int(q[4]), int(q[5])), (int(q[6]), int(q[7]))] is_new = False if q[11] == 1: is_new = True if cd is None: det_xywh2[code][cls] = [[cls, box, score, label_array, color, is_new]] else: det_xywh2[code][cls].append([cls, box, score, label_array, color, is_new]) if len(det_xywh2) > 0: put_queue(image_queue, (1, [det_xywh2, frame, frame_index_list[i], all_frames, draw_config["font_config"]])) push_p = push_stream_result.result(timeout=60) ai_video_file = write_ai_video_result.result(timeout=60) # 接收停止指令 if push_r[0] == 2: logger.info("拉流进程收到控制命令为:{}, requestId: {}",push_r[1] ,request_id) if 'algStart' == push_r[1]: self._algStatus = True;logger.info("算法识别开启, requestId: {}", request_id) if 'algStop' == push_r[1]: self._algStatus = False;logger.info("算法识别关闭, requestId: {}", request_id) if 'stop' == push_r[1]: logger.info("停止推流进程, requestId: {}", request_id) break if 'stop_ex' == push_r[1]: logger.info("停止推流进程, requestId: {}", request_id) ex_status = False break del push_r else: sleep(1) except ServiceException as s: logger.error("推流进程异常:{}, requestId:{}", s.msg, request_id) ex = s.code, s.msg except Exception: logger.error("推流进程异常:{}, requestId:{}", format_exc(), request_id) ex = ExceptionType.SERVICE_INNER_EXCEPTION.value[0], ExceptionType.SERVICE_INNER_EXCEPTION.value[1] finally: # 关闭推流管, 分析视频写对象 close_all_p(push_p, None, ai_video_file, request_id) if ex: code, msg = ex put_queue(push_ex_queue, (1, code, msg), timeout=2) else: if ex_status: # 关闭推流的时候, 等待1分钟图片队列处理完，如果1分钟内没有处理完, 清空图片队列, 丢弃没有上传的图片 c_time = time() while time() - c_time < 60: if image_queue.qsize() == 0 or image_queue.empty(): break sleep(2) for q in [push_queue, image_queue, hb_queue]: clear_queue(q) logger.info("推流进程停止完成！requestId:{}", request_id)