tuoheng_alg_fangke_discard_20240715/util/Cv2Utils.py

# -*- coding: utf-8 -*-
import time

import cv2
import subprocess as sp

import ffmpeg
import numpy as np
from loguru import logger
from exception.CustomerException import ServiceException
from enums.ExceptionEnum import ExceptionType


class Cv2Util():

    def __init__(self, pullUrl, pushUrl=None, orFilePath=None, aiFilePath=None, requestId=None):
        self.pullUrl = pullUrl
        self.pushUrl = pushUrl
        self.orFilePath = orFilePath
        self.aiFilePath = aiFilePath
        self.cap = None
        self.p = None
        self.or_video_file = None
        self.ai_video_file = None
        self.fps = None
        self.width = None
        self.height = None
        self.wh = None
        self.all_frames = None
        self.bit_rate = None
        self.pull_p = None
        self.requestId = requestId
        self.p_push_retry_num = 0

    '''
        获取视频信息
    '''

    def get_video_info(self):
        try:
            if self.pullUrl is None:
                logger.error("拉流地址不能为空, requestId:{}", self.requestId)
                raise ServiceException(ExceptionType.PULL_STREAM_URL_EXCEPTION.value[0],
                                       ExceptionType.PULL_STREAM_URL_EXCEPTION.value[1])
            probe = ffmpeg.probe(self.pullUrl)
            # 视频大小
            # format = probe['format']
            # size = int(format['size'])/1024/1024
            video_stream = next((stream for stream in probe['streams'] if stream['codec_type'] == 'video'), None)
            if video_stream is None:
                logger.error("根据拉流地址未获取到视频流, requestId:{}", self.requestId)
                return
            width = video_stream.get('width')
            height = video_stream.get('height')
            nb_frames = video_stream.get('nb_frames')
            fps = video_stream.get('r_frame_rate')
            # duration = video_stream.get('duration')
            bit_rate = video_stream.get('bit_rate')
            if width:
                self.width = int(width)
            if height:
                self.height = int(height)
            if width is not None and height is not None:
                self.wh = int(width * height * 3)
            if nb_frames:
                self.all_frames = int(nb_frames)
            if fps:
                up, down = str(fps).split('/')
                self.fps = int(eval(up) / eval(down))
            # if duration:
            #    self.duration = float(video_stream['duration'])
            if bit_rate:
                self.bit_rate = int(bit_rate) / 1000
            logger.info("视频信息, width:{}|height:{}|fps:{}|all_frames:{}|bit_rate:{}, requestId:{}", self.width,
                        self.height, self.fps, self.all_frames, self.bit_rate, self.requestId)
        except ServiceException as s:
            logger.exception("获取视频信息异常: {}, requestId:{}", s.msg, self.requestId)
            raise s
        except Exception as e:
            logger.exception("获取视频信息异常：{}, requestId:{}", e, self.requestId)

    '''
        拉取视频
    '''

    def build_pull_p(self):
        try:
            if self.width is None or self.height is None:
                return
            if self.pull_p:
                logger.info("重试, 关闭拉流管道, requestId:{}", self.requestId)
                self.pull_p.stdout.close()
                self.pull_p.terminate()
                self.pull_p.wait()
            if self.pullUrl is None:
                logger.error("拉流地址不能为空, requestId:{}", self.requestId)
                raise ServiceException(ExceptionType.PULL_STREAM_URL_EXCEPTION.value[0],
                                       ExceptionType.PULL_STREAM_URL_EXCEPTION.value[1])
            # command = ['ffmpeg',
            #            # '-b:v', '3000k',
            #            '-i', self.pullUrl,
            #            '-f', 'rawvideo',
            #            '-vcodec', 'rawvideo',
            #            '-pix_fmt', 'bgr24',
            #            # '-s', "{}x{}".format(int(width), int(height)),
            #            '-an',
            #            '-']
            aa = {'loglevel': 'error', 'c:v': 'h264_cuvid'}
            process = (
                ffmpeg
                    .input(self.pullUrl, **aa)
                    .output('pipe:', format='rawvideo', pix_fmt='bgr24', loglevel='error')
                    .overwrite_output()
                    .global_args('-an')
                    .run_async(pipe_stdout=True)
            )
            # self.pull_p = sp.Popen(command, stdout=sp.PIPE, stderr=sp.PIPE)
            self.pull_p = process
        except ServiceException as s:
            logger.exception("构建拉流管道异常: {}, requestId:{}", s, self.requestId)
            raise s
        except Exception as e:
            logger.exception("构建拉流管道异常：{}, requestId:{}", e, self.requestId)

    def checkconfig(self):
        if self.fps is None or self.width is None or self.height is None:
            return True
        return False

    def read(self):
        result = None
        try:
            # if self.pull_p is None:
            #     logger.error("拉流管道为空, requestId:{}", self.requestId)
            #     raise ServiceException(ExceptionType.PULL_PIPELINE_INIT_EXCEPTION.value[0],
            #                            ExceptionType.PULL_PIPELINE_INIT_EXCEPTION.value[1])
            in_bytes = self.pull_p.stdout.read(self.wh)
            if in_bytes is not None and len(in_bytes) > 0:
                result = (np.frombuffer(in_bytes, np.uint8).reshape([int(self.height), int(self.width), 3]))
                result = cv2.resize(result, (int(self.width / 2), int(self.height / 2)), interpolation=cv2.INTER_LINEAR)
        except ServiceException as s:
            logger.exception("读流异常: {}, requestId:{}", s, self.requestId)
        except Exception as e:
            logger.exception("读流异常：{}, requestId:{}", e, self.requestId)
        return result

    def close(self):
        if self.pull_p:
            if self.pull_p.stdout:
                self.pull_p.stdout.close()
            self.pull_p.terminate()
            self.pull_p.wait()
            logger.info("关闭拉流管道完成, requestId:{}", self.requestId)
        if self.p:
            if self.p.stdin:
                self.p.stdin.close()
            self.p.terminate()
            self.p.wait()
            # self.p.communicate()
            # self.p.kill()
            logger.info("关闭管道完成, requestId:{}", self.requestId)
        if self.or_video_file:
            self.or_video_file.release()
            logger.info("关闭原视频写入流完成, requestId:{}", self.requestId)
        if self.ai_video_file:
            self.ai_video_file.release()
            logger.info("关闭AI视频写入流完成, requestId:{}", self.requestId)

    # 构建 cv2
    # def build_cv2(self):
    #     try:
    #         if self.cap is not None:
    #             logger.info("重试, 关闭cap, requestId:{}", self.requestId)
    #             self.cap.release()
    #         if self.p is not None:
    #             logger.info("重试, 关闭管道, requestId:{}", self.requestId)
    #             self.p.stdin.close()
    #             self.p.terminate()
    #             self.p.wait()
    #         if self.pullUrl is None:
    #             logger.error("拉流地址不能为空, requestId:{}", self.requestId)
    #             raise ServiceException(ExceptionType.PULL_STREAM_URL_EXCEPTION.value[0],
    #                                    ExceptionType.PULL_STREAM_URL_EXCEPTION.value[1])
    #         if self.pushUrl is None:
    #             logger.error("推流地址不能为空, requestId:{}", self.requestId)
    #             raise ServiceException(ExceptionType.PUSH_STREAM_URL_EXCEPTION.value[0],
    #                                    ExceptionType.PUSH_STREAM_URL_EXCEPTION.value[1])
    #         self.cap = cv2.VideoCapture(self.pullUrl)
    #         if self.fps is None or self.fps == 0:
    #             self.fps = int(self.cap.get(cv2.CAP_PROP_FPS))
    #         if self.width is None or self.width == 0:
    #             self.width = int(self.cap.get(cv2.CAP_PROP_FRAME_WIDTH))
    #         if self.height is None or self.height == 0:
    #             self.height = int(self.cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    #         command = ['/usr/bin/ffmpeg',
    #                    # '-y',  # 不经过确认，输出时直接覆盖同名文件。
    #                    '-f', 'rawvideo',
    #                    '-vcodec', 'rawvideo',
    #                    '-pix_fmt', 'bgr24',  # 显示可用的像素格式
    #                    # '-s', "{}x{}".format(self.width * 2, self.height),
    #                    '-s', "{}x{}".format(int(self.width), int(self.height/2)),
    #                    # '-r', str(15),
    #                    '-i', '-',  # 指定输入文件
    #                    '-g', '25',
    #                    '-b:v', '3000k',
    #                    '-tune', 'zerolatency',  # 加速编码速度
    #                    '-c:v', 'libx264',  # 指定视频编码器
    #                    '-sc_threshold', '0',
    #                    '-pix_fmt', 'yuv420p',
    #                    '-an',
    #                    '-preset', 'ultrafast',  # 指定输出的视频质量，会影响文件的生成速度，有以下几个可用的值 ultrafast,
    #                    # superfast, veryfast, faster, fast, medium, slow, slower, veryslow。
    #                    '-f', 'flv',
    #                    self.pushUrl]
    #         # 管道配置
    #         logger.info("fps:{}|height:{}|width:{}|requestId:{}", self.fps, self.height, self.width, self.requestId)
    #         self.p = sp.Popen(command, stdin=sp.PIPE)
    #     except ServiceException as s:
    #         logger.exception("构建cv2异常: {}, requestId:{}", s, self.requestId)
    #         raise s
    #     except Exception as e:
    #         logger.exception("初始化cv2异常：{}, requestId:{}", e, self.requestId)

    # 构建 cv2
    def build_p(self):
        try:
            if self.p:
                logger.info("重试, 关闭管道, requestId:{}", self.requestId)
                if self.p.stdin:
                    self.p.stdin.close()
                self.p.terminate()
                self.p.wait()
                # self.p.communicate()
                # self.p.kill()
            if self.pushUrl is None:
                logger.error("推流地址不能为空, requestId:{}", self.requestId)
                raise ServiceException(ExceptionType.PUSH_STREAM_URL_EXCEPTION.value[0],
                                       ExceptionType.PUSH_STREAM_URL_EXCEPTION.value[1])
            command = ['ffmpeg',
                       # '-loglevel', 'debug',
                       '-y',  # 不经过确认，输出时直接覆盖同名文件。
                       '-f', 'rawvideo',
                       '-vcodec', 'rawvideo',
                       '-pix_fmt', 'bgr24',
                       '-thread_queue_size', '16',
                       # '-s', "{}x{}".format(self.width * 2, self.height),
                       '-s', "{}x{}".format(int(self.width), int(self.height / 2)),
                       '-r', str(self.fps),
                       '-i', '-',  # 指定输入文件
                       '-g', str(self.fps),
                       # '-maxrate', '15000k',
                       # '-profile:v', 'high',
                       # '-b:v', '4000k',
                       # '-crf', '18',
                       '-rc:v', 'vbr',
                       '-cq:v', '30',
                       '-qmin', '30',
                       '-qmax', '30',
                       '-c:v', 'h264_nvenc',  #
                       # '-bufsize', '4000k',
                       # '-c:v', 'libx264',  # 指定视频编码器
                       # '-tune', 'zerolatency',  # 加速编码速度
                       # '-sc_threshold', '0',
                       '-pix_fmt', 'yuv420p',
                       "-an",
                       # '-flvflags', 'no_duration_filesize',
                       '-preset', 'fast',  # 指定输出的视频质量，会影响文件的生成速度，有以下几个可用的值 ultrafast,
                       # superfast, veryfast, faster, fast, medium, slow, slower, veryslow。
                       '-f', 'flv',
                       self.pushUrl]
            # command = 'ffmpeg -loglevel debug -y -f rawvideo -vcodec rawvideo -pix_fmt bgr24' +\
            #            ' -s ' +  "{}x{}".format(int(self.width), int(self.height/2))\
            #            + ' -i - ' + '-g ' + str(self.fps)+\
            #            ' -b:v 6000k -tune zerolatency -c:v libx264 -pix_fmt yuv420p -preset ultrafast'+\
            #            ' -f flv ' + self.pushUrl

            # kwargs = {'format': 'rawvideo',
            #           # 'vcodec': 'rawvideo',
            #           'pix_fmt': 'bgr24',
            #           's': '{}x{}'.format(int(self.width), int(self.height/2))}
            # out = {
            #     'r': str(self.fps),
            #     'g': str(self.fps),
            #     'b:v': '5500k', # 恒定码率
            #     # 'maxrate': '15000k',
            #     # 'crf': '18',
            #     'bufsize': '5500k',
            #     'tune': 'zerolatency', # 加速编码速度
            #     'c:v': 'libx264',  # 指定视频编码器
            #     'sc_threshold': '0',
            #     'pix_fmt': 'yuv420p',
            #     # 'flvflags': 'no_duration_filesize',
            #     'preset': 'medium',  # 指定输出的视频质量，会影响文件的生成速度，有以下几个可用的值 ultrafast,
            #     # superfast, veryfast, faster, fast, medium, slow, slower, veryslow。
            #     'format': 'flv'}
            # 管道配置
            # process2 = (
            #     ffmpeg
            #         .input('pipe:', **kwargs)
            #         .output(self.pushUrl, **out)
            #         .global_args('-y', '-an')
            #         .overwrite_output()
            #         .run_async(pipe_stdin=True)
            # )
            logger.info("fps:{}|height:{}|width:{}|requestId:{}", self.fps, self.height, self.width, self.requestId)
            self.p = sp.Popen(command, stdin=sp.PIPE, shell=False)
            # self.p = process2
        except ServiceException as s:
            logger.exception("构建p管道异常: {}, requestId:{}", s, self.requestId)
            raise s
        except Exception as e:
            logger.exception("初始化p管道异常：{}, requestId:{}", e, self.requestId)

    def push_stream(self, frame):
        if self.p is None:
            self.build_p()
        try:
            self.p.stdin.write(frame.tostring())
        except Exception as ex:
            logger.exception("推流进管道异常：{}, requestId: {}", ex, self.requestId)
            current_retry_num = 0
            while True:
                try:
                    if self.p_push_retry_num > 20:
                        logger.info("推流失败重试次数过多, 请检查相关配置信息, 当前重试次数: {}, requestId: {}",
                                    self.p_push_retry_num, self.requestId)
                        current_retry_num = 4
                        break
                    self.p_push_retry_num += 1
                    time.sleep(10)
                    self.build_p()
                    self.p.stdin.write(frame.tostring())
                    logger.info("构建p管道重试成功, 当前重试次数: {}, requestId: {}", current_retry_num,
                                self.requestId)
                    break
                except Exception as e:
                    current_retry_num += 1
                    logger.exception("构建p管道异常：{}, 开始重试, 当前重试次数：{}, requestId: {}", e,
                                     current_retry_num, self.requestId)
                    if current_retry_num > 3:
                        logger.exception("构建p管道重试失败：{}, requestId: {}", e, self.requestId)
                        break
            if current_retry_num > 3:
                raise ServiceException(ExceptionType.PUSH_STREAM_URL_E_EXCEPTION.value[0],
                                       ExceptionType.PUSH_STREAM_URL_E_EXCEPTION.value[1])

    def video_write(self, or_frame, ai_frame):
        try:
            self.build_write()
            if or_frame is not None and len(or_frame) > 0:
                self.or_video_file.write(or_frame)
            if ai_frame is not None and len(ai_frame) > 0:
                self.ai_video_file.write(ai_frame)
        except Exception as ex:
            ai_retry_num = 0
            while True:
                try:
                    if or_frame is not None and len(or_frame) > 0:
                        self.or_video_file.write(or_frame)
                    if ai_frame is not None and len(ai_frame) > 0:
                        self.ai_video_file.write(ai_frame)
                    logger.info("重新写入离线分析后视频到本地, 当前重试次数: {}, requestId: {}", ai_retry_num,
                                self.requestId)
                    break
                except Exception as e:
                    ai_retry_num += 1
                    logger.exception("重新写入离线分析后视频到本地：{}, 开始重试, 当前重试次数：{}, requestId: {}", e,
                                     ai_retry_num, self.requestId)
                    if ai_retry_num > 3:
                        logger.exception("重新写入离线分析后视频到本地,重试失败：{}, requestId: {}", e, self.requestId)
                        raise ServiceException(ExceptionType.SERVICE_INNER_EXCEPTION.value[0],
                                               ExceptionType.SERVICE_INNER_EXCEPTION.value[1])

    def build_write(self):
        try:
            if self.fps is None or self.width is None or self.height is None:
                raise ServiceException(ExceptionType.VIDEO_CONFIG_EXCEPTION.value[0],
                                       ExceptionType.VIDEO_CONFIG_EXCEPTION.value[1])
            if self.orFilePath is not None and self.or_video_file is None:
                self.or_video_file = cv2.VideoWriter(self.orFilePath, cv2.VideoWriter_fourcc(*'mp4v'), self.fps,
                                                     (int(self.width / 2), int(self.height / 2)))
                if self.or_video_file is None:
                    raise ServiceException(ExceptionType.OR_WRITE_OBJECT_EXCEPTION.value[0],
                                           ExceptionType.OR_WRITE_OBJECT_EXCEPTION.value[1])
            if self.aiFilePath is not None and self.ai_video_file is None:
                self.ai_video_file = cv2.VideoWriter(self.aiFilePath, cv2.VideoWriter_fourcc(*'mp4v'), self.fps,
                                                     (int(self.width), int(self.height / 2)))
                if self.ai_video_file is None:
                    raise ServiceException(ExceptionType.AI_WRITE_OBJECT_EXCEPTION.value[0],
                                           ExceptionType.AI_WRITE_OBJECT_EXCEPTION.value[1])
        except ServiceException as s:
            logger.exception("构建文件写对象异常: {}, requestId:{}", s, self.requestId)
            raise s
        except Exception as e:
            logger.exception("构建文件写对象异常: {}, requestId:{}", e, self.requestId)
            raise e

    def video_merge(self, frame1, frame2):
        # frameLeft = cv2.resize(frame1, (int(self.width / 2), int(self.height / 2)), interpolation=cv2.INTER_LINEAR)
        # frameRight = cv2.resize(frame2, (int(self.width / 2), int(self.height / 2)), interpolation=cv2.INTER_LINEAR)
        # frame_merge = np.hstack((frameLeft, frameRight))
        frame_merge = np.hstack((frame1, frame2))
        return frame_merge

    def getP(self):
        if self.p is None:
            logger.error("获取管道为空, requestId:{}", self.requestId)
            raise ServiceException(ExceptionType.PULL_PIPELINE_INIT_EXCEPTION.value[0],
                                   ExceptionType.PULL_PIPELINE_INIT_EXCEPTION.value[1])
        return self.p

    def getCap(self):
        if self.cap is None:
            logger.error("获取cv2为空, requestId:{}", self.requestId)
            raise ServiceException(ExceptionType.CV2_IS_NULL_EXCEPTION.value[0],
                                   ExceptionType.CV2_IS_NULL_EXCEPTION.value[1])
        return self.cap

    def getOrVideoFile(self):
        if self.or_video_file is None:
            logger.error("获取原视频写入对象为空, requestId:{}", self.requestId)
            raise ServiceException(ExceptionType.OR_WRITE_OBJECT_EXCEPTION.value[0],
                                   ExceptionType.OR_WRITE_OBJECT_EXCEPTION.value[1])
        return self.or_video_file

    def getAiVideoFile(self):
        if self.ai_video_file is None:
            logger.error("获取AI视频写入对象为空, requestId:{}", self.requestId)
            raise ServiceException(ExceptionType.AI_WRITE_OBJECT_EXCEPTION.value[0],
                                   ExceptionType.AI_WRITE_OBJECT_EXCEPTION.value[1])
        return self.ai_video_file