Using auxiliary stream intead of paddings in a few samples

samples/opencv_cuda_bg_remover_mog2/bgremover.py

@@ -1,11 +1,15 @@
 """Background remover module."""
 
+from typing import Dict
+
 import cv2
 
+from savant.deepstream.auxiliary_stream import AuxiliaryStream
 from savant.deepstream.meta.frame import NvDsFrameMeta
 from savant.deepstream.opencv_utils import nvds_to_gpu_mat
 from savant.deepstream.pyfunc import NvDsPyFuncPlugin
 from savant.gstreamer import Gst
+from savant.parameter_storage import param_storage
 from savant.utils.artist import Artist
 
 
@@ -16,18 +20,65 @@ class BgRemover(NvDsPyFuncPlugin):
     MOG2 method from openCV is used to remove background.
     """
 
-    def __init__(self, **kwargs):
+    def __init__(
+        self,
+        codec_params: Dict,
+        **kwargs,
+    ):
         super().__init__(**kwargs)
+        self.codec_params = codec_params
+        self.result_aux_streams: Dict[str, AuxiliaryStream] = {}
         self.back_subtractors = {}
 
         self.gaussian_filter = cv2.cuda.createGaussianFilter(
             cv2.CV_8UC4, cv2.CV_8UC4, (9, 9), 2
         )
 
+    def on_source_add(self, source_id: str):
+        """Initialize an auxiliary stream for background removal result."""
+
+        self.logger.info('Source %s added.', source_id)
+        if source_id in self.result_aux_streams:
+            self.logger.info('Source %s already has a result stream.', source_id)
+            return
+
+        result_source_id = f'{source_id}-processed'
+        result_resolution = (
+            param_storage()['frame']['width'] * 2,
+            param_storage()['frame']['height'],
+        )
+        self.logger.info(
+            'Creating result auxiliary stream %s for source %s. Resolution: %s.',
+            result_source_id,
+            source_id,
+            result_resolution,
+        )
+        self.result_aux_streams[source_id] = self.auxiliary_stream(
+            source_id=result_source_id,
+            width=result_resolution[0],
+            height=result_resolution[1],
+            codec_params=self.codec_params,
+        )
+
+        if source_id in self.back_subtractors:
+            self.logger.info(
+                'Source %s already has a background subtractor.', source_id
+            )
+            return
+
+        self.logger.info('Creating background subtractor for source %s.', source_id)
+        self.back_subtractors[source_id] = cv2.cuda.createBackgroundSubtractorMOG2()
+
+    def on_stop(self) -> bool:
+        self.result_aux_streams = {}
+        return super().on_stop()
+
     def on_source_eos(self, source_id: str):
         """On source EOS event callback."""
         if source_id is self.back_subtractors:
             self.back_subtractors.pop(source_id)
+        if source_id in self.result_aux_streams:
+            self.result_aux_streams.get(source_id).eos()
 
     def process_frame(self, buffer: Gst.Buffer, frame_meta: NvDsFrameMeta):
         """Process frame metadata.
