Add depth score

depthai_sdk/src/depthai_sdk/classes/packets.py

@@ -39,6 +39,7 @@ class _TrackingDetection(_Detection):
 class _TwoStageDetection(_Detection):
     nn_data: dai.NNData
 
+
 class NNDataPacket:
     """
     Contains only dai.NNData message
@@ -50,6 +51,7 @@ def __init__(self, name: str, nn_data: dai.NNData):
         self.name = name
         self.msg = nn_data
 
+
 class FramePacket:
     """
     Contains only dai.ImgFrame message and cv2 frame, which is used by visualization logic.
@@ -80,6 +82,7 @@ def __init__(self,
         self.color_frame = color_frame
         self.visualizer = visualizer
 
+
 class DepthPacket(FramePacket):
     mono_frame: dai.ImgFrame
 
@@ -88,6 +91,7 @@ def __init__(self,
                  img_frame: dai.ImgFrame,
                  mono_frame: Optional[dai.ImgFrame],
                  depth_map: Optional[np.ndarray] = None,
+                 confidence_map: Optional[np.ndarray] = None,
                  visualizer: 'Visualizer' = None):
         super().__init__(name=name,
                          msg=img_frame,
@@ -98,6 +102,13 @@ def __init__(self,
             self.mono_frame = mono_frame
 
         self.depth_map = depth_map
+        self.confidence_map = confidence_map
+        self.depth_score = None
+        if self.confidence_map:
+            values = 1 - (self.confidence_map.getData() / 255)
+            values_no_outliers = values[np.logical_and(values > 0.0, values < 1.0)]
+            self.depth_score = np.sqrt(np.mean(values_no_outliers))
+
 
 class SpatialBbMappingPacket(FramePacket):
     """

depthai_sdk/src/depthai_sdk/components/stereo_component.py

@@ -113,6 +113,10 @@ def __init__(self,
             'sigma': None
         }
 
+        # Output config
+        self.enable_depth_score = False
+        self.validate_calibration = False
+
         self._undistortion_offset: Optional[int] = None
 
         if not self._replay:
@@ -307,6 +311,23 @@ def config_wls(self,
             'sigma': wls_sigma,
         }
 
+    def config_output(self,
+                      depth_score: bool = None,
+                      validate_calibration: bool = None
+                      ) -> None:
+        """
+        Configures output streams.
+
+        Args:
+            depth_score: True to include depth score in the output packets.
+            validate_calibration: Check if the calibration is valid during the runtime (done on-host) and warn
+            the user if it's not. This can be used to detect if the calibration is invalid (e.g. due to temperature drift).
+        """
+        if depth_score is not None:
+            self.enable_depth_score = depth_score
+        if validate_calibration is not None:
+            self.validate_calibration = validate_calibration
+
     def set_colormap(self, colormap: dai.Colormap):
         """
         Sets the colormap to use for colorizing the disparity map. Used for on-device postprocessing.
@@ -432,6 +453,7 @@ def disparity(self, pipeline: dai.Pipeline, device: dai.Device) -> XoutBase:
                 colormap=self._comp._postprocess_colormap,
                 wls_config=self._comp.wls_config,
                 ir_settings=self._comp.ir_settings,
+                confidence_map=self._try_get_confidence_map()
             )
 
             return self._comp._create_xout(pipeline, out)
@@ -464,7 +486,8 @@ def depth(self, pipeline: dai.Pipeline, device: dai.Device) -> XoutBase:
                 colorize=self._comp._colorize,
                 colormap=self._comp._postprocess_colormap,
                 wls_config=self._comp.wls_config,
-                ir_settings=self._comp.ir_settings
+                ir_settings=self._comp.ir_settings,
+                confidence_map=self._try_get_confidence_map()
             )
             return self._comp._create_xout(pipeline, out)
 
@@ -489,3 +512,11 @@ def encoded(self, pipeline: dai.Pipeline, device: dai.Device) -> XoutBase:
                                frame_shape=(1200, 800))
 
             return self._comp._create_xout(pipeline, out)
+
+        def _try_get_confidence_map(self):
+            confidence_map = None
+            if self._comp.enable_depth_score:
+                confidence_map = StreamXout(self._comp.node.id,
+                                            self._comp.node.confidenceMap,
+                                            name='depth_score')
+            return confidence_map

depthai_sdk/src/depthai_sdk/oak_outputs/xout/xout_depth.py

@@ -24,7 +24,8 @@ def __init__(self,
                  colorize: StereoColor = None,
                  colormap: int = None,
                  wls_config: dict = None,
-                 ir_settings: dict = None):
+                 ir_settings: dict = None,
+                 confidence_map: StreamXout = None):
         self.name = 'Depth'
         super().__init__(device=device,
                          frames=frames,
@@ -34,7 +35,8 @@ def __init__(self,
                          colorize=colorize,
                          colormap=colormap,
                          wls_config=wls_config,
-                         ir_settings=ir_settings)
+                         ir_settings=ir_settings,
+                         confidence_map=confidence_map)
 
         self.disp_scale_factor = dispScaleFactor
 

