You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

201 lines
13KB

  1. # -*- coding: utf-8 -*-
  2. from concurrent.futures import ThreadPoolExecutor, wait, ALL_COMPLETED
  3. from os.path import join
  4. from threading import Thread
  5. from traceback import format_exc
  6. import cv2
  7. import numpy as np
  8. from loguru import logger
  9. from util.Cv2Utils import write_or_video, write_ai_video, push_video_stream, close_all_p, video_conjuncing
  10. from util.ImageUtils import url2Array, add_water_pic
  11. from util.PlotsUtils import draw_painting_joint
  12. from util.QueUtil import put_queue
  13. class OnPushStreamThread2(Thread):
  14. __slots__ = ('_msg', '_push_queue', '_context', 'ex', '_logo', '_image_queue')
  15. def __init__(self, *args):
  16. super().__init__()
  17. # 传参
  18. self._msg, self._push_queue, self._image_queue, self._context = args
  19. # 自带参数
  20. self.ex = None
  21. self._logo = None
  22. if self._context["video"]["video_add_water"]:
  23. self._logo = self._msg.get("logo_url")
  24. if self._logo:
  25. self._logo = url2Array(self._logo, enable_ex=False)
  26. if not self._logo:
  27. self._logo = cv2.imread(join(self._context['base_dir'], "image/logo.png"), -1)
  28. def run(self):
  29. request_id, push_queue, image_queue = self._msg.get("request_id"), self._push_queue, self._image_queue
  30. orFilePath, aiFilePath, logo = self._context.get("orFilePath"), self._context.get("aiFilePath"), self._logo
  31. or_video_file, ai_video_file, push_p = None, None, None
  32. push_url = self._msg.get("push_url")
  33. try:
  34. logger.info("开始启动推流线程!requestId:{}", request_id)
  35. with ThreadPoolExecutor(max_workers=2) as t:
  36. with ThreadPoolExecutor(max_workers=5) as tt:
  37. p_push_status, or_write_status, ai_write_status = [0, 0], [0, 0], [0, 0]
  38. while True:
  39. push_r = push_queue.get()
  40. if push_r is not None:
  41. # [(1, 原视频帧, 分析视频帧)]
  42. # [(code, retResults[2])]
  43. # [(2, 操作指令)]
  44. if push_r[0] == 1: # 视频帧操作
  45. frame_list, frame_index_list, all_frames = push_r[1]
  46. allowedList, rainbows, label_arrays, font_config = push_r[2]
  47. for i, frame in enumerate(frame_list):
  48. copy_frame = frame.copy()
  49. det_xywh = {}
  50. # 每帧可能存在多模型,多模型问题处理
  51. thread_p = []
  52. for det in push_r[3]:
  53. code, retResults = det
  54. det_xywh[code] = {}
  55. # 如果识别到了检测目标
  56. if len(retResults[i]) > 0:
  57. for qs in retResults[i]:
  58. detect_targets_code = int(qs[6])
  59. if detect_targets_code not in allowedList:
  60. logger.warning("当前检测目标不在检测目标中: {}, requestId: {}", detect_targets_code, request_id)
  61. continue
  62. score = qs[5]
  63. label_array = label_arrays[detect_targets_code]
  64. color = rainbows[detect_targets_code]
  65. if not isinstance(qs[1], (list, tuple, np.ndarray)):
  66. xc, yc, x2, y2 = int(qs[1]), int(qs[2]), int(qs[3]), int(qs[4])
  67. box = [(xc, yc), (x2, yc), (x2, y2), (xc, y2)]
  68. else:
  69. box = qs[1]
  70. # box, img, label_array, score=0.5, color=None, config=None
  71. dp = tt.submit(draw_painting_joint, box, copy_frame, label_array, score,
  72. color, font_config)
  73. thread_p.append(dp)
  74. cd = det_xywh[code].get(detect_targets_code)
  75. if cd is None:
  76. det_xywh[code][detect_targets_code] = [
  77. [detect_targets_code, box, score, label_array, color]]
  78. else:
  79. det_xywh[code][detect_targets_code].append(
  80. [detect_targets_code, box, score, label_array, color])
  81. if logo:
  82. frame = add_water_pic(frame, logo, request_id)
  83. copy_frame = add_water_pic(copy_frame, logo, request_id)
  84. if len(thread_p) > 0:
  85. completed_results = wait(thread_p, timeout=60, return_when=ALL_COMPLETED)
  86. completed_futures = completed_results.done
  87. for r in completed_futures:
  88. if r.exception():
  89. raise r.exception()
  90. frame_merge = video_conjuncing(frame, copy_frame)
  91. # 写原视频到本地
  92. write_or_video_result = t.submit(write_or_video, frame, orFilePath, or_video_file,
  93. or_write_status, request_id)
  94. # 写识别视频到本地
  95. write_ai_video_result = t.submit(write_ai_video, frame_merge, aiFilePath, ai_video_file,
  96. ai_write_status, request_id)
  97. if len(det_xywh) > 0:
  98. put_queue(image_queue, (1, (det_xywh, frame, frame_index_list[i], all_frames,
  99. font_config)))
  100. push_p = push_video_stream(frame_merge, push_p, push_url, p_push_status, request_id)
  101. ai_video_file = write_ai_video_result.result()
  102. or_video_file = write_or_video_result.result()
  103. if push_r[0] == 2:
  104. if 'stop' == push_r[1]:
  105. logger.info("停止推流线程, requestId: {}", request_id)
  106. close_all_p(push_p, or_video_file, ai_video_file, request_id)
  107. or_video_file, ai_video_file, push_p = None, None, None
  108. break
  109. except Exception as e:
  110. logger.error("推流线程异常:{}, requestId:{}", format_exc(), request_id)
  111. self.ex = e
  112. finally:
  113. close_all_p(push_p, or_video_file, ai_video_file, request_id)
  114. logger.info("推流线程停止完成!requestId:{}", request_id)
  115. # class OffPushStreamThread(Thread):
  116. # __slots__ = ('_msg', '_push_queue', '_context', 'ex', '_logo', '_image_queue')
  117. #
  118. # def __init__(self, *args):
  119. # super().__init__()
  120. # # 传参
  121. # self._msg, self._push_queue, self._image_queue, self._context = args
  122. # # 自带参数
  123. # self.ex = None
  124. # self._logo = None
  125. # if self._context["video"]["video_add_water"]:
  126. # self._logo = self._msg.get("logo_url")
  127. # if self._logo:
  128. # self._logo = url2Array(self._logo, enable_ex=False)
  129. # if not self._logo:
  130. # self._logo = cv2.imread(join(self._context['base_dir'], "image/logo.png"), -1)
  131. #
  132. # def run(self):
  133. # request_id, push_queue, image_queue = self._msg.get("request_id"), self._push_queue, self._image_queue
  134. # aiFilePath, logo = self._context.get("aiFilePath"), self._logo
  135. # ai_video_file, push_p = None, None
  136. # push_url = self._msg.get("push_url")
  137. # try:
  138. # logger.info("开始启动推流线程!requestId:{}", request_id)
  139. # with ThreadPoolExecutor(max_workers=1) as t:
  140. # p_push_status, or_write_status, ai_write_status = [0, 0], [0, 0], [0, 0]
  141. # while True:
  142. # push_parm = push_queue.get()
  143. # if push_parm is not None:
  144. # # [(1, 原视频帧, 分析视频帧)]
  145. # # # [视频帧、当前帧数、 总帧数、 [(问题数组、code、allowedList、label_arraylist、rainbows)]]
  146. # # res = (1, (pull_frame[1], pull_frame[2], pull_frame[3], []))
  147. # # [(2, 操作指令)]
  148. # if push_parm[0] == 1: # 视频帧操作
  149. # frame, current_frame, all_frames, ques_list = push_parm[1]
  150. # copy_frame = frame.copy()
  151. # det_xywh = {}
  152. # if len(ques_list) > 0:
  153. # for qs in ques_list:
  154. # det_xywh[qs[1]] = {}
  155. # detect_targets_code = int(qs[0][0])
  156. # score = qs[0][-1]
  157. # label_array = qs[3][detect_targets_code]
  158. # color = qs[4][detect_targets_code]
  159. # if not isinstance(qs[0][1], (list, tuple, np.ndarray)):
  160. # xc, yc, x2, y2 = int(qs[0][1]), int(qs[0][2]), int(qs[0][3]), int(qs[0][4])
  161. # box = [(xc, yc), (x2, yc), (x2, y2), (xc, y2)]
  162. # else:
  163. # box = qs[0][1]
  164. # draw_painting_joint(box, copy_frame, label_array, score, color, "leftTop")
  165. # cd = det_xywh[qs[1]].get(detect_targets_code)
  166. # if cd is None:
  167. # det_xywh[qs[1]][detect_targets_code] = [
  168. # [detect_targets_code, box, score, label_array, color]]
  169. # else:
  170. # det_xywh[qs[1]][detect_targets_code].append(
  171. # [detect_targets_code, box, score, label_array, color])
  172. # if logo:
  173. # frame = add_water_pic(frame, logo, request_id)
  174. # copy_frame = add_water_pic(copy_frame, logo, request_id)
  175. # frame_merge = video_conjuncing(frame, copy_frame)
  176. # # 写识别视频到本地
  177. # write_ai_video_result = t.submit(write_ai_video, frame_merge, aiFilePath, ai_video_file,
  178. # ai_write_status, request_id)
  179. # if len(det_xywh) > 0:
  180. # put_queue(image_queue, (1, (det_xywh, frame, current_frame, all_frames)))
  181. # push_p = push_video_stream(frame_merge, push_p, push_url, p_push_status, request_id)
  182. # ai_video_file = write_ai_video_result.result()
  183. # if push_parm[0] == 2:
  184. # if 'stop' == push_parm[1]:
  185. # logger.info("停止推流线程, requestId: {}", request_id)
  186. # close_all_p(push_p, None, ai_video_file, request_id)
  187. # ai_video_file, push_p = None, None
  188. # break
  189. # except Exception as e:
  190. # logger.error("推流线程异常:{}, requestId:{}", format_exc(), request_id)
  191. # self.ex = e
  192. # finally:
  193. # close_all_p(push_p, None, ai_video_file, request_id)
  194. # logger.info("推流线程停止完成!requestId:{}", request_id)