Depth registration

CMakeLists.txt

@@ -90,8 +90,24 @@ add_library(${PLUGIN_NAME} SHARED
     src/load_image_to_cuda.cu
 )
 
-# Unity Native-plugin API includes TODO: cleanup
-set(UnityPluginAPI_INCLUDE "C:/Program Files/Unity/Hub/Editor/6000.2.8f1/Editor/Data/PluginAPI")
+# Unity Native-plugin API headers (IUnityInterface.h etc.).
+# Auto-detects the newest Unity Hub install; override with -DUnityPluginAPI_INCLUDE=<path>.
+if(NOT UnityPluginAPI_INCLUDE)
+    file(GLOB _unity_editor_dirs "C:/Program Files/Unity/Hub/Editor/*")
+    list(SORT _unity_editor_dirs COMPARE NATURAL ORDER DESCENDING)
+    foreach(_unity_dir ${_unity_editor_dirs})
+        if(EXISTS "${_unity_dir}/Editor/Data/PluginAPI/IUnityInterface.h")
+            set(UnityPluginAPI_INCLUDE "${_unity_dir}/Editor/Data/PluginAPI")
+            break()
+        endif()
+    endforeach()
+endif()
+set(UnityPluginAPI_INCLUDE "${UnityPluginAPI_INCLUDE}" CACHE PATH "Path to Unity's Native Plugin API headers")
+if(NOT EXISTS "${UnityPluginAPI_INCLUDE}/IUnityInterface.h")
+    message(FATAL_ERROR "Unity Native Plugin API headers not found (IUnityInterface.h). "
+        "Set -DUnityPluginAPI_INCLUDE=<Unity>/Editor/Data/PluginAPI for your installed Unity version.")
+endif()
+message(STATUS "Using Unity Plugin API headers: ${UnityPluginAPI_INCLUDE}")
 
 # Include directories
 # TODO: Use target_include_directories instead of include_directories

PySpotObserver/examples/basic_streaming.py

@@ -26,6 +26,7 @@
     build_config_from_args,
     parse_camera_list,
 )
