fix a bug in rtsp stream

ffmpegcv/__init__.py

@@ -6,6 +6,7 @@
 from .ffmpeg_writer_stream_realtime import FFmpegWriterStreamRT
 from .ffmpeg_reader_qsv import FFmpegReaderQSV
 from .ffmpeg_writer_qsv import FFmpegWriterQSV
+from .ffmpeg_reader_pannels import FFmpegReaderPannels
 from .ffmpeg_noblock import noblock, ReadLiveLast
 from .video_info import get_num_NVIDIA_GPUs
 import shutil
@@ -462,3 +463,58 @@ def VideoCaptureStreamRT(
 
 
 VideoReaderStreamRT = VideoCaptureStreamRT
+
+
+def VideoCapturePannels(
+    file:str,
+    crop_xywh_l:list,
+    codec=None,
+    pix_fmt="bgr24",
+    resize=None
+):
+    """
+    Alternative to cv2.VideoCapture
+
+    Parameters
+    ----------
+    file : str
+        Path to video file.
+    crop_xywh_l : list of crop_xywh
+        Crop the frame. [(x0, y0, w0, h0), (x1, y1, w1, h1), ...].
+    codec : str
+        Codec to use. Optional. Default is `None`.
+    pix_fmt : str
+        Pixel format. ['bgr24' | 'rgb24' | 'gray']. Optional. Default is 'bgr24'.
+    resize  : tuple as (w,h)
+        Resize the pannels to identical size. Optional. Default is `None`.
+        Does not keep the ratio.
+
+
+    Examples
+    --------
+    
+
+    ffmpegcv
+    ```
+    w,h = 1280,800
+    cap = ffmpegcv.VideoCapturePannels(file,
+        [[0,0,w,h], [w,0,w,h],[0,h,w,h], [w,h,w,h]])
+    while True:
+        ret, frame_pannels = cap.read()
+        if not ret:
+            break
+        print(frame_pannels.shape)  # shape=(4,h,w,3)
+    ```
+
+    Author: Chenxinfeng 2024-02-15, [email protected]
+    """
+
+    return FFmpegReaderPannels.VideoReader(
+        file, 
+        crop_xywh_l,
+        codec, 
+        pix_fmt,
+        resize,
+    )
+
+VideoReaderPannels = VideoCapturePannels

ffmpegcv/ffmpeg_reader_pannels.py

@@ -0,0 +1,87 @@
+import os
+import numpy as np
+from ffmpegcv.ffmpeg_reader import FFmpegReader, get_outnumpyshape
+
+from ffmpegcv.video_info import (
+    get_info,
+    run_async
+)
+
+
+class FFmpegReaderPannels(FFmpegReader):
+    @staticmethod
+    def VideoReader(
+        filename:str,
+        crop_xywh_l:list,
+        codec,
+        pix_fmt='bgr24',
+        resize=None
+    ):
+        assert os.path.exists(filename) and os.path.isfile(
+            filename
+        ), f"{filename} not exists"
+        assert pix_fmt in ["rgb24", "bgr24", "yuv420p", "nv12", "gray"]
+        vid = FFmpegReaderPannels()
+        crop_xywh_l = np.array(crop_xywh_l)
+        vid.crop_xywh_l = crop_xywh_l
+        videoinfo = get_info(filename)
+        vid.origin_width = videoinfo.width
+        vid.origin_height = videoinfo.height
+        vid.fps = videoinfo.fps
+        vid.count = videoinfo.count
+        vid.duration = videoinfo.duration
+        vid.pix_fmt = pix_fmt
+        vid.codec = codec if codec else videoinfo.codec
+
+        vid.crop_width_l = crop_xywh_l[:,2]
+        vid.crop_height_l = crop_xywh_l[:,3]
+        vid.size_l = crop_xywh_l[:,2:][:,::-1]
+        vid.npannel = len(crop_xywh_l)
+        vid.out_numpy_shape_l = [get_outnumpyshape(s[::-1], pix_fmt) for s in vid.size_l]
+        if len(set(vid.crop_width_l)) == len(set(vid.crop_height_l)) == 1:
+            vid.is_pannel_similar = True
+            vid.crop_width = vid.crop_width_l[0]
+            vid.crop_height = vid.crop_height_l[0]
+            vid.size = vid.size_l[0]
+            vid.out_numpy_shape = (vid.npannel, *vid.out_numpy_shape_l[0])
+        else:
+            vid.is_pannel_similar = False
+            vid.crop_width = vid.crop_height = vid.size = None
+            vid.out_numpy_shape = (np.sum(np.prod(s) for s in vid.out_numpy_shape_l),)
+
+        VINSRCs =''.join(f'[VSRC{i}]' for i in range(vid.npannel))
+        pix_fmtopt = ',extractplanes=y' if pix_fmt=='gray' else ''
+        CROPs = ';'.join(f'[VSRC{i}]crop={w}:{h}:{x}:{y}{pix_fmtopt}[VPANEL{i}]'
+                          for i, (x,y,w,h) in enumerate(vid.crop_xywh_l))
+        filteropt = f' -filter_complex "split={vid.npannel}{VINSRCs};{CROPs}"'
+        outmaps = ''.join(f' -map [VPANEL{i}] -pix_fmt {pix_fmt} -r {vid.fps} -f rawvideo pipe:'
+                          for i in range(vid.npannel))
+
+        vid.ffmpeg_cmd = (
+            f"ffmpeg -loglevel warning "
+            f' -r {vid.fps} -i "{filename}" '
+            f" {filteropt} {outmaps}"
+        )
+        return vid
+
+    def read(self):
+        if self.waitInit:
+            self.process = run_async(self.ffmpeg_cmd)
+            self.waitInit = False
+
+        in_bytes = self.process.stdout.read(np.prod(self.out_numpy_shape))
+        if not in_bytes:
+            self.release()
+            return False, None
+        self.iframe += 1
+        img0 = np.frombuffer(in_bytes, np.uint8)
+        if self.is_pannel_similar:
+            img = img0.reshape(self.out_numpy_shape)
+        else:
+            img = []
+            for out_numpy_shape in self.out_numpy_shape_l:
+                nbuff = np.prod(out_numpy_shape)
+                img.append(img0[:nbuff].reshape(out_numpy_shape))
+                img0 = img0[nbuff:]
+
+        return True, img

ffmpegcv/ffmpeg_reader_stream.py

@@ -35,8 +35,10 @@ def VideoReader(
                 resize_keepratio, resize_keepratioalign)
         vid.size = (vid.width, vid.height)
 
+        rtsp_opt = '-rtsp_transport tcp ' if stream_url.startswith('rtsp://') else ''
         vid.ffmpeg_cmd = (
             f"ffmpeg -loglevel warning "
+            f' {rtsp_opt} '
             f' -vcodec {vid.codec} -i {stream_url} '
             f" {filteropt} -pix_fmt {pix_fmt}  -f rawvideo pipe:"
         )

ffmpegcv/stream_info.py

@@ -5,7 +5,8 @@
 
 
 def get_info(stream_url):
-    cmd = 'ffprobe -v quiet -print_format xml -select_streams v:0 -show_format -show_streams "{}"'.format(stream_url)
+    rtspflag = '-rtsp_transport tcp' if stream_url.startswith('rtsp://') else {}
+    cmd = 'ffprobe -v quiet -print_format xml {} -select_streams v:0 -show_format -show_streams "{}"'.format(rtspflag, stream_url)
     output = subprocess.check_output(shlex.split(cmd), shell=False)
     root = ET.fromstring(output)
     assert (root[0].tag, root[0][0].tag) == ("streams", "stream")

ffmpegcv/version.py

@@ -1 +1 @@
-__version__='0.3.9b'
+__version__='0.3.10'

setup.py

@@ -6,7 +6,7 @@
 
 setup(
     name='ffmpegcv', # 应用名
-    version='0.3.9b', # 版本号
+    version='0.3.10', # 版本号
     packages=find_packages(include=['ffmpegcv*']), # 包括在安装包内的 Python 包
     author='chenxf',
     author_email='[email protected]',