Merge pull request #936 from luxonis/sdk/on_host_stereo_colorization

On-device stereo colorization

Author: daniilpastukhov

depthai_sdk/examples/StereoComponent/stereo.py

@@ -5,13 +5,10 @@
 from depthai_sdk.visualize.configs import StereoColor
 
 with OakCamera() as oak:
-    stereo = oak.create_stereo('800p', fps=30)
-    stereo.config_postprocessing(
-        colorize=StereoColor.GRAY,
-        colormap=cv2.COLORMAP_BONE,
-        wls_filter=True,
-        wls_level=WLSLevel.HIGH
-    )
+    stereo = oak.create_stereo('800p', fps=30, encode='h264')
+
+    stereo.config_postprocessing(colorize=StereoColor.RGBD, colormap=cv2.COLORMAP_MAGMA)
+    stereo.config_wls(wls_level=WLSLevel.MEDIUM)
 
     oak.visualize(stereo.out.disparity)
     oak.start(blocking=True)

depthai_sdk/setup.py

@@ -9,7 +9,7 @@
 
 setup(
     name='depthai-sdk',
-    version='1.9.3.1',
+    version='1.9.4',
     description='This package provides an abstraction of the DepthAI API library.',
     long_description=io.open("README.md", encoding="utf-8").read(),
     long_description_content_type="text/markdown",

depthai_sdk/src/depthai_sdk/components/stereo_component.py

@@ -1,3 +1,4 @@
+import warnings
 from enum import Enum
 from typing import Optional, Union, Any, Dict
 
@@ -6,11 +7,11 @@
 from depthai_sdk.components.camera_component import CameraComponent
 from depthai_sdk.components.component import Component
 from depthai_sdk.components.parser import parse_cam_socket, parse_median_filter, parse_encode
-from depthai_sdk.oak_outputs.xout.xout_disparity import XoutDisparity
+from depthai_sdk.oak_outputs.xout.xout_base import XoutBase, StreamXout
 from depthai_sdk.oak_outputs.xout.xout_depth import XoutDepth
+from depthai_sdk.oak_outputs.xout.xout_disparity import XoutDisparity
 from depthai_sdk.oak_outputs.xout.xout_h26x import XoutH26x
 from depthai_sdk.oak_outputs.xout.xout_mjpeg import XoutMjpeg
-from depthai_sdk.oak_outputs.xout.xout_base import XoutBase, StreamXout
 from depthai_sdk.replay import Replay
 from depthai_sdk.visualize.configs import StereoColor
 
@@ -32,8 +33,7 @@ def __init__(self,
                  replay: Optional[Replay] = None,
                  args: Any = None,
                  name: Optional[str] = None,
-                 encode: Union[None, str, bool, dai.VideoEncoderProperties.Profile] = None
-                 ):
+                 encode: Union[None, str, bool, dai.VideoEncoderProperties.Profile] = None):
         """
         Args:
             pipeline (dai.Pipeline): DepthAI pipeline
@@ -44,15 +44,18 @@ def __init__(self,
             replay (Replay object, optional): Replay
             args (Any, optional): Use user defined arguments when constructing the pipeline
             name (str, optional): Name of the output stream
+            encode (str/bool/Profile, optional): Encode the output stream
         """
         super().__init__()
         self.out = self.Out(self)
 
-        left: Union[None, CameraComponent, dai.node.MonoCamera, dai.node.ColorCamera, dai.Node.Output]
-        right: Union[None, CameraComponent, dai.node.MonoCamera, dai.node.ColorCamera, dai.Node.Output]
+        self.left: Union[None, CameraComponent, dai.node.MonoCamera, dai.node.ColorCamera, dai.Node.Output]
+        self.right: Union[None, CameraComponent, dai.node.MonoCamera, dai.node.ColorCamera, dai.Node.Output]
 