@@ -37,18 +88,22 @@ def process_frame(self, buffer: Gst.Buffer, frame_meta: NvDsFrameMeta):
         """
         stream = self.get_cuda_stream(frame_meta)
         with nvds_to_gpu_mat(buffer, frame_meta.frame_meta) as frame_mat:
-            with Artist(frame_mat, stream) as artist:
-                if frame_meta.source_id in self.back_subtractors:
+            result_stream = self.result_aux_streams[frame_meta.source_id]
+            # Created frame will be sent automatically
+            result_frame, result_buffer = result_stream.create_frame(
+                pts=frame_meta.pts,
+                duration=frame_meta.duration,
+            )
+            with nvds_to_gpu_mat(result_buffer, batch_id=0) as result_mat:
+                with Artist(result_mat, stream) as artist:
+                    frame_mat_copy = frame_mat.clone()
+
                     back_sub = self.back_subtractors[frame_meta.source_id]
-                else:
-                    back_sub = cv2.cuda.createBackgroundSubtractorMOG2()
-                    self.back_subtractors[frame_meta.source_id] = back_sub
-                ref_frame = cv2.cuda_GpuMat(
-                    frame_mat,
-                    (0, 0, int(frame_meta.roi.width), int(frame_meta.roi.height)),
-                )
-                cropped = ref_frame.clone()
-                self.gaussian_filter.apply(cropped, cropped, stream=stream)
-                cu_mat_fg = back_sub.apply(cropped, -1, stream)
-                res_image = ref_frame.copyTo(cu_mat_fg, stream)
-                artist.add_graphic(res_image, (int(frame_meta.roi.width), 0))
+                    self.gaussian_filter.apply(
+                        frame_mat_copy, frame_mat_copy, stream=stream
+                    )
+                    cu_mat_fg = back_sub.apply(frame_mat_copy, -1, stream)
+                    res_image = frame_mat_copy.copyTo(cu_mat_fg, stream)
+
+                    artist.add_graphic(frame_mat, (0, 0))
+                    artist.add_graphic(res_image, (int(frame_meta.roi.width), 0))

samples/opencv_cuda_bg_remover_mog2/demo.yml

@@ -7,19 +7,10 @@ parameters:
   frame:
     width: 1280
     height: 720
-    # Add paddings to the frame before processing
-    padding:
-      # Paddings are kept on the output frame
-      keep: true
-      left: 0
-      right: 1280
-      top: 0
-      bottom: 0
-  output_frame:
-    # Frame is output without any encoding
-    # this is to circumvent 3 hardware decoding processes limit on NVIDIA consumer hardware
-    codec: ${oc.env:CODEC, 'raw-rgba'}
   batch_size: 1
+  # to check auxiliary streams' encoder
+  auxiliary_encoders:
+    - ${pipeline.elements[0].kwargs.codec_params}
 
 # pipeline definition
 pipeline:
@@ -32,4 +23,9 @@ pipeline:
       module: samples.opencv_cuda_bg_remover_mog2.bgremover
       # specify the pyfunc's python class from the module
       class_name: BgRemover
+      kwargs:
+        # codec parameters for result stream (auxiliary stream)
+        codec_params:
+          codec: ${oc.env:CODEC, 'h264'}
+
   # sink definition is skipped, zeromq sink is used by default to connect with sink adapters

samples/opencv_cuda_bg_remover_mog2/docker-compose.l4t.yml

@@ -10,7 +10,7 @@ services:
       - LOCATION=https://eu-central-1.linodeobjects.com/savant-data/demo/road_traffic.mp4
       - DOWNLOAD_PATH=/tmp/video-loop-source-downloads
       - ZMQ_ENDPOINT=pub+connect:ipc:///tmp/zmq-sockets/input-video.ipc
-      - SOURCE_ID=road-traffic-processed
+      - SOURCE_ID=road-traffic
       - SYNC_OUTPUT=True
     entrypoint: /opt/savant/adapters/gst/sources/video_loop.sh
     depends_on:

samples/opencv_cuda_bg_remover_mog2/docker-compose.x86.yml

@@ -10,7 +10,7 @@ services:
       - LOCATION=https://eu-central-1.linodeobjects.com/savant-data/demo/road_traffic.mp4
       - DOWNLOAD_PATH=/tmp/video-loop-source-downloads
       - ZMQ_ENDPOINT=pub+connect:ipc:///tmp/zmq-sockets/input-video.ipc
-      - SOURCE_ID=road-traffic-processed
+      - SOURCE_ID=road-traffic
       - SYNC_OUTPUT=True
     entrypoint: /opt/savant/adapters/gst/sources/video_loop.sh
     depends_on:

samples/super_resolution/README.md

@@ -41,8 +41,8 @@ The demo uses models that are compiled into TensorRT engines the first time the
 
 docker compose -f samples/super_resolution/docker-compose.x86.yml up
 
-# open 'rtsp://127.0.0.1:554/stream/video' in your player
-# or visit 'http://127.0.0.1:888/stream/video/' (LL-HLS)
+# open 'rtsp://127.0.0.1:554/stream/video-super-resolution' in your player
+# or visit 'http://127.0.0.1:888/stream/video-super-resolution/' (LL-HLS)
 
 # Ctrl+C to stop running the compose bundle
 ```

samples/super_resolution/docker-compose.x86.yml

@@ -52,7 +52,7 @@ services:
       - ../assets/stub_imgs:/stub_imgs
     environment:
       - ZMQ_ENDPOINT=sub+connect:ipc:///tmp/zmq-sockets/output-video.ipc
-      - SOURCE_ID=video
+      - SOURCE_IDS=video-super-resolution
       - FRAMERATE=25/1
       - STUB_FILE_LOCATION=/stub_imgs/smpte100_3840x1080.jpeg
       - DEV_MODE=True

samples/super_resolution/module.yml

@@ -4,21 +4,12 @@ parameters:
   # set the super resolution model scale (x2/3/4) and name (ninasr_b0/1)
   sr_scale: 3
   sr_model: ninasr_b0
+  # needs to retrieve the model output in pyfunc element
+  sr_attribute: sr_frame
+  # pipeline processing frame parameters
   frame:
     width: 640
     height: 360
-    padding:
-      keep: true
-      left: 0
-      top: 0
-      # to output super resolution only
-#      right: ${calc:"arg_0*arg_1-arg_0", ${parameters.frame.width}, ${parameters.sr_scale}}
-#      bottom: ${calc:"arg_0*arg_1-arg_0", ${parameters.frame.height}, ${parameters.sr_scale}}
-      # to output scaled original + super resolution
-      right: ${calc:"arg_0*arg_1*2-arg_0", ${parameters.frame.width}, ${parameters.sr_scale}}
-      bottom: ${calc:"arg_0*arg_1-arg_0", ${parameters.frame.height}, ${parameters.sr_scale}}
-  output_frame:
-    codec: hevc
   batch_size: 1
 
 pipeline:
@@ -43,9 +34,13 @@ pipeline:
             module: savant.converter.raw_output
             class_name: ModelRawOutputConverter
           attributes:
-            - name: sr_frame
+            - name: ${parameters.sr_attribute}
               # just a way to save model output before place on frame, no need to output
               internal: true
     - element: pyfunc
       module: samples.super_resolution.overlay
       class_name: SROverlay
+      kwargs:
+        # codec parameters for result stream (auxiliary stream)
+        codec_params:
+          codec: hevc

samples/super_resolution/overlay.py

@@ -1,74 +1,123 @@
 """Overlay."""
 
+from typing import Dict
+
 import cv2
 import numpy as np
 
+from savant.deepstream import opencv_utils
 from savant.deepstream.meta.frame import NvDsFrameMeta
 from savant.deepstream.opencv_utils import nvds_to_gpu_mat
 from savant.deepstream.pyfunc import NvDsPyFuncPlugin
 from savant.gstreamer import Gst
 from savant.parameter_storage import param_storage
-from savant.utils.artist import Artist
 
 SR_MODEL_NAME = param_storage()['sr_model']
-SR_ATTR_NAME = 'sr_frame'
-INPUT_RESOLUTION = (
-    param_storage()['frame']['width'],
-    param_storage()['frame']['height'],
-)
+SR_ATTR_NAME = param_storage()['sr_attribute']
 SUPER_RESOLUTION = (
     param_storage()['frame']['width'] * param_storage()['sr_scale'],
     param_storage()['frame']['height'] * param_storage()['sr_scale'],
 )
 
 
 class SROverlay(NvDsPyFuncPlugin):
-    """Super resolution overlay pyfunc."""
+    def __init__(
+        self,
+        codec_params: Dict,
+        **kwargs,
+    ):
+        self.codec_params = codec_params
+        self.result_aux_stream = None
+
+        super().__init__(**kwargs)
+
+    def on_source_add(self, source_id: str):
+        """Initialize an auxiliary stream for super resolution result."""
+
+        self.logger.info('Source %s added.', source_id)
+
+        result_source_id = f'{source_id}-super-resolution'
+        result_resolution = (SUPER_RESOLUTION[0] * 2, SUPER_RESOLUTION[1])
+        self.logger.info(
+            'Creating result auxiliary stream %s for source %s. Resolution: %s.',
+            result_source_id,
+            source_id,
+            result_resolution,
+        )
+        self.result_aux_stream = self.auxiliary_stream(
+            source_id=result_source_id,
+            width=result_resolution[0],
+            height=result_resolution[1],
+            codec_params=self.codec_params,
+        )
 
     def process_frame(self, buffer: Gst.Buffer, frame_meta: NvDsFrameMeta):
         """Process frame metadata.
 
         :param buffer: Gstreamer buffer with this frame's data.
         :param frame_meta: This frame's metadata.
         """
-        cuda_stream = self.get_cuda_stream(frame_meta)
-        with nvds_to_gpu_mat(buffer, frame_meta.frame_meta) as frame_mat, Artist(
-            frame_mat, cuda_stream
-        ) as artist:
-            # TODO: original + super resolution mix
-            sr_lt = (0, 0)  # super resolution left, top
-
-            # place origin, then super resolution
-            if frame_mat.size()[0] > SUPER_RESOLUTION[0]:
-                # scale original image and place first
-                source_image = cv2.cuda_GpuMat(
-                    frame_mat,
-                    (0, 0, INPUT_RESOLUTION[0], INPUT_RESOLUTION[1]),
-                )
-                scaled_image = cv2.cuda.resize(
-                    src=source_image,
-                    dsize=SUPER_RESOLUTION,
-                    # interpolation=cv2.INTER_LINEAR,
-                    stream=cuda_stream,
-                )
-                artist.add_graphic(scaled_image, (0, 0))
-                sr_lt = (scaled_image.size()[0], 0)
 
-            # check super resolution attr
+        # Get CUDA stream for asynchronous processing
+        cuda_stream = self.get_cuda_stream(frame_meta)
+        with nvds_to_gpu_mat(buffer, frame_meta.frame_meta) as frame_mat:
+            # Check super resolution attr
             sr_attr = None
             for obj_meta in frame_meta.objects:
                 if obj_meta.is_primary:
                     sr_attr = obj_meta.get_attr_meta(SR_MODEL_NAME, SR_ATTR_NAME)
                     break
 
-            # transform super resolution and place on the frame
+            # Transform super resolution image
             if sr_attr:
+                # Normalize array values to be within the range [0.0, 1.0]
                 sr_image_np = sr_attr.value.clip(0.0, 1.0)
+                # Convert the normalized array to 8-bit unsigned integer format
                 sr_image_np = (sr_image_np * 255).astype(np.uint8)
-                # chw => hwc
+                # CHW => HWC
                 sr_image_np = np.transpose(sr_image_np, (1, 2, 0))
-                # rgb => rgba
+                # RGB => RGBA
                 sr_image_np = np.dstack(
-                    (sr_image_np, np.full(SUPER_RESOLUTION[::-1], 255, dtype=np.uint8))
+                    (
+                        sr_image_np,
+                        np.full(SUPER_RESOLUTION[::-1], 255, dtype=np.uint8),
+                    )
+                )
+
+                # Create frame for the auxiliary stream.
+                # The frame will be sent automatically
+                aux_frame, aux_buffer = self.result_aux_stream.create_frame(
+                    pts=frame_meta.pts,
+                    duration=frame_meta.duration,
                 )
-                artist.add_graphic(sr_image_np, sr_lt)
+                with nvds_to_gpu_mat(aux_buffer, batch_id=0) as aux_mat:
+                    # Scale the image to display it alongside the super resolution result.
+                    scaled_image = cv2.cuda.resize(
+                        src=frame_mat,
+                        dsize=SUPER_RESOLUTION,
+                        stream=cuda_stream,
+                    )
+
+                    # Place original frame and super resolution frame side by side
+                    opencv_utils.alpha_comp(
+                        aux_mat,
+                        scaled_image,
+                        (0, 0),
+                        stream=cuda_stream,
+                    )
+                    opencv_utils.alpha_comp(
+                        aux_mat,
+                        sr_image_np,
+                        (sr_image_np.shape[1], 0),
+                        stream=cuda_stream,
+                    )
+            else:
+                self.logger.warning('Super resolution attribute not found.')
+
+    def on_stop(self) -> bool:
+        self.result_aux_stream = None
+        return super().on_stop()
+
+    def on_source_eos(self, source_id: str):
+        self.logger.info('Got EOS from source %s.', source_id)
+        self.result_aux_stream.eos()