depthai_sdk/src/depthai_sdk/oak_outputs/xout/xout_disparity.py

@@ -2,6 +2,7 @@
 import logging
 import warnings
 from collections import defaultdict
+from functools import cached_property
 from typing import List, Optional
 
 import depthai as dai
@@ -29,13 +30,17 @@ def __init__(self,
                  colorize: StereoColor = None,
                  colormap: int = None,
                  wls_config: dict = None,
-                 ir_settings: dict = None):
+                 ir_settings: dict = None,
+                 confidence_map: StreamXout = None):
         self.mono_frames = mono_frames
         self.multiplier = disp_factor
         self.fps = fps
         self.name = 'Disparity'
         self.device = device
 
+        self.confidence_map = confidence_map
+        self.fig, self.axes = None, None  # for depth score visualization
+
         self.colorize = colorize
         self.colormap = colormap
         self.use_wls_filter = wls_config['enabled']
@@ -127,12 +132,47 @@ def visualize(self, packet: DepthPacket):
                 self._visualizer.add_circle(coords=(x, y), radius=3, color=(255, 255, 255), thickness=-1)
                 self._visualizer.add_text(text=text, coords=(x, y - 10))
 
+        if self._visualizer.config.stereo.depth_score and packet.confidence_map:
+            self.fig.canvas.draw()
+            self.axes.clear()
+            self.axes.hist(255 - packet.confidence_map.getData(), bins=3, color='blue', alpha=0.5)
+            self.axes.set_title(f'Depth score: {packet.depth_score:.2f}')
+            self.axes.set_xlabel('Depth score')
+            self.axes.set_ylabel('Frequency')
+
+            # self.axes.text(0.5, 0.9, f'Overall depth score: {packet.depth_score:.2f}', ha='center', va='center',
+            #                transform=self.axes.transAxes, fontsize=20)
+            # Convert plot to numpy array
+            img = np.fromstring(self.fig.canvas.tostring_rgb(), dtype=np.uint8, sep='')
+            img = img.reshape(self.fig.canvas.get_width_height()[::-1] + (3,))
+            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
+            cv2.imshow(f'Depth score ({self.name})', img)
+
         super().visualize(packet)
 
     def xstreams(self) -> List[StreamXout]:
-        if self.mono_frames is None:
-            return [self.frames]
-        return [self.frames, self.mono_frames]
+        streams = [self.frames]
+        if self.mono_frames is not None:
+            streams.append(self.mono_frames)
+        if self.confidence_map is not None:
+            streams.append(self.confidence_map)
+
+        return streams
+
+    def setup_visualize(self,
+                        visualizer: 'Visualizer',
+                        visualizer_enabled: bool,
+                        name: str = None
+                        ) -> None:
+        super().setup_visualize(visualizer, visualizer_enabled, name)
+
+        if self.confidence_map:
+            from matplotlib import pyplot as plt
+            self.fig, self.axes = plt.subplots(1, 1, figsize=(8, 4), constrained_layout=False)
+
+    @cached_property
+    def stream_names(self) -> List[str]:
+        return [s.name for s in self.xstreams()]
 
     def new_msg(self, name: str, msg: dai.Buffer) -> None:
         if name not in self._streams:
@@ -144,9 +184,7 @@ def new_msg(self, name: str, msg: dai.Buffer) -> None:
         if seq not in self.msgs:
             self.msgs[seq] = dict()
 
-        if name == self.frames.name:
-            self.msgs[seq][name] = msg
-        elif name == self.mono_frames.name:
+        if name in self.stream_names:
             self.msgs[seq][name] = msg
         else:
             raise ValueError('Message from unknown stream name received by TwoStageSeqSync!')
@@ -156,13 +194,16 @@ def new_msg(self, name: str, msg: dai.Buffer) -> None:
             if self.queue.full():
                 self.queue.get()  # Get one, so queue isn't full
 
-            mono_frame = None
+            mono_frame, confidence_map = None, None
             if self.mono_frames is not None:
                 mono_frame = self.msgs[seq][self.mono_frames.name]
+            if self.confidence_map is not None:
+                confidence_map = self.msgs[seq][self.confidence_map.name]
 
             packet = DepthPacket(
                 self.get_packet_name(),
                 img_frame=self.msgs[seq][self.frames.name],
+                confidence_map=confidence_map,
                 mono_frame=mono_frame,
                 visualizer=self._visualizer
             )

depthai_sdk/src/depthai_sdk/visualize/configs.py

@@ -55,6 +55,7 @@ class StereoConfig:
     wls_filter: bool = False
     wls_lambda: float = 8000
     wls_sigma: float = 1.5
+    depth_score: bool = False
 
 @dataclass
 class DetectionConfig:

depthai_sdk/src/depthai_sdk/visualize/visualizer.py

@@ -318,7 +318,8 @@ def stereo(self,
                colormap: int = None,
                wls_filter: bool = None,
                wls_lambda: float = None,
-               wls_sigma: float = None):
+               wls_sigma: float = None,
+               depth_score: bool = None):
         kwargs = self._process_kwargs(locals())
 
         if len(kwargs) > 0: