115 lines
5.0 KiB
Python
115 lines
5.0 KiB
Python
|
import subprocess as sp
|
|||
|
|
from traceback import format_exc
|
|||
|
|
import cv2, time
|
|||
|
|
import numpy as np
|
|||
|
|
from loguru import logger
|
|||
|
|
from DrGraph.utils.Helper import Helper
|
|||
|
|
from DrGraph.utils.Exception import ServiceException
|
|||
|
|
from DrGraph.utils.Constant import Constant
|
|||
|
|
|
|||
|
|
class NetStream:
|
|||
|
|
def __init__(self, pull_url, push_url, request_id):
|
|||
|
|
self.pull_url = pull_url
|
|||
|
|
self.push_url = push_url
|
|||
|
|
self.request_id = request_id
|
|||
|
|
self.pull_p = None
|
|||
|
|
|
|||
|
|
self.width = 1920
|
|||
|
|
self.height = 1080 * 3 // 2
|
|||
|
|
self.width_height_3 = 1920 * 1080 * 3 // 2
|
|||
|
|
self.w_2 = 960
|
|||
|
|
self.h_2 = 540
|
|||
|
|
|
|||
|
|
self.frame_count = 0
|
|||
|
|
self.start_time = time.time();
|
|||
|
|
|
|||
|
|
def clear_pull_p(self,pull_p, requestId):
|
|||
|
|
try:
|
|||
|
|
if pull_p and pull_p.poll() is None:
|
|||
|
|
logger.info("关闭拉流管道, requestId:{}", requestId)
|
|||
|
|
if pull_p.stdout:
|
|||
|
|
pull_p.stdout.close()
|
|||
|
|
pull_p.terminate()
|
|||
|
|
pull_p.wait(timeout=30)
|
|||
|
|
logger.info("拉流管道已关闭, requestId:{}", requestId)
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error("关闭拉流管道异常: {}, requestId:{}", format_exc(), requestId)
|
|||
|
|
if pull_p and pull_p.poll() is None:
|
|||
|
|
pull_p.kill()
|
|||
|
|
pull_p.wait(timeout=30)
|
|||
|
|
raise e
|
|||
|
|
|
|||
|
|
def start_pull_p(self, pull_url, requestId):
|
|||
|
|
try:
|
|||
|
|
command = ['D:/DrGraph/DSP/ffmpeg.exe']
|
|||
|
|
# if pull_url.startswith("rtsp://"):
|
|||
|
|
# command.extend(['-timeout', '20000000', '-rtsp_transport', 'tcp'])
|
|||
|
|
# if pull_url.startswith("http") or pull_url.startswith("rtmp"):
|
|||
|
|
# command.extend(['-rw_timeout', '20000000'])
|
|||
|
|
command.extend(['-re',
|
|||
|
|
'-y',
|
|||
|
|
'-an',
|
|||
|
|
# '-hwaccel', 'cuda', cuvid
|
|||
|
|
'-c:v', 'h264_cuvid',
|
|||
|
|
# '-resize', self.wah,
|
|||
|
|
'-i', pull_url,
|
|||
|
|
'-f', 'rawvideo',
|
|||
|
|
# '-pix_fmt', 'bgr24',
|
|||
|
|
'-r', '25',
|
|||
|
|
'-'])
|
|||
|
|
self.pull_p = sp.Popen(command, stdout=sp.PIPE)
|
|||
|
|
return self.pull_p
|
|||
|
|
except ServiceException as s:
|
|||
|
|
logger.error("构建拉流管道ServiceException异常: url={}, {}, requestId:{}", pull_url, s.msg, requestId)
|
|||
|
|
raise s
|
|||
|
|
except Exception as e:
|
|||
|
|
logger.error("构建拉流管道Exception异常：url={}, {}, requestId:{}", pull_url, format_exc(), requestId)
|
|||
|
|
raise e
|
|||
|
|
|
|||
|
|
def pull_read_video_stream(self):
|
|||
|
|
result = None
|
|||
|
|
try:
|
|||
|
|
if self.pull_p is None:
|
|||
|
|
self.start_pull_p(self.pull_url, self.request_id)
|
|||
|
|
in_bytes = self.pull_p.stdout.read(self.width_height_3)
|
|||
|
|
if in_bytes is not None and len(in_bytes) > 0:
|
|||
|
|
try:
|
|||
|
|
# result = (np.frombuffer(in_bytes, np.uint8).reshape([height * 3 // 2, width, 3]))
|
|||
|
|
# ValueError: cannot reshape array of size 3110400 into shape (1080,1920)
|
|||
|
|
result = (np.frombuffer(in_bytes, np.uint8)).reshape((self.height, self.width))
|
|||
|
|
result = cv2.cvtColor(result, cv2.COLOR_YUV2BGR_NV12)
|
|||
|
|
# result = cv2.cvtColor(result, cv2.COLOR_RGB2BGR)
|
|||
|
|
if result.shape[1] > Constant.pull_frame_width:
|
|||
|
|
result = cv2.resize(result, (result.shape[1] // 2, result.shape[0] // 2), interpolation=cv2.INTER_LINEAR)
|
|||
|
|
except Exception:
|
|||
|
|
logger.error("视频格式异常：{}, requestId:{}", format_exc(), requestId)
|
|||
|
|
raise ServiceException(ExceptionType.VIDEO_RESOLUTION_EXCEPTION.value[0],
|
|||
|
|
ExceptionType.VIDEO_RESOLUTION_EXCEPTION.value[1])
|
|||
|
|
except ServiceException as s:
|
|||
|
|
logger.error("ServiceException 读流异常: {}, requestId:{}", s.msg, self.request_id)
|
|||
|
|
self.clear_pull_p(self.pull_p, self.request_id)
|
|||
|
|
self.pull_p = None
|
|||
|
|
result = None
|
|||
|
|
raise s
|
|||
|
|
except Exception:
|
|||
|
|
logger.error("Exception 读流异常：{}, requestId:{}", format_exc(), self.request_id)
|
|||
|
|
self.clear_pull_p(self.pull_p, self.request_id)
|
|||
|
|
self.pull_p = None
|
|||
|
|
self.width = None
|
|||
|
|
self.height = None
|
|||
|
|
self.width_height_3 = None
|
|||
|
|
result = None
|
|||
|
|
logger.error("读流异常：{}, requestId:{}", format_exc(), self.request_id)
|
|||
|
|
return result
|
|||
|
|
|
|||
|
|
def prepare_pull(self):
|
|||
|
|
if self.pull_p is None:
|
|||
|
|
self.start_time = time.time();
|
|||
|
|
self.start_pull_p(self.pull_url, self.request_id)
|
|||
|
|
|
|||
|
|
def next_pull_frame(self):
|
|||
|
|
if self.pull_p is None:
|
|||
|
|
logger.error(f'pull_p is None, requestId: {self.request_id}')
|
|||
|
|
return None
|
|||
|
|
frame = self.pull_read_video_stream()
|
|||
|
|
return frame
|