-        _left_stream: dai.Node.Output
-        _right_stream: dai.Node.Output
+        self._left_stream: dai.Node.Output
+        self._right_stream: dai.Node.Output
+
+        self.colormap = None
 
         self._replay: Optional[Replay] = replay
         self._resolution: Optional[Union[str, dai.MonoCameraProperties.SensorResolution]] = resolution
@@ -74,8 +77,8 @@ def __init__(self,
 
         # Configuration variables
         self._colorize = None
-        self._colormap = None
-        self._use_wls_filter = None
+        self._postprocess_colormap = None
+        self.wls_enabled = None
         self._wls_level = None
         self._wls_lambda = None
         self._wls_sigma = None
@@ -129,14 +132,25 @@ def on_init(self, pipeline: dai.Pipeline, device: dai.Device, version: dai.OpenV
             self._left_stream.link(self.node.left)
             self._right_stream.link(self.node.right)
 
+        colormap_manip = None
+        if self.colormap:
+            colormap_manip = pipeline.create(dai.node.ImageManip)
+            colormap_manip.initialConfig.setColormap(self.colormap, self.node.initialConfig.getMaxDisparity())
+            colormap_manip.initialConfig.setFrameType(dai.ImgFrame.Type.NV12)
+            colormap_manip.setMaxOutputFrameSize(1200 * 800 * 3)
+            self.node.disparity.link(colormap_manip.inputImage)
+
         if self.encoder:
             try:
                 fps = self.left.get_fps()  # CameraComponent
             except AttributeError:
                 fps = self.left.getFps()  # MonoCamera
 
             self.encoder.setDefaultProfilePreset(fps, self._encoderProfile)
-            self.node.disparity.link(self.encoder.input)
+            if colormap_manip:
+                colormap_manip.out.link(self.encoder.input)
+            else:
+                self.node.disparity.link(self.encoder.input)
 
         self.node.setOutputSize(1200, 800)
 
@@ -181,22 +195,14 @@ def config_stereo(self,
 
     def config_postprocessing(self,
                               colorize: Union[StereoColor, bool] = None,
-                              colormap: int = None,
-                              wls_filter: bool = None,
-                              wls_level: WLSLevel = None,
-                              wls_lambda: float = None,
-                              wls_sigma: float = None
+                              colormap: int = None
                               ) -> None:
         """
         Configures postprocessing options.
 
         Args:
             colorize: Colorize the disparity map. Can be either a StereoColor enum, string or bool.
             colormap: Colormap to use for colorizing the disparity map.
-            wls_filter: Enable WLS filter. If enabled, the output will be filtered using WLS filter.
-            wls_level: WLS filter level. Can be either a WLSLevel enum or string.
-            wls_lambda: WLS filter lambda.
-            wls_sigma: WLS filter sigma.
         """
         if colorize is None:
             self._colorize = StereoColor.GRAY
@@ -207,8 +213,22 @@ def config_postprocessing(self,
         elif isinstance(colorize, str):
             self._colorize = StereoColor[colorize.upper()]
 
-        self._colormap = colormap
-        self._use_wls_filter = wls_filter
+        self._postprocess_colormap = colormap
+
+    def config_wls(self,
+                   wls_level: Union[WLSLevel, str] = None,
+                   wls_lambda: float = None,
+                   wls_sigma: float = None
+                   ) -> None:
+        """
+        Configures WLS filter options.
+
+        Args:
+            wls_level: WLS filter level. Can be either a WLSLevel enum or string.
+            wls_lambda: WLS filter lambda.
+            wls_sigma: WLS filter sigma.
+        """
+        self.wls_enabled = True if wls_level else False
 
         if isinstance(wls_level, WLSLevel):
             self._wls_level = wls_level
@@ -218,6 +238,15 @@ def config_postprocessing(self,
         self._wls_lambda = wls_lambda
         self._wls_sigma = wls_sigma
 
+    def set_colormap(self, colormap: dai.Colormap):
+        """
+        Sets the colormap to use for colorizing the disparity map. Used for on-device postprocessing.
+
+        Args:
+            colormap: Colormap to use for colorizing the disparity map.
+        """
+        self.colormap = colormap
+
     def _get_disparity_factor(self, device: dai.Device) -> float:
         """
         Calculates the disparity factor used to calculate depth from disparity.
@@ -254,6 +283,10 @@ def main(self, pipeline: dai.Pipeline, device: dai.Device) -> XoutBase:
             return self.depth(pipeline, device)
 
         def disparity(self, pipeline: dai.Pipeline, device: dai.Device) -> XoutBase:
+            if self._comp.colormap:
+                warnings.warn('Colormap set with `set_colormap` is ignored when using disparity output. '
+                              'Please use `configure_postprocessing` instead.')
+
             fps = self._comp.left.get_fps() if self._comp._replay is None else self._comp._replay.getFps()
 
             out = XoutDisparity(
@@ -262,8 +295,8 @@ def disparity(self, pipeline: dai.Pipeline, device: dai.Device) -> XoutBase:
                 max_disp=self._comp.node.getMaxDisparity(),
                 fps=fps,
                 colorize=self._comp._colorize,
-                colormap=self._comp._colormap,
-                use_wls_filter=self._comp._use_wls_filter,
+                colormap=self._comp._postprocess_colormap,
+                use_wls_filter=self._comp.wls_enabled,
                 wls_level=self._comp._wls_level,
                 wls_lambda=self._comp._wls_lambda,
                 wls_sigma=self._comp._wls_sigma
@@ -272,15 +305,19 @@ def disparity(self, pipeline: dai.Pipeline, device: dai.Device) -> XoutBase:
             return self._comp._create_xout(pipeline, out)
 
         def depth(self, pipeline: dai.Pipeline, device: dai.Device) -> XoutBase:
+            if self._comp.colormap:
+                warnings.warn('Colormap set with `set_colormap` is ignored when using depth output. '
+                              'Please use `configure_postprocessing` instead.')
+
             fps = self._comp.left.get_fps() if self._comp._replay is None else self._comp._replay.getFps()
             out = XoutDepth(
                 device=device,
                 frames=StreamXout(self._comp.node.id, self._comp.depth, name=self._comp.name),
                 mono_frames=StreamXout(self._comp.node.id, self._comp._right_stream, name=self._comp.name),
                 fps=fps,
                 colorize=self._comp._colorize,
-                colormap=self._comp._colormap,
-                use_wls_filter=self._comp._use_wls_filter,
+                colormap=self._comp._postprocess_colormap,
+                use_wls_filter=self._comp.wls_enabled,
                 wls_level=self._comp._wls_level,
                 wls_lambda=self._comp._wls_lambda,
                 wls_sigma=self._comp._wls_sigma
@@ -291,15 +328,18 @@ def encoded(self, pipeline: dai.Pipeline, device: dai.Device) -> XoutBase:
             if not self._comp.encoder:
                 raise RuntimeError('Encoder not enabled, cannot output encoded frames')
 
+            if self._comp.wls_enabled:
+                warnings.warn('WLS filter is enabled, but cannot be applied to encoded frames.')
+
             if self._comp._encoderProfile == dai.VideoEncoderProperties.Profile.MJPEG:
                 out = XoutMjpeg(frames=StreamXout(self._comp.encoder.id, self._comp.encoder.bitstream),
-                                color=False,
+                                color=self._comp.colormap is not None,
                                 lossless=self._comp.encoder.getLossless(),
                                 fps=self._comp.encoder.getFrameRate(),
                                 frame_shape=(1200, 800))
             else:
                 out = XoutH26x(frames=StreamXout(self._comp.encoder.id, self._comp.encoder.bitstream),
-                               color=False,
+                               color=self._comp.colormap is not None,
                                profile=self._comp._encoderProfile,
                                fps=self._comp.encoder.getFrameRate(),
                                frame_shape=(1200, 800))