+
 from pyspotobserver import CameraType, SpotConfig, SpotConnection
 
 logging.basicConfig(

PySpotObserver/pyspotobserver/__init__.py

@@ -5,17 +5,17 @@
 and streaming camera data with support for both synchronous and asynchronous patterns.
 """
 
-from .config import SpotConfig, CameraType
-from .connection import SpotAuthenticationError, SpotConnection, SpotConnectionError
 from .camera_stream import SpotCamStream, SpotCamStreamError
+from .config import CameraType, SpotConfig
+from .connection import SpotAuthenticationError, SpotConnection, SpotConnectionError
 
 __version__ = "0.1.0"
 __all__ = [
-    "SpotConfig",
     "CameraType",
-    "SpotConnection",
-    "SpotConnectionError",
     "SpotAuthenticationError",
     "SpotCamStream",
     "SpotCamStreamError",
+    "SpotConfig",
+    "SpotConnection",
+    "SpotConnectionError",
 ]

PySpotObserver/pyspotobserver/camera_stream.py

@@ -20,6 +20,11 @@
 
 from .color_correction import _ROBOT_CCMS
 from .config import CameraType, SpotConfig
+from .depth_registration import (
+    DepthRegistrationParams,
+    extract_registration_params,
+    register_depth,
+)
 from .stitch import (
     STITCH_OUT_H,
     STITCH_OUT_W,
@@ -53,7 +58,8 @@ class ImageFrame:
 
     Attributes:
         rgb_images: List of RGB images as numpy arrays (H, W, 3) in float32 [0, 1]
-        depth_images: List of depth images as numpy arrays (H, W) in float32 (meters)
+        depth_images: List of depth images as numpy arrays (H, W) in float32 (meters),
+            registered into the RGB camera frame (same H, W as the RGB images)
         camera_order: List of CameraType enums indicating order of images
         timestamp: Time when frame was captured (monotonic clock)
         acquisition_time: Robot's acquisition time from image response
@@ -148,6 +154,10 @@ def __init__(
         self._frame_pool_index: int = 0
         self._image_requests: list[image_pb2.ImageRequest] = []
 
+        # Per-camera raw-depth -> RGB-frame registration params, cached from the
+        # first response batch (fixed by the mechanical mounting).
+        self._depth_reg_params: list[DepthRegistrationParams] | None = None
+
         # Optional per-stream vision pipeline, imported lazily.
         self._vision_pipeline: VisionPipeline | None = None
 
@@ -240,6 +250,7 @@ def start_streaming(self, camera_mask: int) -> None:
         self._frame_pool = []
         self._frame_pool_index = 0
         self._ccm_scratch_by_shape = {}
+        self._depth_reg_params = None
         self._frame_count = 0
         self._error_count = 0
         self._last_body_to_worlds = None
@@ -828,10 +839,13 @@ def _fill_frame_from_responses(
                 responses[rgb_idx], is_depth=False, out_array=frame.rgb_images[frame_i], ccm=ccm
             )
 
-            self._convert_image_response_inplace(
-                responses[depth_idx],
-                is_depth=True,
-                out_array=frame.depth_images[frame_i],
+            if self._depth_reg_params is None:
+                raise SpotCamStreamError("Depth registration params not initialized")
+            raw_depth = self._convert_image_response_alloc(responses[depth_idx], is_depth=True)
+            register_depth(
+                raw_depth,
+                self._depth_reg_params[i],
+                out=frame.depth_images[frame_i],
             )
 
         # Update timestamps and camera extrinsics params
@@ -870,6 +884,10 @@ def _decode_initial_responses(
     ) -> list[np.ndarray]:
         """
         Decode initial responses once for shape inference and first frame fill.
+
+        Also caches the per-camera depth registration params: the raw depth
+        sources live in the depth sensor's frame, so every depth image is
+        reprojected into the RGB camera frame (RGB-sized) before use.
         """
         n_cameras = len(self._sdk_camera_order)
         expected_count = n_cameras * 2  # RGB + depth per camera
@@ -878,14 +896,32 @@ def _decode_initial_responses(
             raise SpotCamStreamError(f"Expected {expected_count} responses, got {len(responses)}")
 
         decoded: list[np.ndarray] = []
+        reg_params: list[DepthRegistrationParams] = []
         for i in range(n_cameras):
             rgb_idx = i * 2
             depth_idx = i * 2 + 1
             rgb = self._convert_image_response_alloc(responses[rgb_idx], is_depth=False)
             if self._ccms is not None:
                 self._apply_ccm_inplace(rgb, self._ccms[self._sdk_camera_order[i]])
             decoded.append(rgb)
-            decoded.append(self._convert_image_response_alloc(responses[depth_idx], is_depth=True))
+
+            raw_depth = self._convert_image_response_alloc(responses[depth_idx], is_depth=True)
+            try:
+                params = extract_registration_params(
+                    responses[rgb_idx], responses[depth_idx], dst_shape=rgb.shape[:2]
+                )
+            except ValueError as exc:
+                raise SpotCamStreamError(
+                    f"Failed to build depth registration params for "
+                    f"{self._sdk_camera_order[i].name}: {exc}"
+                ) from exc
+            reg_params.append(params)
+
+            registered = np.zeros(rgb.shape[:2], dtype=np.float32)
+            register_depth(raw_depth, params, registered)
+            decoded.append(registered)
+
+        self._depth_reg_params = reg_params
         return decoded
 
     def _cache_stitch_params(self, responses: list[image_pb2.ImageResponse]) -> None:

PySpotObserver/pyspotobserver/color_correction.py

@@ -1,4 +1,5 @@
 import numpy as np
+
 from .config import CameraType
 
 # Known robot IPs

PySpotObserver/pyspotobserver/config.py

@@ -36,14 +36,16 @@ def get_source_name(cls, camera: "CameraType", depth: bool = False) -> str:
             # Hand camera uses color sensor naming, not fisheye naming.
             cls.HAND: "hand_color_image",
         }
+        # Raw depth-sensor-frame sources: substantially smaller on the wire than
+        # the robot-registered *_depth_in_visual_frame sources. Alignment with the
+        # RGB image is done client-side (see depth_registration.register_depth).
         depth_names = {
-            cls.BACK: "back_depth_in_visual_frame",
-            cls.FRONTLEFT: "frontleft_depth_in_visual_frame",
-            cls.FRONTRIGHT: "frontright_depth_in_visual_frame",
-            cls.LEFT: "left_depth_in_visual_frame",
-            cls.RIGHT: "right_depth_in_visual_frame",
-            # Hand depth is aligned to hand color frame.
-            cls.HAND: "hand_depth_in_hand_color_frame",
+            cls.BACK: "back_depth",
+            cls.FRONTLEFT: "frontleft_depth",
+            cls.FRONTRIGHT: "frontright_depth",
+            cls.LEFT: "left_depth",
+            cls.RIGHT: "right_depth",
+            cls.HAND: "hand_depth",
         }
 
         if camera not in rgb_names:

PySpotObserver/pyspotobserver/depth_registration.py

@@ -0,0 +1,156 @@
+"""Client-side registration of raw depth images into the RGB camera frame.
+
+Spot's ``*_depth_in_visual_frame`` sources are the robot doing this registration
+onboard, at the cost of shipping RGB-sized depth images over the wire. Requesting
+the raw depth-sensor sources (e.g. ``frontleft_depth``) is substantially cheaper,
+but the pixels live in the depth sensor's frame. This module reprojects them into
+the RGB camera's image plane so downstream consumers see the same aligned,
+RGB-sized depth they would get from the robot-registered sources.
+
+Mirrors the C++ implementation (``register_depth_to_rgb`` in cuda_kernels.cu).
+"""
+
+from dataclasses import dataclass
+
+import numpy as np
+from bosdyn.api import image_pb2  # type: ignore[import-untyped]
+from bosdyn.client.frame_helpers import (  # type: ignore[import-untyped]
+    BODY_FRAME_NAME,
+    get_a_tform_b,
+)
+
+
+@dataclass
+class DepthRegistrationParams:
+    """Folded reprojection from a raw depth image into an RGB image plane.
+
+    For a depth pixel (u, v) with depth z (meters):
+
+        p = z * rays[v, u] + t
+
+    where ``rays[v, u] = M @ [u, v, 1]`` with ``M = K_rgb @ R @ inv(K_depth)`` and
+    ``t = K_rgb @ translation`` (R, translation from rgb_T_depth). Then
+    ``(p[0] / p[2], p[1] / p[2])`` is the RGB pixel and ``p[2]`` is the depth in
+    the RGB camera's frame.
+    """
+
+    rays: np.ndarray  # (H_src, W_src, 3) float32
+    t: np.ndarray  # (3,) float32
+    src_shape: tuple[int, int]  # (H_src, W_src)
+    dst_shape: tuple[int, int]  # (H_dst, W_dst)
+
+
+def _pinhole_intrinsics(source: image_pb2.ImageSource) -> np.ndarray:
+    """Return the 3x3 pinhole camera matrix of an ImageSource."""
+    which = source.WhichOneof("camera_models")
+    if which == "pinhole":
+        intr = source.pinhole.intrinsics
+    elif which == "pinhole_brown_conrady":
+        intr = source.pinhole_brown_conrady.intrinsics.pinhole_intrinsics
+    elif which == "kannala_brandt":
+        intr = source.kannala_brandt.intrinsics.pinhole_intrinsics
+    else:
+        raise ValueError(f"No pinhole intrinsics for {source.name!r} (model={which!r})")
+
+    return np.array(
+        [
+            [intr.focal_length.x, intr.skew.x, intr.principal_point.x],
+            [intr.skew.y, intr.focal_length.y, intr.principal_point.y],
+            [0.0, 0.0, 1.0],
+        ],
+        dtype=np.float64,
+    )
+
+
+def _body_t_sensor(response: image_pb2.ImageResponse) -> np.ndarray:
+    snapshot = response.shot.transforms_snapshot
+    frame_name = response.shot.frame_name_image_sensor
+    pose = get_a_tform_b(snapshot, BODY_FRAME_NAME, frame_name)
+    if pose is None:
+        raise ValueError(f"No transform from body to {frame_name!r}")
+    return pose.to_matrix()
+
+
+def extract_registration_params(
+    rgb_response: image_pb2.ImageResponse,
+    depth_response: image_pb2.ImageResponse,
+    dst_shape: tuple[int, int],
+) -> DepthRegistrationParams:
+    """Build registration params from a time-synced RGB + raw depth response pair.
+
+    Intrinsics and extrinsics are fixed by the mechanical mounting, so this is
+    computed once per stream and reused for every frame.
+    """
+    body_T_rgb = _body_t_sensor(rgb_response)
+    body_T_depth = _body_t_sensor(depth_response)
+    rgb_T_depth = np.linalg.inv(body_T_rgb) @ body_T_depth
+
+    k_rgb = _pinhole_intrinsics(rgb_response.source)
+    k_depth = _pinhole_intrinsics(depth_response.source)
+
+    M = k_rgb @ rgb_T_depth[:3, :3] @ np.linalg.inv(k_depth)
+    t = k_rgb @ rgb_T_depth[:3, 3]
+
+    src_h = depth_response.shot.image.rows
+    src_w = depth_response.shot.image.cols
+    v, u = np.indices((src_h, src_w), dtype=np.float64)
+    uv1 = np.stack([u, v, np.ones_like(u)], axis=-1)  # (H_src, W_src, 3)
+    rays = (uv1 @ M.T).astype(np.float32)
+
+    return DepthRegistrationParams(
+        rays=rays,
+        t=t.astype(np.float32),
+        src_shape=(src_h, src_w),
+        dst_shape=(int(dst_shape[0]), int(dst_shape[1])),
+    )
+
+
+def register_depth(
+    raw_depth: np.ndarray,
+    params: DepthRegistrationParams,
+    out: np.ndarray,
+) -> None:
+    """Reproject a raw depth image into the RGB frame, in-place into ``out``.
+
+    Args:
+        raw_depth: (H_src, W_src) float32 depth in meters, 0 = invalid.
+        params: Registration params for this camera pair.
+        out: (H_dst, W_dst) float32 output, fully overwritten. Pixels that
+            receive no depth sample are 0.
+    """
+    if raw_depth.shape != params.src_shape:
+        raise ValueError(f"Raw depth shape {raw_depth.shape} != expected {params.src_shape}")
+    if out.shape != params.dst_shape:
+        raise ValueError(f"Output shape {out.shape} != expected {params.dst_shape}")
+
+    out.fill(0.0)
+
+    valid = raw_depth > 0
+    if not valid.any():
+        return
+
+    z = raw_depth[valid]
+    p = z[:, None] * params.rays[valid] + params.t
+    pz = p[:, 2]
+
+    front = pz > 0  # discard points behind the RGB camera
+    p, pz = p[front], pz[front]
+    if pz.size == 0:
+        return
+
+    u = np.floor(p[:, 0] / pz).astype(np.int32)
+    v = np.floor(p[:, 1] / pz).astype(np.int32)
+
+    # 2x2 splat: the raw depth image is lower resolution than the RGB target, so
+    # covering the four nearest pixels reduces holes without a dilation pass.
+    xs = np.concatenate((u, u + 1, u, u + 1))
+    ys = np.concatenate((v, v, v + 1, v + 1))
+    zs = np.concatenate((pz, pz, pz, pz))
+
+    dst_h, dst_w = params.dst_shape
+    in_bounds = (xs >= 0) & (xs < dst_w) & (ys >= 0) & (ys < dst_h)
+    xs, ys, zs = xs[in_bounds], ys[in_bounds], zs[in_bounds]
+
+    # Paint far-to-near so the nearest surface wins where samples overlap.
+    order = np.argsort(zs)[::-1]
+    out[ys[order], xs[order]] = zs[order]

PySpotObserver/tests/test_config.py

@@ -27,19 +27,10 @@ def test_get_source_name_rgb(self):
         assert CameraType.get_source_name(CameraType.HAND) == "hand_color_image"
 
     def test_get_source_name_depth(self):
-        """Test depth camera source name generation."""
-        assert (
-            CameraType.get_source_name(CameraType.FRONTLEFT, depth=True)
-            == "frontleft_depth_in_visual_frame"
-        )
-        assert (
-            CameraType.get_source_name(CameraType.RIGHT, depth=True)
-            == "right_depth_in_visual_frame"
-        )
-        assert (
-            CameraType.get_source_name(CameraType.HAND, depth=True)
-            == "hand_depth_in_hand_color_frame"
-        )
+        """Test depth camera source name generation (raw depth-sensor sources)."""
+        assert CameraType.get_source_name(CameraType.FRONTLEFT, depth=True) == "frontleft_depth"
+        assert CameraType.get_source_name(CameraType.RIGHT, depth=True) == "right_depth"
+        assert CameraType.get_source_name(CameraType.HAND, depth=True) == "hand_depth"
 
 
 class TestSpotConfig: