diff --git a/samples/opencv_cuda_bg_remover_mog2/bgremover.py b/samples/opencv_cuda_bg_remover_mog2/bgremover.py index ceb159a78..711ca175c 100644 --- a/samples/opencv_cuda_bg_remover_mog2/bgremover.py +++ b/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,20 @@ 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)) diff --git a/samples/opencv_cuda_bg_remover_mog2/demo.yml b/samples/opencv_cuda_bg_remover_mog2/demo.yml index 8281e2d1e..e669dcf3f 100644 --- a/samples/opencv_cuda_bg_remover_mog2/demo.yml +++ b/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 diff --git a/samples/opencv_cuda_bg_remover_mog2/docker-compose.l4t.yml b/samples/opencv_cuda_bg_remover_mog2/docker-compose.l4t.yml index ebde910f4..f3627a1cc 100644 --- a/samples/opencv_cuda_bg_remover_mog2/docker-compose.l4t.yml +++ b/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: diff --git a/samples/opencv_cuda_bg_remover_mog2/docker-compose.x86.yml b/samples/opencv_cuda_bg_remover_mog2/docker-compose.x86.yml index a7a153f1a..a91d2fa17 100644 --- a/samples/opencv_cuda_bg_remover_mog2/docker-compose.x86.yml +++ b/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: