Merge pull request #829 from insight-platform/raw-output-gpu-converter

Adding a new converter to process raw output on GPU

samples/super_resolution/module.yml

@@ -32,7 +32,7 @@ pipeline:
           layer_names: [output]
           converter:
             module: savant.converter.raw_output
-            class_name: ModelRawOutputConverter
+            class_name: ModelCudaRawOutputConverter
           attributes:
             - name: ${parameters.sr_attribute}
               # just a way to save model output before place on frame, no need to output

samples/super_resolution/overlay.py

@@ -2,15 +2,16 @@
 
 from typing import Dict
 
+import cupy as cp
 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.memory_repr import cupy_array_as_opencv_gpu_mat
 
 SR_MODEL_NAME = param_storage()['sr_model']
 SR_ATTR_NAME = param_storage()['sr_attribute']
@@ -71,18 +72,20 @@ def process_frame(self, buffer: Gst.Buffer, frame_meta: NvDsFrameMeta):
             # 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)
+                sr_image_cp = 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)
+                sr_image_cp = (sr_image_cp * 255).astype(cp.uint8)
                 # CHW => HWC
-                sr_image_np = np.transpose(sr_image_np, (1, 2, 0))
+                sr_image_cp = cp.transpose(sr_image_cp, (1, 2, 0))
                 # RGB => RGBA
-                sr_image_np = np.dstack(
+                sr_image_cp = cp.dstack(
                     (
-                        sr_image_np,
-                        np.full(SUPER_RESOLUTION[::-1], 255, dtype=np.uint8),
+                        sr_image_cp,
+                        cp.full(SUPER_RESOLUTION[::-1], 255, dtype=cp.uint8),
                     )
                 )
+                sr_image_start_point = (sr_image_cp.shape[1], 0)
+                sr_image_mat = cupy_array_as_opencv_gpu_mat(sr_image_cp)
 
                 # Create frame for the auxiliary stream.
                 # The frame will be sent automatically
@@ -107,8 +110,8 @@ def process_frame(self, buffer: Gst.Buffer, frame_meta: NvDsFrameMeta):
                     )
                     opencv_utils.alpha_comp(
                         aux_mat,
-                        sr_image_np,
-                        (sr_image_np.shape[1], 0),
+                        sr_image_mat,
+                        sr_image_start_point,
                         stream=cuda_stream,
                     )
             else:

savant/converter/raw_output.py

@@ -1,13 +1,41 @@
-"""Model raw output converter."""
+"""Model raw output converters. They provide different options for processing output tensors: on the GPU or the host.
+See `savant.base.converter.TensorFormat`"""
 
 from typing import Any, List, Optional, Tuple
 
+import cupy as cp
 import numpy as np
 
-from savant.base.converter import BaseAttributeModelOutputConverter
+from savant.base.converter import BaseAttributeModelOutputConverter, TensorFormat
 from savant.base.model import AttributeModel
 
 
+class ModelCudaRawOutputConverter(BaseAttributeModelOutputConverter):
+    """Model raw output converter."""
+
+    tensor_format: TensorFormat = TensorFormat.CuPy
+
+    def __call__(
+        self,
+        *output_layers: cp.ndarray,
+        model: AttributeModel,
+        roi: Tuple[float, float, float, float],
+    ) -> List[Tuple[str, Any, Optional[float]]]:
+        """Returns raw model output tensors as attributes.
+
+        :param output_layers: Model output layer tensors
+        :param model: Attribute model
+        :param roi: ``[top, left, width, height]`` of the rectangle
+            on which the model infers
+        :return: list of attributes values with confidences
+            ``(attr_name, value, confidence)``
+        """
+        return [
+            (model.output.attributes[i].name, output, 1.0)
+            for i, output in enumerate(output_layers)
+        ]
+
+
 class ModelRawOutputConverter(BaseAttributeModelOutputConverter):
     """Model raw output converter."""