diff --git a/samples/opencv_cuda_bg_remover_mog2/bgremover.py b/samples/opencv_cuda_bg_remover_mog2/bgremover.py index ceb159a7..c5992b78 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,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)) diff --git a/samples/opencv_cuda_bg_remover_mog2/demo.yml b/samples/opencv_cuda_bg_remover_mog2/demo.yml index 8281e2d1..e669dcf3 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 ebde910f..f3627a1c 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 a7a153f1..a91d2fa1 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: diff --git a/samples/super_resolution/README.md b/samples/super_resolution/README.md index 386ca075..9a6d20d8 100644 --- a/samples/super_resolution/README.md +++ b/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 ``` diff --git a/samples/super_resolution/docker-compose.x86.yml b/samples/super_resolution/docker-compose.x86.yml index 1a803b9c..ebea0e6f 100644 --- a/samples/super_resolution/docker-compose.x86.yml +++ b/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 diff --git a/samples/super_resolution/module.yml b/samples/super_resolution/module.yml index f4da74ed..98d26b3c 100644 --- a/samples/super_resolution/module.yml +++ b/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 diff --git a/samples/super_resolution/overlay.py b/samples/super_resolution/overlay.py index 3fe81c1c..e108ff08 100644 --- a/samples/super_resolution/overlay.py +++ b/samples/super_resolution/overlay.py @@ -1,21 +1,19 @@ """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'], @@ -23,7 +21,35 @@ 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. @@ -31,44 +57,67 @@ def process_frame(self, buffer: Gst.Buffer, frame_meta: NvDsFrameMeta): :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()