201 lines
13 KiB
Python
201 lines
13 KiB
Python
|
# -*- coding: utf-8 -*-
|
|||
|
|
from concurrent.futures import ThreadPoolExecutor, wait, ALL_COMPLETED
|
|||
|
|
from os.path import join
|
|||
|
|
from threading import Thread
|
|||
|
|
from traceback import format_exc
|
|||
|
|
|
|||
|
|
import cv2
|
|||
|
|
import numpy as np
|
|||
|
|
from loguru import logger
|
|||
|
|
from util.Cv2Utils import write_or_video, write_ai_video, push_video_stream, close_all_p, video_conjuncing
|
|||
|
|
from util.ImageUtils import url2Array, add_water_pic
|
|||
|
|
from util.PlotsUtils import draw_painting_joint
|
|||
|
|
from util.QueUtil import put_queue
|
|||
|
|
|
|||
|
|
|
|||
|
|
class OnPushStreamThread2(Thread):
|
|||
|
|
__slots__ = ('_msg', '_push_queue', '_context', 'ex', '_logo', '_image_queue')
|
|||
|
|
|
|||
|
|
def __init__(self, *args):
|
|||
|
|
super().__init__()
|
|||
|
|
# 传参
|
|||
|
|
self._msg, self._push_queue, self._image_queue, self._context = args
|
|||
|
|
# 自带参数
|
|||
|
|
self.ex = None
|
|||
|
|
self._logo = None
|
|||
|
|
if self._context["video"]["video_add_water"]:
|
|||
|
|
self._logo = self._msg.get("logo_url")
|
|||
|
|
if self._logo:
|
|||
|
|
self._logo = url2Array(self._logo, enable_ex=False)
|
|||
|
|
if not self._logo:
|
|||
|
|
self._logo = cv2.imread(join(self._context['base_dir'], "image/logo.png"), -1)
|
|||
|
|
|
|||
|
|
def run(self):
|
|||
|
|
request_id, push_queue, image_queue = self._msg.get("request_id"), self._push_queue, self._image_queue
|
|||
|
|
orFilePath, aiFilePath, logo = self._context.get("orFilePath"), self._context.get("aiFilePath"), self._logo
|
|||
|
|
or_video_file, ai_video_file, push_p = None, None, None
|
|||
|
|
push_url = self._msg.get("push_url")
|
|||
|
|
try:
|
|||
|
|
logger.info("开始启动推流线程！requestId:{}", request_id)
|
|||
|
|
with ThreadPoolExecutor(max_workers=2) as t:
|
|||
|
|
with ThreadPoolExecutor(max_workers=5) as tt:
|
|||
|
|
p_push_status, or_write_status, ai_write_status = [0, 0], [0, 0], [0, 0]
|
|||
|
|
while True:
|
|||
|
|
push_r = push_queue.get()
|
|||
|
|
if push_r is not None:
|
|||
|
|
# [(1, 原视频帧, 分析视频帧)]
|
|||
|
|
# [(code, retResults[2])]
|
|||
|
|
# [(2, 操作指令)]
|
|||
|
|
if push_r[0] == 1: # 视频帧操作
|
|||
|
|
frame_list, frame_index_list, all_frames = push_r[1]
|
|||
|
|
allowedList, rainbows, label_arrays, font_config = push_r[2]
|
|||
|
|
for i, frame in enumerate(frame_list):
|
|||
|
|
copy_frame = frame.copy()
|
|||
|
|
det_xywh = {}
|
|||
|
|
# 每帧可能存在多模型，多模型问题处理
|
|||
|
|
thread_p = []
|
|||
|
|
for det in push_r[3]:
|
|||
|
|
code, retResults = det
|
|||
|
|
det_xywh[code] = {}
|
|||
|
|
# 如果识别到了检测目标
|
|||
|
|
if len(retResults[i]) > 0:
|
|||
|
|
for qs in retResults[i]:
|
|||
|
|
detect_targets_code = int(qs[6])
|
|||
|
|
if detect_targets_code not in allowedList:
|
|||
|
|
logger.warning("当前检测目标不在检测目标中: {}, requestId: {}", detect_targets_code, request_id)
|
|||
|
|
continue
|
|||
|
|
score = qs[5]
|
|||
|
|
label_array = label_arrays[detect_targets_code]
|
|||
|
|
color = rainbows[detect_targets_code]
|
|||
|
|
if not isinstance(qs[1], (list, tuple, np.ndarray)):
|
|||
|
|
xc, yc, x2, y2 = int(qs[1]), int(qs[2]), int(qs[3]), int(qs[4])
|
|||
|
|
box = [(xc, yc), (x2, yc), (x2, y2), (xc, y2)]
|
|||
|
|
else:
|
|||
|
|
box = qs[1]
|
|||
|
|
# box, img, label_array, score=0.5, color=None, config=None
|
|||
|
|
dp = tt.submit(draw_painting_joint, box, copy_frame, label_array, score,
|
|||
|
|
color, font_config)
|
|||
|
|
thread_p.append(dp)
|
|||
|
|
cd = det_xywh[code].get(detect_targets_code)
|
|||
|
|
if cd is None:
|
|||
|
|
det_xywh[code][detect_targets_code] = [
|
|||
|
|
[detect_targets_code, box, score, label_array, color]]
|
|||
|
|
else:
|
|||
|
|
det_xywh[code][detect_targets_code].append(
|
|||
|
|
[detect_targets_code, box, score, label_array, color])
|
|||
|
|
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:
|
|||
|
|
completed_results = wait(thread_p, timeout=60, return_when=ALL_COMPLETED)
|
|||
|
|
completed_futures = completed_results.done
|
|||
|
|
for r in completed_futures:
|
|||
|
|
if r.exception():
|
|||
|
|
raise r.exception()
|
|||
|
|
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)
|
|||
|
|
if len(det_xywh) > 0:
|
|||
|
|
put_queue(image_queue, (1, (det_xywh, frame, frame_index_list[i], all_frames,
|
|||
|
|
font_config)))
|
|||
|
|
push_p = push_video_stream(frame_merge, push_p, push_url, p_push_status, request_id)
|
|||
|
|
ai_video_file = write_ai_video_result.result()
|
|||
|
|
or_video_file = write_or_video_result.result()
|
|||
|
|
if push_r[0] == 2:
|
|||
|
|
if 'stop' == push_r[1]:
|
|||
|
|
logger.info("停止推流线程, requestId: {}", request_id)
|
|||
|
|
close_all_p(push_p, or_video_file, ai_video_file, request_id)
|
|||
|
|
or_video_file, ai_video_file, push_p = None, None, None
|
|||
|
|
break
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error("推流线程异常:{}, requestId:{}", format_exc(), request_id)
|
|||
|
|
self.ex = e
|
|||
|
|
finally:
|
|||
|
|
close_all_p(push_p, or_video_file, ai_video_file, request_id)
|
|||
|
|
logger.info("推流线程停止完成！requestId:{}", request_id)
|
|||
|
|
|
|||
|
|
|
|||
|
|
# class OffPushStreamThread(Thread):
|
|||
|
|
# __slots__ = ('_msg', '_push_queue', '_context', 'ex', '_logo', '_image_queue')
|
|||
|
|
#
|
|||
|
|
# def __init__(self, *args):
|
|||
|
|
# super().__init__()
|
|||
|
|
# # 传参
|
|||
|
|
# self._msg, self._push_queue, self._image_queue, self._context = args
|
|||
|
|
# # 自带参数
|
|||
|
|
# self.ex = None
|
|||
|
|
# self._logo = None
|
|||
|
|
# if self._context["video"]["video_add_water"]:
|
|||
|
|
# self._logo = self._msg.get("logo_url")
|
|||
|
|
# if self._logo:
|
|||
|
|
# self._logo = url2Array(self._logo, enable_ex=False)
|
|||
|
|
# if not self._logo:
|
|||
|
|
# self._logo = cv2.imread(join(self._context['base_dir'], "image/logo.png"), -1)
|
|||
|
|
#
|
|||
|
|
# def run(self):
|
|||
|
|
# request_id, push_queue, image_queue = self._msg.get("request_id"), self._push_queue, self._image_queue
|
|||
|
|
# aiFilePath, logo = self._context.get("aiFilePath"), self._logo
|
|||
|
|
# ai_video_file, push_p = None, None
|
|||
|
|
# push_url = self._msg.get("push_url")
|
|||
|
|
# try:
|
|||
|
|
# logger.info("开始启动推流线程！requestId:{}", request_id)
|
|||
|
|
# with ThreadPoolExecutor(max_workers=1) as t:
|
|||
|
|
# p_push_status, or_write_status, ai_write_status = [0, 0], [0, 0], [0, 0]
|
|||
|
|
# while True:
|
|||
|
|
# push_parm = push_queue.get()
|
|||
|
|
# if push_parm is not None:
|
|||
|
|
# # [(1, 原视频帧, 分析视频帧)]
|
|||
|
|
# # # [视频帧、当前帧数、 总帧数、 [(问题数组、code、allowedList、label_arraylist、rainbows)]]
|
|||
|
|
# # res = (1, (pull_frame[1], pull_frame[2], pull_frame[3], []))
|
|||
|
|
# # [(2, 操作指令)]
|
|||
|
|
# if push_parm[0] == 1: # 视频帧操作
|
|||
|
|
# frame, current_frame, all_frames, ques_list = push_parm[1]
|
|||
|
|
# copy_frame = frame.copy()
|
|||
|
|
# det_xywh = {}
|
|||
|
|
# if len(ques_list) > 0:
|
|||
|
|
# for qs in ques_list:
|
|||
|
|
# det_xywh[qs[1]] = {}
|
|||
|
|
# detect_targets_code = int(qs[0][0])
|
|||
|
|
# score = qs[0][-1]
|
|||
|
|
# label_array = qs[3][detect_targets_code]
|
|||
|
|
# color = qs[4][detect_targets_code]
|
|||
|
|
# if not isinstance(qs[0][1], (list, tuple, np.ndarray)):
|
|||
|
|
# xc, yc, x2, y2 = int(qs[0][1]), int(qs[0][2]), int(qs[0][3]), int(qs[0][4])
|
|||
|
|
# box = [(xc, yc), (x2, yc), (x2, y2), (xc, y2)]
|
|||
|
|
# else:
|
|||
|
|
# box = qs[0][1]
|
|||
|
|
# draw_painting_joint(box, copy_frame, label_array, score, color, "leftTop")
|
|||
|
|
# cd = det_xywh[qs[1]].get(detect_targets_code)
|
|||
|
|
# if cd is None:
|
|||
|
|
# det_xywh[qs[1]][detect_targets_code] = [
|
|||
|
|
# [detect_targets_code, box, score, label_array, color]]
|
|||
|
|
# else:
|
|||
|
|
# det_xywh[qs[1]][detect_targets_code].append(
|
|||
|
|
# [detect_targets_code, box, score, label_array, color])
|
|||
|
|
# if logo:
|
|||
|
|
# frame = add_water_pic(frame, logo, request_id)
|
|||
|
|
# copy_frame = add_water_pic(copy_frame, logo, request_id)
|
|||
|
|
# 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)
|
|||
|
|
# if len(det_xywh) > 0:
|
|||
|
|
# put_queue(image_queue, (1, (det_xywh, frame, current_frame, all_frames)))
|
|||
|
|
# push_p = push_video_stream(frame_merge, push_p, push_url, p_push_status, request_id)
|
|||
|
|
# ai_video_file = write_ai_video_result.result()
|
|||
|
|
# if push_parm[0] == 2:
|
|||
|
|
# if 'stop' == push_parm[1]:
|
|||
|
|
# logger.info("停止推流线程, requestId: {}", request_id)
|
|||
|
|
# close_all_p(push_p, None, ai_video_file, request_id)
|
|||
|
|
# ai_video_file, push_p = None, None
|
|||
|
|
# break
|
|||
|
|
# except Exception as e:
|
|||
|
|
# logger.error("推流线程异常:{}, requestId:{}", format_exc(), request_id)
|
|||
|
|
# self.ex = e
|
|||
|
|
# finally:
|
|||
|
|
# close_all_p(push_p, None, ai_video_file, request_id)
|
|||
|
|
# logger.info("推流线程停止完成！requestId:{}", request_id)
|