Added align_frame to packets, added aligned_to to tof packet (dispari…

…tyDepth). Created tof align with color example

depthai_sdk/examples/ToFComponent/tof_align.py

@@ -0,0 +1,19 @@
+from depthai_sdk import OakCamera
+from depthai_sdk.classes.packets import DisparityDepthPacket
+import cv2
+from depthai_sdk.visualize.visualizer import Visualizer
+
+with OakCamera() as oak:
+    cam_c = oak.create_camera('CAM_C')
+    tof = oak.create_tof("CAM_A", align_to=cam_c)
+    depth_q = oak.queue(tof.out.depth).queue
+
+    vis = Visualizer() # Only for depth colorization
+    oak.start()
+    while oak.running():
+        depth: DisparityDepthPacket = depth_q.get()
+        colored_depth = depth.get_colorized_frame(vis)
+        cv2.imshow("depth", colored_depth)
+        cv2.imshow('Weighted', cv2.addWeighted(depth.aligned_frame.getCvFrame(), 0.5, colored_depth, 0.5, 0))
+        if cv2.waitKey(1) == ord('q'):
+            break

depthai_sdk/src/depthai_sdk/classes/packets.py

@@ -119,14 +119,14 @@ def __init__(self,
                  disparity_map: Optional[np.ndarray] = None,
                  colorize: StereoColor = None,
                  colormap: int = None,
-                 mono_frame: Optional[dai.ImgFrame] = None,
+                 aligned_frame: Optional[dai.ImgFrame] = None,
                  confidence_map: Optional[np.ndarray] = None
                  ):
         """
         disparity_map might be filtered, eg. if WLS filter is enabled
         """
         super().__init__(name=name, msg=img)
-        self.mono_frame = mono_frame
+        self.aligned_frame = aligned_frame
         self.disparity_map = disparity_map
         self.multiplier = multiplier
         self.colorize = colorize
@@ -150,9 +150,9 @@ def get_colorized_frame(self, visualizer) -> np.ndarray:
         colorized_disp = frame * self.multiplier
 
         try:
-            mono_frame = self.mono_frame.getCvFrame()
+            aligned_frame = self.aligned_frame.getCvFrame()
         except AttributeError:
-            mono_frame = None
+            aligned_frame = None
 
         stereo_config = visualizer.config.stereo
 
@@ -163,7 +163,7 @@ def get_colorized_frame(self, visualizer) -> np.ndarray:
             colormap = stereo_config.colormap
             colormap[0] = [0, 0, 0]  # Invalidate pixels 0 to be black
 
-        if mono_frame is not None and colorized_disp.ndim == 2 and mono_frame.ndim == 3:
+        if aligned_frame is not None and colorized_disp.ndim == 2 and aligned_frame.ndim == 3:
             colorized_disp = colorized_disp[..., np.newaxis]
 
         if colorize == StereoColor.GRAY:
@@ -172,7 +172,7 @@ def get_colorized_frame(self, visualizer) -> np.ndarray:
             colorized_disp = cv2.applyColorMap(colorized_disp.astype(np.uint8), colormap)
         elif colorize == StereoColor.RGBD:
             colorized_disp = cv2.applyColorMap(
-                (colorized_disp + mono_frame * 0.5).astype(np.uint8), colormap
+                (colorized_disp + aligned_frame * 0.5).astype(np.uint8), colormap
             )
         return colorized_disp
 
@@ -190,7 +190,7 @@ def __init__(self,
                  img_frame: dai.ImgFrame,
                  colorize: StereoColor = None,
                  colormap: int = None,
-                 mono_frame: Optional[dai.ImgFrame] = None,
+                 aligned_frame: Optional[dai.ImgFrame] = None,
                  disp_scale_factor=255 / 95,
                  confidence_map=None
                  ):
@@ -202,7 +202,7 @@ def __init__(self,
             multiplier=255 / 95,
             colorize=colorize,
             colormap=colormap,
-            mono_frame=mono_frame,
+            aligned_frame=aligned_frame,
             confidence_map=confidence_map
         )
         self.disp_scale_factor = disp_scale_factor

depthai_sdk/src/depthai_sdk/components/camera_helper.py

