tuoheng_alg_move_20240715/concurrency/PushVideoStreamProcess.py

# -*- coding: utf-8 -*-

from concurrent.futures import ThreadPoolExecutor
from multiprocessing import Process
from os import getpid

from os.path import join
from time import time, sleep
from traceback import format_exc

import cv2
import psutil

from loguru import logger

from enums.ExceptionEnum import ExceptionType
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, xywh2xyxy2

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

    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
        high_score_image = {}
        # 相似度, 默认值0.65
        similarity = context["service"]["filter"]["similarity"]
        # 图片相似度开关
        picture_similarity = bool(context["service"]["filter"]["picture_similarity"])
        frame_step = int(context["service"]["filter"]["frame_step"])
        try:
            init_log(base_dir, env)
            logger.info("开始实时启动推流进程！requestId:{}", request_id)
            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):
                                # 复制帧用来画图
                                copy_frame = frame.copy()
                                det_xywh, thread_p = {}, []
                                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]
                                            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 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 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()
                                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 high_score_image.get("current_frame") is None:
                                    high_score_image["current_frame"] = frame_index_list[i]
                                diff_frame_num = frame_index_list[i] - high_score_image["current_frame"]
                                # 指定帧内处理逻辑
                                if frame_step >= diff_frame_num and high_score_image.get("code") is not None and len(det_xywh) > 0:
                                    # 所有的模型编号
                                    cache_codes = set(high_score_image["code"].keys())
                                    # 遍历当前模型
                                    for det_code, det_clss in det_xywh.items():
                                        # 如果模型编号不在缓存中
                                        if det_code not in cache_codes:
                                            high_score_image["code"][det_code] = {}
                                            for cls, cls_list in det_xywh[det_code].items():
                                                if len(cls_list) > 0:
                                                    high_score_image["code"][det_code][cls] = (frame, frame_index_list[i], cls_list)
                                        else:
                                            # 如果模型编号在缓存中, 检查检测目标的情况
                                            cache_clss = set(high_score_image["code"][det_code].keys())
                                            for cls, cls_list in det_xywh[det_code].items():
                                                if len(cls_list) > 0:
                                                    if cls not in cache_clss:
                                                        high_score_image["code"][det_code][cls] = (frame, frame_index_list[i], cls_list)
                                                    else:
                                                        if len(det_xywh[det_code][cls]) > len(high_score_image["code"][det_code][cls][2]):
                                                            high_score_image["code"][det_code][cls] = (frame, frame_index_list[i], cls_list)
                                                        elif len(det_xywh[det_code][cls]) == len(high_score_image["code"][det_code][cls][2]):
                                                            # [cls, box, score, label_array, color]
                                                            hs = 0
                                                            for ii, s in enumerate(cls_list):
                                                                if s[2] >= high_score_image["code"][det_code][cls][2][ii][2]:
                                                                    hs += 1
                                                            if hs == len(cls_list):
                                                                high_score_image["code"][det_code][cls] = (frame, frame_index_list[i], cls_list)
                                if diff_frame_num > frame_step:
                                    if high_score_image.get("code") is not None:
                                        high_score_image["or_frame"] = frame
                                        put_queue(image_queue, (1, [high_score_image["code"], all_frames, draw_config["font_config"]]))
                                    high_score_image["code"] = None
                                    high_score_image["current_frame"] = None
                                # 如果有问题, 走下面的逻辑
                                if len(det_xywh) > 0:
                                    flag = True
                                    # 比较上一次识别到问题的帧和当前问题帧的相似度
                                    if picture_similarity and high_score_image.get("or_frame") is not None:
                                        hash1 = ImageUtils.dHash(high_score_image["or_frame"])
                                        hash2 = ImageUtils.dHash(frame)
                                        dist = ImageUtils.Hamming_distance(hash1, hash2)
                                        similarity_1 = 1 - dist * 1.0 / 64
                                        if similarity_1 >= similarity:
                                            flag = False
                                    # 如果缓存问题是空的，给缓存添加问题
                                    if flag and high_score_image.get("code") is None:
                                        high_score_image["code"] = {}
                                        for code, det_list in det_xywh.items():
                                            if high_score_image["code"].get(code) is None:
                                                high_score_image["code"][code] = {}
                                            for cls, cls_list in det_list.items():
                                                high_score_image["code"][code][cls] = (frame, frame_index_list[i], cls_list)
                                push_p = push_stream_result.result(timeout=5)
                                ai_video_file = write_ai_video_result.result(timeout=5)
                                or_video_file = write_or_video_result.result(timeout=5)
                        # 接收停止指令
                        if push_r[0] == 2:
                            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, push_ex_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
        high_score_image = {}
        # 相似度, 默认值0.65
        similarity = context["service"]["filter"]["similarity"]
        # 图片相似度开关
        picture_similarity = bool(context["service"]["filter"]["picture_similarity"])
        frame_step = int(context["service"]["filter"]["frame_step"])
        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):
                                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 = {}, []
                                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]
                                            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 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 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()
                                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 high_score_image.get("current_frame") is None:
                                    high_score_image["current_frame"] = frame_index_list[i]
                                diff_frame_num = frame_index_list[i] - high_score_image["current_frame"]
                                # 指定帧内处理逻辑
                                if frame_step >= diff_frame_num and high_score_image.get("code") is not None and len(det_xywh) > 0:
                                    # 所有的模型编号
                                    cache_codes = set(high_score_image["code"].keys())
                                    # 遍历当前模型
                                    for det_code, det_clss in det_xywh.items():
                                        # 如果模型编号不在缓存中
                                        if det_code not in cache_codes:
                                            high_score_image["code"][det_code] = {}
                                            for cls, cls_list in det_xywh[det_code].items():
                                                if len(cls_list) > 0:
                                                    high_score_image["code"][det_code][cls] = (frame, frame_index_list[i], cls_list)
                                        else:
                                            # 如果模型编号在缓存中, 检查检测目标的情况
                                            cache_clss = set(high_score_image["code"][det_code].keys())
                                            for cls, cls_list in det_xywh[det_code].items():
                                                if len(cls_list) > 0:
                                                    # 如果检测目标在缓存中
                                                    if cls not in cache_clss:
                                                        high_score_image["code"][det_code][cls] = (frame, frame_index_list[i], cls_list)
                                                    else:
                                                        if len(det_xywh[det_code][cls]) > len(high_score_image["code"][det_code][cls][2]):
                                                            high_score_image["code"][det_code][cls] = (frame, frame_index_list[i], cls_list)
                                                        elif len(det_xywh[det_code][cls]) == len(high_score_image["code"][det_code][cls][2]):
                                                            # [cls, box, score, label_array, color]
                                                            hs = 0
                                                            for ii, s in enumerate(cls_list):
                                                                if s[2] >= high_score_image["code"][det_code][cls][2][ii][2]:
                                                                    hs += 1
                                                            if hs == len(cls_list):
                                                                high_score_image["code"][det_code][cls] = (frame, frame_index_list[i], cls_list)
                                if diff_frame_num > frame_step:
                                    if high_score_image.get("code") is not None:
                                        high_score_image["or_frame"] = frame
                                        put_queue(image_queue, (1, [high_score_image["code"], all_frames, draw_config["font_config"]]))
                                    high_score_image["code"] = None
                                    high_score_image["current_frame"] = None
                                # 如果有问题, 走下面的逻辑
                                if len(det_xywh) > 0:
                                    flag = True
                                    # 比较上一次识别到问题的帧和当前问题帧的相似度
                                    if picture_similarity and high_score_image.get("or_frame") is not None:
                                        hash1 = ImageUtils.dHash(high_score_image["or_frame"])
                                        hash2 = ImageUtils.dHash(frame)
                                        dist = ImageUtils.Hamming_distance(hash1, hash2)
                                        similarity_1 = 1 - dist * 1.0 / 64
                                        if similarity_1 >= similarity:
                                            flag = False
                                    # 如果缓存问题是空的，给缓存添加问题
                                    if flag and high_score_image.get("code") is None:
                                        high_score_image["code"] = {}
                                        for code, det_list in det_xywh.items():
                                            if high_score_image["code"].get(code) is None:
                                                high_score_image["code"][code] = {}
                                            for cls, cls_list in det_list.items():
                                                high_score_image["code"][code][cls] = (frame, frame_index_list[i], cls_list)
                                push_p = push_stream_result.result(timeout=5)
                                ai_video_file = write_ai_video_result.result(timeout=5)
                        # 接收停止指令
                        if push_r[0] == 2:
                            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, push_ex_queue, hb_queue]:
                clear_queue(q)
        logger.info("推流进程停止完成！requestId:{}", request_id)