tuoheng_algN/concurrency/PushVideoStreamProcess.py

488 lines
30 KiB
Python
Raw Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

# -*- 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']
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())
else:
frame_merge = video_conjuncing(frame, 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()
print('----line384:',self._algSwitch,self._algStatus)
if self._algSwitch and (not self._algStatus):
frame_merge = video_conjuncing(frame, frame.copy())
else:
frame_merge = video_conjuncing(frame, 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)