@@ -221,6 +221,7 @@ def getClosesResolution(sensor: dai.CameraFeatures,
     desired, i = (width, 0) if width is not None else (height, 1)
 
     resolutions = [get_sensor_resolution(type, conf.width, conf.height) for conf in sensor.configs if conf.type == type]
+    resolutions = [res for res in resolutions if res is not None]
 
     for (res, size) in resolutions:
         err = abs(size[i] - desired)

depthai_sdk/src/depthai_sdk/components/stereo_component.py

@@ -450,14 +450,14 @@ def get_fourcc(self) -> Optional[str]:
     Available outputs (to the host) of this component
     """
 
-    def _mono_frames(self):
+    def _aligned_frames(self):
         """
-        Create mono frames output if WLS filter is enabled or colorize is set to RGBD
+        Create aligned frames output if WLS filter is enabled or colorize is set to RGBD
         """
-        mono_frames = None
+        aligned_frame = None
         if self.wls_config['enabled'] or self._colorize == StereoColor.RGBD:
-            mono_frames = StreamXout(self._right_stream)
-        return mono_frames
+            aligned_frame = StreamXout(self._right_stream)
+        return aligned_frame
 
     def _try_get_confidence_map(self):
         confidence_map = None
@@ -472,7 +472,7 @@ def __call__(self, device: dai.Device) -> XoutBase:
                     device=device,
                     frames=StreamXout(self._comp.depth),
                     dispScaleFactor=depth_to_disp_factor(device, self._comp.node),
-                    mono_frames=self._comp._mono_frames(),
+                    aligned_frame=self._comp._aligned_frames(),
                     colorize=self._comp._colorize,
                     colormap=self._comp._postprocess_colormap,
                     ir_settings=self._comp.ir_settings,
@@ -487,7 +487,7 @@ def __call__(self, device: dai.Device, fourcc: Optional[str] = None) -> XoutBase
                     frames=StreamXout(self._comp.encoder.bitstream) if fourcc else
                     StreamXout(self._comp.disparity),
                     disp_factor=255.0 / self._comp.node.getMaxDisparity(),
-                    mono_frames=self._comp._mono_frames(),
+                    aligned_frame=self._comp._aligned_frames(),
                     colorize=self._comp._colorize,
                     fourcc=fourcc,
                     colormap=self._comp._postprocess_colormap,

depthai_sdk/src/depthai_sdk/components/tof_component.py

@@ -48,8 +48,9 @@ def __init__(
 
         if align_to is not None:
             self._align = pipeline.create(dai.node.ImageAlign)
+            self._align_to_output = align_to.node.isp
             self.node.depth.link(self._align.input)
-            align_to.node.isp.link(self._align.inputAlignTo)
+            self._align_to_output.link(self._align.inputAlignTo)
 
     def _find_tof(self, device: dai.Device) -> dai.CameraBoardSocket:
         # Use the first ToF sensor, usually, there will only be one
@@ -89,8 +90,8 @@ def __call__(self, device: dai.Device) -> XoutBase:
                         ),
                         "tof_depth",
                     ),
+                    aligned_frame=StreamXout(self._comp._align_to_output, "aligned_stream") if self._comp._align else None,
                     dispScaleFactor=9500,
-                    mono_frames=None,
                     ir_settings={"auto_mode": False},
                 ).set_comp_out(self)
 

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

@@ -13,7 +13,7 @@ def __init__(self,
                  device: dai.Device,
                  frames: StreamXout,
                  dispScaleFactor: float,
-                 mono_frames: Optional[StreamXout],
+                 aligned_frame: Optional[StreamXout],
                  colorize: StereoColor = None,
                  colormap: int = None,
                  ir_settings: dict = None,
@@ -22,7 +22,7 @@ def __init__(self,
         super().__init__(device=device,
                          frames=frames,
                          disp_factor=255 / 95,
-                         mono_frames=mono_frames,
+                         aligned_frame=aligned_frame,
                          colorize=colorize,
                          colormap=colormap,
                          ir_settings=ir_settings,
@@ -31,14 +31,14 @@ def __init__(self,
         self.disp_scale_factor = dispScaleFactor
 
     def package(self, msgs: Dict) -> DisparityDepthPacket:
-        mono_frame = msgs[self.mono_frames.name] if self.mono_frames else None
+        aligned_frame = msgs[self.aligned_frame.name] if self.aligned_frame else None
         confidence_map = msgs[self.confidence_map.name] if self.confidence_map else None
         return DisparityDepthPacket(
             self.get_packet_name(),
             msgs[self.frames.name],
             colorize=self.colorize,
             colormap=self.colormap,
-            mono_frame=mono_frame,
+            aligned_frame=aligned_frame,
             disp_scale_factor=self.disp_scale_factor,
             confidence_map=confidence_map
         )

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

@@ -30,14 +30,14 @@ def __init__(self,
                  device: dai.Device,
                  frames: StreamXout,
                  disp_factor: float,
-                 mono_frames: Optional[StreamXout],
+                 aligned_frame: Optional[StreamXout],
                  fourcc: str = None,
                  colorize: StereoColor = None,
                  colormap: int = None,
                  wls_config: dict = None,
                  ir_settings: dict = None,
                  confidence_map: StreamXout = None):
-        self.mono_frames = mono_frames
+        self.aligned_frame = aligned_frame
         self.name = 'Disparity'
         self.multiplier = disp_factor
         self.device = device
@@ -96,16 +96,16 @@ def on_callback(self, packet) -> None:
 
     def xstreams(self) -> List[StreamXout]:
         streams = [self.frames]
-        if self.mono_frames is not None:
-            streams.append(self.mono_frames)
+        if self.aligned_frame is not None:
+            streams.append(self.aligned_frame)
         if self.confidence_map is not None:
             streams.append(self.confidence_map)
 
         return streams
 
     def package(self, msgs: Dict) -> DisparityPacket:
         img_frame = msgs[self.frames.name]
-        mono_frame = msgs[self.mono_frames.name] if self.mono_frames else None
+        aligned_frame = msgs[self.aligned_frame.name] if self.aligned_frame else None
         confidence_map = msgs[self.confidence_map.name] if self.confidence_map else None
         # TODO: refactor the mess below
         packet = DisparityPacket(
@@ -116,7 +116,7 @@ def package(self, msgs: Dict) -> DisparityPacket:
             colorize=self.colorize,
             colormap=self.colormap,
             confidence_map=confidence_map,
-            mono_frame=mono_frame,
+            aligned_frame=aligned_frame,
         )
         packet._get_codec = self.get_codec
 
@@ -127,10 +127,10 @@ def package(self, msgs: Dict) -> DisparityPacket:
             if disparity_frame is None:
                 return None
 
-        if mono_frame and self.use_wls_filter:
+        if aligned_frame and self.use_wls_filter:
             # Perform WLS filtering
             # If we have wls enabled, it means CV2 is installed
-            disparity_frame = self.wls_filter.filter(disparity_frame, mono_frame.getCvFrame())
+            disparity_frame = self.wls_filter.filter(disparity_frame, aligned_frame.getCvFrame())
 
         packet.disparity_map = disparity_frame