From 6c158803973da85c2e04ea18b3822c150a857283 Mon Sep 17 00:00:00 2001 From: Faisal Zaghloul Date: Tue, 7 Jul 2026 19:32:39 -0400 Subject: [PATCH 1/6] Support depth registration This is essentially what Spot does onboard for its *_depth_in_visual_frame sources. This allows us to request the *_depth targets instead, which are ~1/3 the size (with no additional useful depth info). Also Added a README file. --- CMakeLists.txt | 20 +- PySpotObserver/examples/config_example.yaml | 4 +- .../pyspotobserver/camera_stream.py | 48 +++- PySpotObserver/pyspotobserver/config.py | 16 +- .../pyspotobserver/depth_registration.py | 156 +++++++++++++ PySpotObserver/tests/test_config.py | 17 +- .../tests/test_depth_registration.py | 154 +++++++++++++ .../tests/test_depth_registration_robot.py | 174 +++++++++++++++ README.md | 105 +++++++++ src/cuda_kernels.cu | 85 +++++++ src/include/cuda_kernels.cuh | 23 ++ src/include/spot-connection.h | 14 +- src/include/vision-pipeline.h | 4 - src/spot-connection.cpp | 211 +++++++++++++++--- tests/integ-test/integ-test.cpp | 4 +- 15 files changed, 972 insertions(+), 63 deletions(-) create mode 100644 PySpotObserver/pyspotobserver/depth_registration.py create mode 100644 PySpotObserver/tests/test_depth_registration.py create mode 100644 PySpotObserver/tests/test_depth_registration_robot.py create mode 100644 README.md diff --git a/CMakeLists.txt b/CMakeLists.txt index 6992205..606a7b7 100644 --- a/CMakeLists.txt +++ b/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=. +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=/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 diff --git a/PySpotObserver/examples/config_example.yaml b/PySpotObserver/examples/config_example.yaml index bec26cb..e1f635a 100644 --- a/PySpotObserver/examples/config_example.yaml +++ b/PySpotObserver/examples/config_example.yaml @@ -4,8 +4,8 @@ # TUSKER: "128.148.138.22" # GOUGER: "128.148.138.21" robot_ip: "128.148.138.22" -username: "" -password: "" +username: "user" +password: "bigbubbabigbubba" # Streaming settings image_buffer_size: 5 diff --git a/PySpotObserver/pyspotobserver/camera_stream.py b/PySpotObserver/pyspotobserver/camera_stream.py index 538c95b..4ddf94c 100644 --- a/PySpotObserver/pyspotobserver/camera_stream.py +++ b/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,6 +896,7 @@ 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 @@ -885,7 +904,24 @@ def _decode_initial_responses( 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: diff --git a/PySpotObserver/pyspotobserver/config.py b/PySpotObserver/pyspotobserver/config.py index 58b33e8..8aa1464 100644 --- a/PySpotObserver/pyspotobserver/config.py +++ b/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: diff --git a/PySpotObserver/pyspotobserver/depth_registration.py b/PySpotObserver/pyspotobserver/depth_registration.py new file mode 100644 index 0000000..7c40396 --- /dev/null +++ b/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] diff --git a/PySpotObserver/tests/test_config.py b/PySpotObserver/tests/test_config.py index 50dd4a1..c6feb9d 100644 --- a/PySpotObserver/tests/test_config.py +++ b/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: diff --git a/PySpotObserver/tests/test_depth_registration.py b/PySpotObserver/tests/test_depth_registration.py new file mode 100644 index 0000000..d9489fb --- /dev/null +++ b/PySpotObserver/tests/test_depth_registration.py @@ -0,0 +1,154 @@ +""" +Unit tests for client-side depth registration. +""" + +import numpy as np +import pytest +from bosdyn.api import image_pb2 # type: ignore[import-untyped] +from pyspotobserver.depth_registration import ( + DepthRegistrationParams, + extract_registration_params, + register_depth, +) + + +def make_params(M, t, src_shape, dst_shape) -> DepthRegistrationParams: + """Build params directly from a projection matrix, bypassing proto extraction.""" + v, u = np.indices(src_shape, dtype=np.float64) + uv1 = np.stack([u, v, np.ones_like(u)], axis=-1) + rays = (uv1 @ np.asarray(M, dtype=np.float64).T).astype(np.float32) + return DepthRegistrationParams( + rays=rays, + t=np.asarray(t, dtype=np.float32), + src_shape=src_shape, + dst_shape=dst_shape, + ) + + +class TestRegisterDepth: + def test_identity_maps_constant_plane(self): + """With an identity projection, a constant plane fills the whole output.""" + params = make_params(np.eye(3), np.zeros(3), src_shape=(6, 8), dst_shape=(6, 8)) + raw = np.full((6, 8), 2.0, dtype=np.float32) + out = np.empty((6, 8), dtype=np.float32) + + register_depth(raw, params, out) + + np.testing.assert_allclose(out, 2.0) + + def test_nearest_surface_wins(self): + """Two samples landing on the same pixel keep the nearer depth.""" + # Both source pixels project along the optical axis to dst pixel (0, 0). + params = make_params(np.zeros((3, 3)), np.zeros(3), src_shape=(1, 2), dst_shape=(2, 2)) + params.rays[0, 0] = (0.0, 0.0, 1.0) + params.rays[0, 1] = (0.0, 0.0, 1.0) + raw = np.array([[5.0, 1.0]], dtype=np.float32) + out = np.empty((2, 2), dtype=np.float32) + + register_depth(raw, params, out) + + np.testing.assert_allclose(out, 1.0) # 2x2 splat covers the full output + + def test_invalid_depth_leaves_holes(self): + params = make_params(np.eye(3), np.zeros(3), src_shape=(4, 4), dst_shape=(4, 4)) + raw = np.zeros((4, 4), dtype=np.float32) + raw[1, 1] = 3.0 + out = np.empty((4, 4), dtype=np.float32) + + register_depth(raw, params, out) + + expected = np.zeros((4, 4), dtype=np.float32) + expected[1:3, 1:3] = 3.0 # the sample plus its 2x2 splat + np.testing.assert_allclose(out, expected) + + def test_points_behind_camera_are_skipped(self): + params = make_params(-np.eye(3), np.zeros(3), src_shape=(4, 4), dst_shape=(4, 4)) + raw = np.full((4, 4), 2.0, dtype=np.float32) + out = np.empty((4, 4), dtype=np.float32) + + register_depth(raw, params, out) + + np.testing.assert_allclose(out, 0.0) + + def test_shape_mismatch_raises(self): + params = make_params(np.eye(3), np.zeros(3), src_shape=(4, 4), dst_shape=(4, 4)) + with pytest.raises(ValueError, match="Raw depth shape"): + register_depth(np.zeros((2, 2), dtype=np.float32), params, np.zeros((4, 4), np.float32)) + with pytest.raises(ValueError, match="Output shape"): + register_depth(np.zeros((4, 4), dtype=np.float32), params, np.zeros((2, 2), np.float32)) + + +def _make_response( + frame_name: str, + *, + rows: int, + cols: int, + fx: float, + fy: float, + cx: float, + cy: float, + translation: tuple[float, float, float] = (0.0, 0.0, 0.0), +) -> image_pb2.ImageResponse: + """Build a minimal ImageResponse with pinhole intrinsics and a body->sensor edge.""" + response = image_pb2.ImageResponse() + response.shot.frame_name_image_sensor = frame_name + response.shot.image.rows = rows + response.shot.image.cols = cols + + intr = response.source.pinhole.intrinsics + intr.focal_length.x = fx + intr.focal_length.y = fy + intr.principal_point.x = cx + intr.principal_point.y = cy + + snapshot = response.shot.transforms_snapshot + snapshot.child_to_parent_edge_map["body"].SetInParent() # tree root + edge = snapshot.child_to_parent_edge_map[frame_name] + edge.parent_frame_name = "body" + edge.parent_tform_child.rotation.w = 1.0 + edge.parent_tform_child.position.x = translation[0] + edge.parent_tform_child.position.y = translation[1] + edge.parent_tform_child.position.z = translation[2] + return response + + +class TestExtractRegistrationParams: + def test_coincident_cameras_reproject_in_place(self): + """Same pose and intrinsics for both cameras => registration is a no-op.""" + rgb = _make_response("rgb", rows=20, cols=30, fx=50.0, fy=50.0, cx=0.0, cy=0.0) + depth = _make_response("depth", rows=20, cols=30, fx=50.0, fy=50.0, cx=0.0, cy=0.0) + + params = extract_registration_params(rgb, depth, dst_shape=(20, 30)) + + raw = np.full((20, 30), 2.0, dtype=np.float32) + out = np.empty((20, 30), dtype=np.float32) + register_depth(raw, params, out) + + np.testing.assert_allclose(out, 2.0, rtol=1e-5) + + def test_translated_depth_camera_shifts_projection(self): + """A depth camera offset by baseline b shifts pixels by fx * b / z.""" + fx = 100.0 + rgb = _make_response("rgb", rows=20, cols=40, fx=fx, fy=fx, cx=0.0, cy=0.0) + depth = _make_response( + "depth", + rows=20, + cols=30, + fx=fx, + fy=fx, + cx=0.0, + cy=0.0, + translation=(0.1, 0.0, 0.0), + ) + + params = extract_registration_params(rgb, depth, dst_shape=(20, 40)) + + raw = np.zeros((20, 30), dtype=np.float32) + raw[7, 10] = 2.0 + out = np.empty((20, 40), dtype=np.float32) + register_depth(raw, params, out) + + # Expected shift: fx * 0.1 / 2.0 = 5 pixels in u, none in v. + expected = np.zeros((20, 40), dtype=np.float32) + expected[7:9, 15:17] = 2.0 # sample at (u=15, v=7) plus its 2x2 splat + np.testing.assert_allclose(out, expected, rtol=1e-5) diff --git a/PySpotObserver/tests/test_depth_registration_robot.py b/PySpotObserver/tests/test_depth_registration_robot.py new file mode 100644 index 0000000..74092a0 --- /dev/null +++ b/PySpotObserver/tests/test_depth_registration_robot.py @@ -0,0 +1,174 @@ +""" +Robot-in-the-loop validation of client-side depth registration. + +Registers the raw ``*_depth`` sources into the RGB frame with +``pyspotobserver.depth_registration`` and compares the result against the +robot's own ``*_depth_in_visual_frame`` output for the same (time-synced) +capture. Both images come from a single GetImage RPC, so they describe the +exact same sensor readings. + +Skipped unless a robot is reachable. To run: + + SPOT_ROBOT_IP=192.168.80.3 SPOT_USERNAME=user SPOT_PASSWORD=pass \ + pytest tests/test_depth_registration_robot.py -v -s +""" + +import os + +import numpy as np +import pytest +from bosdyn.api import image_pb2 # type: ignore[import-untyped] +from bosdyn.client.auth import InvalidLoginError # type: ignore[import-untyped] +from bosdyn.client.exceptions import ( # type: ignore[import-untyped] + ResponseError, + RpcError, +) +from bosdyn.client.image import ( # type: ignore[import-untyped] + ImageClient, + UnknownImageSourceError, + build_image_request, +) +from pyspotobserver.depth_registration import extract_registration_params, register_depth + +ROBOT_IP = os.environ.get("SPOT_ROBOT_IP") +USERNAME = os.environ.get("SPOT_USERNAME") +PASSWORD = os.environ.get("SPOT_PASSWORD") + +_REQUIRED_ENV_VARS = ("SPOT_ROBOT_IP", "SPOT_USERNAME", "SPOT_PASSWORD") +_MISSING_ENV_VARS = [name for name in _REQUIRED_ENV_VARS if not os.environ.get(name)] + +requires_robot = pytest.mark.skipif( + bool(_MISSING_ENV_VARS), + reason=( + f"Robot validation skipped: missing environment variable(s) " + f"{', '.join(_MISSING_ENV_VARS)}. Run with e.g.: " + f"SPOT_ROBOT_IP=192.168.80.3 SPOT_USERNAME=user SPOT_PASSWORD=pass " + f"pytest tests/test_depth_registration_robot.py -v -s" + ), +) + +# camera -> (rgb source, raw depth source, robot-registered reference source) +_SOURCES = { + "back": ("back_fisheye_image", "back_depth", "back_depth_in_visual_frame"), + "frontleft": ("frontleft_fisheye_image", "frontleft_depth", "frontleft_depth_in_visual_frame"), + "frontright": ( + "frontright_fisheye_image", + "frontright_depth", + "frontright_depth_in_visual_frame", + ), + "left": ("left_fisheye_image", "left_depth", "left_depth_in_visual_frame"), + "right": ("right_fisheye_image", "right_depth", "right_depth_in_visual_frame"), + "hand": ("hand_color_image", "hand_depth", "hand_depth_in_hand_color_frame"), +} + + +@pytest.fixture(scope="module") +def image_client() -> ImageClient: + import bosdyn.client # type: ignore[import-untyped] + + sdk = bosdyn.client.create_standard_sdk("DepthRegistrationValidation") + robot = sdk.create_robot(ROBOT_IP) + try: + robot.authenticate(USERNAME, PASSWORD) + except InvalidLoginError: + pytest.fail( + f"Authentication failed for user {USERNAME!r} on robot {ROBOT_IP} — " + f"check SPOT_USERNAME and SPOT_PASSWORD" + ) + except RpcError as exc: + pytest.fail( + f"Could not reach robot at {ROBOT_IP} ({exc}) — " + f"check SPOT_ROBOT_IP and that the robot is powered on and on the network" + ) + return robot.ensure_client(ImageClient.default_service_name) + + +def _decode_depth(response: image_pb2.ImageResponse) -> np.ndarray: + img = response.shot.image + src_name = response.source.name + if img.pixel_format != image_pb2.Image.PIXEL_FORMAT_DEPTH_U16: + raise AssertionError(f"{src_name}: unexpected depth pixel format: {img.pixel_format}") + data = np.frombuffer(img.data, dtype=np.uint16) + if data.size != img.rows * img.cols: + raise AssertionError( + f"{src_name}: depth payload has {data.size} pixels, " + f"expected {img.rows}x{img.cols} = {img.rows * img.cols}" + ) + depth_scale = response.source.depth_scale if response.source.depth_scale > 0 else 1.0 + return data.reshape(img.rows, img.cols).astype(np.float32) / depth_scale + + +@requires_robot +@pytest.mark.parametrize("camera", list(_SOURCES)) +def test_registration_matches_robot_registered_depth( + image_client: ImageClient, camera: str +) -> None: + rgb_src, raw_src, ref_src = _SOURCES[camera] + requests = [ + build_image_request( + rgb_src, + quality_percent=100.0, + image_format=image_pb2.Image.FORMAT_JPEG, + pixel_format=image_pb2.Image.PIXEL_FORMAT_RGB_U8, + ), + build_image_request( + raw_src, + image_format=image_pb2.Image.FORMAT_RAW, + pixel_format=image_pb2.Image.PIXEL_FORMAT_DEPTH_U16, + ), + build_image_request( + ref_src, + image_format=image_pb2.Image.FORMAT_RAW, + pixel_format=image_pb2.Image.PIXEL_FORMAT_DEPTH_U16, + ), + ] + + try: + responses = image_client.get_image(requests) + except UnknownImageSourceError: + # Expected for cameras the robot doesn't have (e.g. hand without a gripper). + pytest.skip( + f"{camera}: source(s) not available on this robot: {rgb_src}, {raw_src}, {ref_src}" + ) + except ResponseError as exc: + pytest.fail(f"{camera}: robot rejected the image request: {exc}") + except RpcError as exc: + pytest.fail(f"{camera}: lost connection to robot at {ROBOT_IP} during request: {exc}") + + assert len(responses) == 3 + rgb_response, raw_response, ref_response = responses + + raw = _decode_depth(raw_response) + ref = _decode_depth(ref_response) + + ref_valid = ref > 0 + if ref_valid.sum() < 1000: + pytest.skip(f"{camera}: reference depth too sparse to validate ({ref_valid.sum()} px)") + + params = extract_registration_params(rgb_response, raw_response, dst_shape=ref.shape) + out = np.empty(ref.shape, dtype=np.float32) + register_depth(raw, params, out) + + out_valid = out > 0 + both_valid = out_valid & ref_valid + + # Alignment/coverage: valid regions should mostly coincide. Normalizing by the + # smaller region tolerates density differences between the two splats. + overlap = both_valid.sum() / min(out_valid.sum(), ref_valid.sum()) + + # Agreement where both have data: same sensor readings, so depths should + # match closely except at occlusion boundaries. + abs_diff = np.abs(out[both_valid] - ref[both_valid]) + median_abs = float(np.median(abs_diff)) + rel_close = abs_diff < np.maximum(0.1, 0.05 * ref[both_valid]) + frac_close = float(np.mean(rel_close)) + + print( + f"[{camera}] raw_valid={int((raw > 0).sum())} ref_valid={int(ref_valid.sum())} " + f"out_valid={int(out_valid.sum())} overlap={overlap:.3f} " + f"median_abs={median_abs * 100:.2f}cm frac_close={frac_close:.3f}" + ) + + assert overlap > 0.6, f"{camera}: valid-pixel overlap {overlap:.3f} suggests misalignment" + assert median_abs < 0.05, f"{camera}: median depth error {median_abs:.3f}m too high" + assert frac_close > 0.85, f"{camera}: only {frac_close:.1%} of pixels within tolerance" diff --git a/README.md b/README.md new file mode 100644 index 0000000..fdcb992 --- /dev/null +++ b/README.md @@ -0,0 +1,105 @@ +# SpotObserver + +SpotObserver is a C++/CUDA library for streaming and processing camera data from +[Boston Dynamics Spot](https://bostondynamics.com/products/spot/) robots. It is built +as a native Unity plugin (`SpotObserverLib.dll`) and exposes a C API for: + +- Connecting to one or more Spot robots and streaming RGB + depth feeds from any of the + onboard cameras (back, front-left, front-right, left, right, hand) +- A GPU vision pipeline (for depth completion at the moment) powered by LibTorch and ONNX Runtime +- Low-overhead handoff of images into Unity via CUDA/DX12 interop. +- Per-frame body-to-world transforms alongside the images + +The full C API is declared in [`include/spot-observer.h`](include/spot-observer.h). + +## Python Implementation (PySpotObserver) + +A Pythonic implementation of the same functionality lives in +[`PySpotObserver/`](PySpotObserver/) — see its [README](PySpotObserver/README.md) and +[QUICKSTART](PySpotObserver/QUICKSTART.md). It provides a clean, type-hinted API with +sync and async usage, context-managed connections, YAML configuration, and an optional +ONNX vision pipeline: + +```python +from pyspotobserver import SpotConfig, SpotConnection, CameraType + +with SpotConnection(SpotConfig(robot_ip="192.168.80.3", username="...", password="...")) as conn: + stream = conn.create_cam_stream() + stream.start_streaming(CameraType.FRONTLEFT | CameraType.FRONTRIGHT) + rgb_images, depth_images, body_T_worlds = stream.get_current_images() +``` + +## Repository layout + +| Path | Contents | +| --- | --- | +| `include/` | Public C API header (`spot-observer.h`) | +| `src/` | Library implementation (connection, vision pipeline, CUDA kernels, Unity/DX12 interop) | +| `PySpotObserver/` | Python bindings and examples | +| `extern/` | Third-party dependencies (Spot C++ SDK, OpenCV, LibTorch, ONNX Runtime) | +| `tests/` | Integration tests (`integ-test`, `integ-test-dx12`) and standalone tests | +| `scripts/` | Utility scripts (e.g. timing statistics) | + +## Building SpotObserver + +### Prerequisites + +- Windows 10/11 x64 with Visual Studio 2022 (MSVC v143 toolset) +- CMake **3.28+** +- NVIDIA CUDA Toolkit (currently only tested on Ampere and Ada GPUs, compute 8.6 / 8.9) +- A Unity Editor installation (the CUDA/DX12 interop module uses Unity's Native Plugin API headers; + the build auto-detects the newest Unity Hub install — override with + `-DUnityPluginAPI_INCLUDE=/Editor/Data/PluginAPI` if needed) + +### 1. Set up dependencies + +All third-party dependencies live under `extern/`. Follow +[`extern/README.md`](extern/README.md) for the detailed steps. + +### 2. Configure and build + +Use `CMAKE_PREFIX_PATH` (pointing at the vcpkg installed tree), **not** +`CMAKE_TOOLCHAIN_FILE`: + +```bash +cmake -B build \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_PREFIX_PATH="/installed/x64-windows" \ + -DCMAKE_FIND_PACKAGE_PREFER_CONFIG=TRUE +cmake --build build --config Release +``` + +This produces `SpotObserverLib.dll` and automatically copies all required runtime DLLs +(vcpkg, OpenCV, LibTorch, ONNX Runtime) next to the built binaries. Set +`-DCOPY_DLLS=OFF` to disable the automatic copying. + +Useful options: + +| Option | Default | Effect | +| --- | --- | --- | +| `BUILD_TESTS` | `ON` | Build the test executables under `tests/` | +| `INSTALL_TESTS` | `ON` | Install test executables | +| `COPY_DLLS` | `ON` | Copy required DLLs next to the built library | + +### 3. Install (optional) + +```bash +cmake --install build --config Release +``` + +By default this installs the library, public headers, and runtime DLLs into +`/install`. Override with `-DCMAKE_INSTALL_PREFIX=` at configure time. + +## Tests + +With `BUILD_TESTS=ON` (the default), the build also produces: + +- `tests/integ-test` — end-to-end streaming test against a real robot +- `tests/integ-test-dx12` — DirectX 12 interop / readback test +- `tests/standalone-tests` — smaller isolated tests + +## Notes + +- Robot credentials are passed to `SOb_ConnectToSpot()` at runtime; nothing is + hardcoded in the library. +- The Boston Dynamics C++ SDK is still in beta \ No newline at end of file diff --git a/src/cuda_kernels.cu b/src/cuda_kernels.cu index abbdccb..aff1cfc 100644 --- a/src/cuda_kernels.cu +++ b/src/cuda_kernels.cu @@ -2,6 +2,7 @@ #include "cuda_kernels.cuh" #include +#include #include #include #include @@ -923,6 +924,90 @@ cudaError_t preprocess_depth_image2( return err; } +// ---- Depth-to-RGB registration --------------------------------------------- +// Reprojects a raw depth image (depth sensor frame) into the RGB camera's image +// plane, producing the same layout as Spot's *_depth_in_visual_frame sources. + +__global__ void fill_float_kernel(float* __restrict__ buf, int n, float value) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n) buf[i] = value; +} + +__global__ void register_depth_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + DepthRegistrationParams p +) { + int u = blockIdx.x * blockDim.x + threadIdx.x; + int v = blockIdx.y * blockDim.y + threadIdx.y; + if (u >= p.src_width || v >= p.src_height) return; + + float z = src[v * p.src_width + u]; + if (z <= 0.0f) return; // invalid depth sample + + float uf = float(u), vf = float(v); + float px = z * (p.M[0] * uf + p.M[1] * vf + p.M[2]) + p.t[0]; + float py = z * (p.M[3] * uf + p.M[4] * vf + p.M[5]) + p.t[1]; + float pz = z * (p.M[6] * uf + p.M[7] * vf + p.M[8]) + p.t[2]; + if (pz <= 0.0f) return; // behind the RGB camera + + float inv_z = 1.0f / pz; + int u0 = __float2int_rd(px * inv_z); + int v0 = __float2int_rd(py * inv_z); + + // 2x2 splat: the depth image is lower resolution than the RGB target, so + // covering the four nearest pixels reduces holes without a dilation pass. + #pragma unroll + for (int dv = 0; dv <= 1; ++dv) { + #pragma unroll + for (int du = 0; du <= 1; ++du) { + int x = u0 + du; + int y = v0 + dv; + if (x < 0 || y < 0 || x >= p.dst_width || y >= p.dst_height) continue; + // atomicMin on the raw bits keeps the nearest surface when several + // depth pixels land on one RGB pixel; ordering is correct because + // positive IEEE-754 floats compare like unsigned ints. + atomicMin(reinterpret_cast(&dst[y * p.dst_width + x]), + __float_as_uint(pz)); + } + } +} + +__global__ void finalize_registered_depth_kernel(float* __restrict__ dst, int n, float sentinel) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < n && dst[i] == sentinel) dst[i] = 0.0f; +} + +cudaError_t register_depth_to_rgb( + const float* d_depth_in, + float* d_depth_out, + const DepthRegistrationParams& params, + cudaStream_t stream +) { + const int n_dst = params.dst_width * params.dst_height; + if (n_dst <= 0 || params.src_width <= 0 || params.src_height <= 0) { + return cudaErrorInvalidValue; + } + constexpr float kSentinel = FLT_MAX; + constexpr int threads = 256; + + fill_float_kernel<<<(n_dst + threads - 1) / threads, threads, 0, stream>>>( + d_depth_out, n_dst, kSentinel); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) return err; + + dim3 block(32, 8); + dim3 grid((params.src_width + block.x - 1) / block.x, + (params.src_height + block.y - 1) / block.y); + register_depth_kernel<<>>(d_depth_in, d_depth_out, params); + err = cudaGetLastError(); + if (err != cudaSuccess) return err; + + finalize_registered_depth_kernel<<<(n_dst + threads - 1) / threads, threads, 0, stream>>>( + d_depth_out, n_dst, kSentinel); + return cudaGetLastError(); +} + // Dummy kernel for testing purposes (keeping original) __global__ void setOutputToOnes(float* input, float* output, int size) { int idx = blockIdx.x * blockDim.x + threadIdx.x; diff --git a/src/include/cuda_kernels.cuh b/src/include/cuda_kernels.cuh index 8021ccc..13c75f8 100644 --- a/src/include/cuda_kernels.cuh +++ b/src/include/cuda_kernels.cuh @@ -8,6 +8,29 @@ namespace SOb { cudaError_t setOutputToOnes_launcher(float* d_input, float* d_output, int size); +// Parameters for reprojecting a raw depth image into an RGB camera's image plane +// (client-side equivalent of Spot's *_depth_in_visual_frame sources). +// Folded projection, built on the CPU from intrinsics + extrinsics: +// p = z * M * [u, v, 1]^T + t, with M = K_rgb * R * K_depth^-1 and t = K_rgb * translation +// where (u, v, z) is a depth pixel and (p.x/p.z, p.y/p.z, p.z) is its RGB-frame pixel + depth. +struct DepthRegistrationParams { + float M[9]; // Row-major 3x3 + float t[3]; + int src_width{0}; + int src_height{0}; + int dst_width{0}; + int dst_height{0}; +}; + +// Reprojects d_depth_in (src-sized, meters) into d_depth_out (dst-sized, meters, +// fully overwritten; pixels with no depth sample are 0). +cudaError_t register_depth_to_rgb( + const float* d_depth_in, + float* d_depth_out, + const DepthRegistrationParams& params, + cudaStream_t stream = 0 +); + // size_t depth_preprocessor_get_workspace_size(int width, int height); // cudaError_t preprocess_depth_image( // float* depth_image, diff --git a/src/include/spot-connection.h b/src/include/spot-connection.h index 625bae7..567d468 100644 --- a/src/include/spot-connection.h +++ b/src/include/spot-connection.h @@ -6,6 +6,7 @@ #include "spot-observer.h" #include "model.h" +#include "cuda_kernels.cuh" #include #include @@ -41,6 +42,13 @@ class ReaderWriterCBuf { float* depth_data_{nullptr}; float* cached_depth_{nullptr}; + // Client-side depth registration: per-camera reprojection params. Raw depth + // images are staged in raw_depth_scratch_ (CUDA memory) and registered into + // the depth ring at RGB resolution. + std::vector depth_registration_; + size_t n_elems_per_raw_depth_{0}; + float* raw_depth_scratch_{nullptr}; + bool first_run_{true}; const size_t max_size_; // Maximum size of the queue @@ -60,7 +68,9 @@ class ReaderWriterCBuf { size_t n_bytes_per_rgb, size_t n_bytes_per_depth, std::vector cameras, - std::vector camera_dump_subdirs + std::vector camera_dump_subdirs, + std::vector depth_registration, + size_t n_elems_per_raw_depth ); /** @@ -113,6 +123,8 @@ class SpotCamStream { TensorShape current_rgb_shape_{0, 0, 0, 0}; TensorShape current_depth_shape_{0, 0, 0, 0}; + TimingInfo timing_info_; + bosdyn::api::GetImageRequest _createImageRequest( const std::vector& rgb_sources, const std::vector& depth_sources diff --git a/src/include/vision-pipeline.h b/src/include/vision-pipeline.h index 756f648..6ebf86b 100644 --- a/src/include/vision-pipeline.h +++ b/src/include/vision-pipeline.h @@ -10,10 +10,6 @@ #include #include #include -#include -#include -#include -#include #include #include diff --git a/src/spot-connection.cpp b/src/spot-connection.cpp index d9811cd..acf88d3 100644 --- a/src/spot-connection.cpp +++ b/src/spot-connection.cpp @@ -12,6 +12,7 @@ #include #include +#include #include #include @@ -89,6 +90,68 @@ static std::array se3_pose_to_row_major_matrix(const bosdyn::api::SE3 }; } +// Folds K_rgb * rgb_T_depth * K_depth^-1 into the single matrix + translation the +// registration kernel consumes (one multiply-add chain per depth pixel). +// Skew is assumed zero; Spot's body cameras report none. +static bool build_depth_registration_params( + const bosdyn::api::ImageSource& rgb_source, + const bosdyn::api::ImageSource& depth_source, + const bosdyn::api::SE3Pose& rgb_T_depth, + DepthRegistrationParams& out +) { + if (!rgb_source.has_pinhole() || !depth_source.has_pinhole()) { + LogMessage("build_depth_registration_params: Missing pinhole intrinsics " + "(rgb '{}': {}, depth '{}': {})", + rgb_source.name(), rgb_source.has_pinhole(), + depth_source.name(), depth_source.has_pinhole()); + return false; + } + + const auto& k_rgb = rgb_source.pinhole().intrinsics(); + const auto& k_depth = depth_source.pinhole().intrinsics(); + + const double fx_r = k_rgb.focal_length().x(); + const double fy_r = k_rgb.focal_length().y(); + const double cx_r = k_rgb.principal_point().x(); + const double cy_r = k_rgb.principal_point().y(); + const double fx_d = k_depth.focal_length().x(); + const double fy_d = k_depth.focal_length().y(); + const double cx_d = k_depth.principal_point().x(); + const double cy_d = k_depth.principal_point().y(); + + if (fx_r == 0.0 || fy_r == 0.0 || fx_d == 0.0 || fy_d == 0.0) { + LogMessage("build_depth_registration_params: Degenerate intrinsics " + "(rgb '{}', depth '{}')", rgb_source.name(), depth_source.name()); + return false; + } + + const std::array T = se3_pose_to_row_major_matrix(rgb_T_depth); + + // A = R * K_depth^-1 + double A[3][3]; + for (int i = 0; i < 3; ++i) { + const double r0 = T[i * 4 + 0]; + const double r1 = T[i * 4 + 1]; + const double r2 = T[i * 4 + 2]; + A[i][0] = r0 / fx_d; + A[i][1] = r1 / fy_d; + A[i][2] = -r0 * cx_d / fx_d - r1 * cy_d / fy_d + r2; + } + const double tx = T[3], ty = T[7], tz = T[11]; + + // M = K_rgb * A, t = K_rgb * translation + for (int j = 0; j < 3; ++j) { + out.M[0 + j] = static_cast(fx_r * A[0][j] + cx_r * A[2][j]); + out.M[3 + j] = static_cast(fy_r * A[1][j] + cy_r * A[2][j]); + out.M[6 + j] = static_cast(A[2][j]); + } + out.t[0] = static_cast(fx_r * tx + cx_r * tz); + out.t[1] = static_cast(fy_r * ty + cy_r * tz); + out.t[2] = static_cast(tz); + + return true; +} + static cv::Mat convert_image_to_cv_mat(const bosdyn::api::Image& img, double depth_scale = 1.0) { if (img.format() == bosdyn::api::Image::FORMAT_JPEG) { // Decode JPEG data @@ -159,6 +222,10 @@ ReaderWriterCBuf::~ReaderWriterCBuf() { cudaFree(cached_depth_); cached_depth_ = nullptr; } + if (raw_depth_scratch_) { + cudaFree(raw_depth_scratch_); + raw_depth_scratch_ = nullptr; + } // Do not destroy cuda_stream_ here; SpotConnection owns it. LogMessage("Destroyed ReaderWriterCBuf"); } @@ -167,19 +234,29 @@ bool ReaderWriterCBuf::initialize( size_t n_elems_per_rgb, size_t n_elems_per_depth, std::vector cameras, - std::vector camera_dump_subdirs + std::vector camera_dump_subdirs, + std::vector depth_registration, + size_t n_elems_per_raw_depth ) { n_elems_per_rgb_ = n_elems_per_rgb; n_elems_per_depth_ = n_elems_per_depth; n_images_per_response_ = cameras.size(); cameras_ = std::move(cameras); camera_dump_subdirs_ = std::move(camera_dump_subdirs); + depth_registration_ = std::move(depth_registration); + n_elems_per_raw_depth_ = n_elems_per_raw_depth; if (camera_dump_subdirs_.size() != cameras_.size()) { LogMessage("ReaderWriterCBuf::initialize: Expected {} camera dump subdirs, got {}", cameras_.size(), camera_dump_subdirs_.size()); return false; } + if (depth_registration_.size() != cameras_.size() || n_elems_per_raw_depth_ == 0) { + LogMessage("ReaderWriterCBuf::initialize: Expected {} depth registration params " + "(got {}) and a nonzero raw depth size (got {})", + cameras_.size(), depth_registration_.size(), n_elems_per_raw_depth_); + return false; + } if (rgb_data_ != nullptr || depth_data_ != nullptr) { if (!rgb_data_) checkCudaError(cudaFree(rgb_data_), "cudaFree for RGB data"); @@ -199,6 +276,15 @@ bool ReaderWriterCBuf::initialize( checkCudaError(cudaMalloc(&depth_data_, total_size_depth), "cudaMalloc for Depth data"); checkCudaError(cudaMalloc(&cached_depth_, size_depth_per_response), "cudaMalloc for Cached Depth data"); + if (raw_depth_scratch_) { + checkCudaError(cudaFree(raw_depth_scratch_), "cudaFree for raw depth scratch"); + raw_depth_scratch_ = nullptr; + } + checkCudaError( + cudaMalloc(&raw_depth_scratch_, n_elems_per_raw_depth_ * sizeof(float)), + "cudaMalloc for raw depth scratch" + ); + LogMessage("Allocated {} bytes for RGB data and {} bytes for Depth data in ReaderWriterCBuf" "({} bytes per RGB, {} bytes per Depth, {} images per response, max size {})", total_size_rgb, total_size_depth, n_elems_per_rgb * sizeof(uint8_t), n_elems_per_depth * sizeof(float), n_images_per_response_, max_size_); @@ -342,18 +428,25 @@ void ReaderWriterCBuf::push(const google::protobuf::RepeatedPtrField get_depth_cam_names_from_bit_mask(uint32_t bitma std::vector cam_names; cam_names.reserve(__num_set_bits(bitmask)); - if (bitmask & SpotCamera::BACK) cam_names.emplace_back("back_depth_in_visual_frame"); - if (bitmask & SpotCamera::FRONTLEFT) cam_names.emplace_back("frontleft_depth_in_visual_frame"); - if (bitmask & SpotCamera::FRONTRIGHT) cam_names.emplace_back("frontright_depth_in_visual_frame"); - if (bitmask & SpotCamera::LEFT) cam_names.emplace_back("left_depth_in_visual_frame"); - if (bitmask & SpotCamera::RIGHT) cam_names.emplace_back("right_depth_in_visual_frame"); - if (bitmask & SpotCamera::HAND) cam_names.emplace_back("hand_depth_in_hand_color_frame"); + // 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 on the GPU (see register_depth_to_rgb). + if (bitmask & SpotCamera::BACK) cam_names.emplace_back("back_depth"); + if (bitmask & SpotCamera::FRONTLEFT) cam_names.emplace_back("frontleft_depth"); + if (bitmask & SpotCamera::FRONTRIGHT) cam_names.emplace_back("frontright_depth"); + if (bitmask & SpotCamera::LEFT) cam_names.emplace_back("left_depth"); + if (bitmask & SpotCamera::RIGHT) cam_names.emplace_back("right_depth"); + if (bitmask & SpotCamera::HAND) cam_names.emplace_back("hand_depth"); return cam_names; } @@ -561,7 +670,7 @@ bosdyn::api::GetImageRequest SpotCamStream::_createImageRequest( bosdyn::api::ImageRequest* image_request = request.add_image_requests(); image_request->set_image_source_name(source); image_request->set_quality_percent(100.0); - // TODO: Use FORMAT_RLE format for depth (compresses 0s) + // FORMAT_RLE is rejected by the robot (STATUS_UNSUPPORTED_IMAGE_FORMAT_REQUESTED). image_request->set_image_format(bosdyn::api::Image_Format_FORMAT_RAW); image_request->set_pixel_format(bosdyn::api::Image::PIXEL_FORMAT_DEPTH_U16); } @@ -752,23 +861,71 @@ bool SpotCamStream::streamCameras(uint32_t cam_mask) { return false; } } - // Same thing for depth images + // Same thing for the raw depth images + size_t raw_depth_ref_size = image_responses[num_cams_requested].shot().image().rows() + * image_responses[num_cams_requested].shot().image().cols(); + for (int32_t j = num_cams_requested + 1; j < image_responses.size(); j++) { + const auto& img_response = image_responses[j]; + size_t depth_size = img_response.shot().image().rows() * img_response.shot().image().cols(); + if (raw_depth_ref_size != depth_size) { + LogMessage("SpotCamStream::streamCameras: Inconsistent depth image sizes" + "(expected {}, got {})", raw_depth_ref_size, depth_size); + return false; + } + } + + // Raw depth is registered into the RGB frame on the GPU, so the depth + // ring buffer (and everything downstream) is RGB-sized. current_depth_shape_ = TensorShape{ size_t(num_cams_requested), 1, - size_t(image_responses[num_cams_requested].shot().image().rows()), - size_t(image_responses[num_cams_requested].shot().image().cols()) + current_rgb_shape_.H, + current_rgb_shape_.W }; - size_t depth_ref_size = current_depth_shape_.H * current_depth_shape_.W; - for (int32_t j = num_cams_requested + 1; j < image_responses.size(); j++) { - const auto& img_response = image_responses[j]; - size_t depth_size = img_response.shot().image().rows() * img_response.shot().image().cols(); - if (depth_ref_size != depth_size) { - LogMessage("SpotCamStream::streamCameras: Inconsistent depth image sizes" - "(expected {}, got {})", depth_ref_size, depth_size); + + // Build the per-camera depth -> RGB reprojection from the intrinsics and + // the frame tree. Both are fixed by the mechanical mounting, so this is + // computed once per stream. + std::vector depth_registration; + depth_registration.reserve(num_cams_requested); + + for (int32_t j = 0; j < num_cams_requested; j++) { + const auto& rgb_response = image_responses[j]; + const auto& depth_response = image_responses[num_cams_requested + j]; + const std::string& depth_frame = depth_response.shot().frame_name_image_sensor(); + + bosdyn::api::SE3Pose body_T_depth; + if (!bosdyn::api::get_a_tform_b( + depth_response.shot().transforms_snapshot(), + bosdyn::api::kBodyFrame, + depth_frame, + &body_T_depth)) { + LogMessage("SpotCamStream::streamCameras: Failed to get body_T_depth for frame {}", + depth_frame); + return false; + } + + const bosdyn::api::SE3Pose rgb_T_depth = ~body_T_cameras_[j] * body_T_depth; + + DepthRegistrationParams reg; + if (!build_depth_registration_params( + rgb_response.source(), depth_response.source(), rgb_T_depth, reg)) { + LogMessage("SpotCamStream::streamCameras: Failed to build depth registration " + "params for camera {}", get_dump_name_for_camera(camera_order_[j])); return false; } + reg.src_width = depth_response.shot().image().cols(); + reg.src_height = depth_response.shot().image().rows(); + reg.dst_width = static_cast(current_rgb_shape_.W); + reg.dst_height = static_cast(current_rgb_shape_.H); + depth_registration.push_back(reg); + + LogMessage("Depth registration for camera {}: {}x{} ({}) -> {}x{} ({})", + get_dump_name_for_camera(camera_order_[j]), + reg.src_width, reg.src_height, depth_frame, + reg.dst_width, reg.dst_height, + rgb_response.shot().frame_name_image_sensor()); } // (Re)initialize circular buffer @@ -781,7 +938,9 @@ bool SpotCamStream::streamCameras(uint32_t cam_mask) { rgb_ref_size, depth_ref_size, camera_order_, - std::move(camera_dump_subdirs) + std::move(camera_dump_subdirs), + std::move(depth_registration), + raw_depth_ref_size )) { LogMessage("SpotCamStream::streamCameras: Failed to initialize image buffer"); return false; diff --git a/tests/integ-test/integ-test.cpp b/tests/integ-test/integ-test.cpp index d346676..c0d04d1 100644 --- a/tests/integ-test/integ-test.cpp +++ b/tests/integ-test/integ-test.cpp @@ -66,13 +66,13 @@ int main(int argc, char* argv[]) { std::string password = argv[4]; //SOb_ToggleDebugDumps("./spot_dump"); - SOb_ToggleLogging(false); + SOb_SetLogLevel(1); int32_t spot_ids[2] = {-1, -1}; std::unordered_map> cam_stream_ids; - std::vector cam_bitmasks = {FRONTRIGHT | FRONTLEFT};//, HAND }; + std::vector cam_bitmasks = {FRONTRIGHT | FRONTLEFT, HAND }; for (size_t i = 0; i < 2; i++) { if (robot_ips[i] == "0") { spot_ids[i] = -1; From 963da72f4220e7589d0a60c69837f02b75c33b34 Mon Sep 17 00:00:00 2001 From: Faisal Zaghloul Date: Fri, 10 Jul 2026 16:31:58 -0400 Subject: [PATCH 2/6] Enable naive EMA depth averaging --- include/spot-observer.h | 5 +- src/cuda_kernels.cu | 124 +++++++++++++++---------------- src/include/cuda_kernels.cuh | 25 +++---- src/include/unity-cuda-interop.h | 22 +----- src/include/vision-pipeline.h | 7 ++ src/spot-observer.cpp | 25 ++++++- src/unity-cuda-interop.cpp | 19 +++++ src/vision-pipeline.cpp | 52 ++++++++----- 8 files changed, 155 insertions(+), 124 deletions(-) diff --git a/include/spot-observer.h b/include/spot-observer.h index 667d3f2..634ff06 100644 --- a/include/spot-observer.h +++ b/include/spot-observer.h @@ -102,8 +102,11 @@ void UNITY_INTERFACE_API SOb_UnloadModel(SObModel model); // Config calls UNITY_INTERFACE_EXPORT bool UNITY_INTERFACE_API SOb_ToggleDepthCompletion(bool enable); +// Enable/disable EMA depth averaging for the vision pipeline running on the +// given camera stream (default: enabled). Takes effect on the next processed +// frame. Fails if no vision pipeline exists for the given ids. UNITY_INTERFACE_EXPORT -bool UNITY_INTERFACE_API SOb_ToggleDepthAveraging(bool enable); +bool UNITY_INTERFACE_API SOb_SetDepthAveraging(int32_t robot_id, int32_t cam_stream_id, bool enable); UNITY_INTERFACE_EXPORT bool UNITY_INTERFACE_API SOb_ToggleDepthAveragingWithOpticalFlow(bool enable); UNITY_INTERFACE_EXPORT diff --git a/src/cuda_kernels.cu b/src/cuda_kernels.cu index aff1cfc..47b7db4 100644 --- a/src/cuda_kernels.cu +++ b/src/cuda_kernels.cu @@ -5,8 +5,8 @@ #include #include #include -#include -#include +//#include +//#include #include #include @@ -84,8 +84,8 @@ __global__ void downscale_optimized_kernel( int original_height ) { // Shared memory: input tile + intermediate results for separable convolution - __shared__ float s_input[TILE_SZ * SCALE_FACTOR][TILE_SZ * SCALE_FACTOR + 1]; - __shared__ float s_horizontal[TILE_SZ][TILE_SZ * SCALE_FACTOR + 1]; + __shared__ float s_input[TILE_SZ * SCALE_FACTOR][TILE_SZ * SCALE_FACTOR]; + __shared__ float s_horizontal[TILE_SZ][TILE_SZ * SCALE_FACTOR]; // Handle rotation: if rotating 90 CW, output dimensions are swapped const int new_width = ROTATE_90_CW ? (original_height / SCALE_FACTOR) : (original_width / SCALE_FACTOR); @@ -649,39 +649,6 @@ __device__ __forceinline__ float fast_rcp(float x) { // return cudaGetLastError(); // } -cudaError_t postprocess_depth_image( - float* depth_image, - int width, - int height, - float* workspace, - bool rotate_90_ccw, - cudaStream_t stream -) { - if (!rotate_90_ccw) { - return cudaSuccess; // No rotation needed - } - - // Apply 270 CW rotation - dim3 block(TILE_SIZE, TILE_SIZE); - dim3 grid((width + TILE_SIZE - 1) / TILE_SIZE, - (height + TILE_SIZE - 1) / TILE_SIZE); - - // Rotate to workspace buffer - rotateDepth_fast_kernel<<>>( - depth_image, workspace, width, height); - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) return err; - - // Copy rotated result back - size_t image_size = width * height * sizeof(float); - err = cudaMemcpyAsync(depth_image, workspace, image_size, cudaMemcpyDeviceToDevice, stream); - if (err != cudaSuccess) return err; - - err = cudaGetLastError(); - return err; -} - struct int2_ { int x, y; }; __device__ __forceinline__ @@ -1366,15 +1333,16 @@ cudaError_t prefill_invalid_depth( return cudaGetLastError(); } -// CUDA kernel for maintaining running average of depth values -// new_depth: current frame depth values -// avg_depth: running average depth buffer (input/output) -// width, height: image dimensions - -__global__ void update_depth_cache_kernel( - const float* __restrict__ generated_depth, +// Fused output postprocess + EMA depth-cache update. One pass over the +// image: reads the generated depth, folds it into the running cache, and +// writes the blended result out as the published depth. +// generated_depth and output_depth may alias (each thread reads and writes +// the same index), so those two pointers must not be __restrict__. +__global__ void postprocess_depth_image( + const float* generated_depth, const float* __restrict__ sparse_depth, float* __restrict__ cached_depth, + float* output_depth, float alpha_valid, float alpha_invalid, int width, @@ -1384,55 +1352,79 @@ __global__ void update_depth_cache_kernel( ) { int x = blockIdx.x * blockDim.x + threadIdx.x; int y = blockIdx.y * blockDim.y + threadIdx.y; - + if (x >= width || y >= height) return; - + int idx = y * width + x; float generated_val = generated_depth[idx]; - float new_val = sparse_depth[idx]; + //float new_val = sparse_depth[idx]; float old_val = cached_depth[idx]; - + // Check if new depth value is valid - bool new_valid = (new_val >= min_valid_depth && new_val <= max_valid_depth); + //bool new_valid = (new_val >= min_valid_depth && new_val <= max_valid_depth); bool old_valid = (old_val >= min_valid_depth && old_val <= max_valid_depth); - - if (new_valid) { - if (old_valid) { - cached_depth[idx] = old_val * (1.f-alpha_valid) + new_val * alpha_valid; - } else { - cached_depth[idx] = new_val; - } - } else if (old_valid) { - cached_depth[idx] = old_val * (1.f-alpha_invalid) + generated_val * alpha_invalid; + + float result; + // if (new_valid) { + // if (old_valid) { + // result = old_val * (1.f-alpha_valid) + new_val * alpha_valid; + // } else { + // result = new_val; + // } + // } else + if (old_valid) { + result = old_val * (1.f-alpha_invalid) + generated_val * alpha_invalid; } else { - cached_depth[idx] = generated_val; + result = generated_val; } + + cached_depth[idx] = result; + output_depth[idx] = result; } -// Host function wrapper for depth cache update -cudaError_t update_depth_cache( - const float* generated_depth, +// Host wrapper: optionally rotates the generated depth back into the sensor +// orientation (via workspace), then runs the fused EMA update, writing the +// blended result back over generated_depth as the published output. +cudaError_t postprocess_depth_image( + float* generated_depth, const float* sparse_depth, float* cached_depth, float alpha_valid, float alpha_invalid, int width, int height, + float* workspace, + bool rotate_90_ccw, float min_valid_depth, float max_valid_depth, cudaStream_t stream ) { + const float* d_generated = generated_depth; + if (rotate_90_ccw) { + // Rotate into the workspace; the fused kernel then reads from there + // and can safely overwrite generated_depth with the final result. + dim3 rot_block(TILE_SIZE, TILE_SIZE); + dim3 rot_grid((width + TILE_SIZE - 1) / TILE_SIZE, + (height + TILE_SIZE - 1) / TILE_SIZE); + rotateDepth_fast_kernel<<>>( + generated_depth, workspace, width, height); + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) return err; + d_generated = workspace; + } + dim3 block(32, 8); dim3 grid((width + block.x - 1) / block.x, (height + block.y - 1) / block.y); - - update_depth_cache_kernel<<>>( - generated_depth, sparse_depth, cached_depth, + + postprocess_depth_image<<>>( + d_generated, sparse_depth, cached_depth, generated_depth, alpha_valid, alpha_invalid, width, height, min_valid_depth, max_valid_depth ); - + return cudaGetLastError(); } diff --git a/src/include/cuda_kernels.cuh b/src/include/cuda_kernels.cuh index 13c75f8..85049d8 100644 --- a/src/include/cuda_kernels.cuh +++ b/src/include/cuda_kernels.cuh @@ -54,15 +54,6 @@ cudaError_t preprocess_depth_image2( cudaStream_t stream = 0 ); -cudaError_t postprocess_depth_image( - float* depth_image, - int width, - int height, - float* workspace, - bool rotate_90_ccw = false, - cudaStream_t stream = 0 -); - void convert_uint8_img_to_float_img( const uint8_t* d_in, // [N,H,W,4] or [N,H,W,3] float* d_out, // [N,3,H,W] or [N,3,W,H] if rotated @@ -85,14 +76,20 @@ cudaError_t prefill_invalid_depth( float max_valid_depth = 100.0f, cudaStream_t stream = 0 ); -cudaError_t update_depth_cache( - const float* generated_depth, +// Fused output postprocess + EMA cache update: optionally rotates the +// generated depth back into the sensor orientation (via workspace), folds it +// into the running cache, and writes the blended result back over +// generated_depth as the published output — one pass, no extra copies. +cudaError_t postprocess_depth_image( + float* generated_depth, // Input: model output; overwritten with the blended result const float* sparse_depth, - float* cached_depth, // Input/output: running average buffer - float alpha_valid, // If old and new depth are valid, use this alpha - float alpha_invalid, // If old depth is valid but new depth is invalid, use this alpha (between old depth and generated depth) + float* cached_depth, // Input/output: running average buffer + float alpha_valid, // If old and new depth are valid, use this alpha + float alpha_invalid, // If old depth is valid but new depth is invalid, use this alpha (between old depth and generated depth) int width, int height, + float* workspace, // >= width*height floats; only used when rotating + bool rotate_90_ccw = false, float min_valid_depth = 0.01f, float max_valid_depth = 100.0f, cudaStream_t stream = 0 diff --git a/src/include/unity-cuda-interop.h b/src/include/unity-cuda-interop.h index 9c21e1d..fbe5e38 100644 --- a/src/include/unity-cuda-interop.h +++ b/src/include/unity-cuda-interop.h @@ -4,26 +4,8 @@ #pragma once -#include "logger.h" -#include "utils.h" - -#include -#include - -#ifdef _WIN32 -#include -#include -#include -#include - -#else -#error "Only Windows is supported for D3D12" -#endif - -// This ordering is required unfortunately -#include "spot-observer.h" -#include "IUnityGraphics.h" -#include "IUnityGraphicsD3D12.h" +#include +#include "IUnityInterface.h" namespace SOb { diff --git a/src/include/vision-pipeline.h b/src/include/vision-pipeline.h index 6ebf86b..c1dafb2 100644 --- a/src/include/vision-pipeline.h +++ b/src/include/vision-pipeline.h @@ -28,6 +28,7 @@ class VisionPipeline { MLModel& model_; const SpotCamStream& spot_cam_stream_; bool first_run_{true}; + std::atomic depth_averaging_enabled_{true}; // Threading std::unique_ptr pipeline_thread_; @@ -76,6 +77,12 @@ class VisionPipeline { void stop(); bool isRunning() const { return running_.load(); } + // Runtime toggle for EMA depth averaging (default: enabled). + // Takes effect on the next processed frame. + void setDepthAveraging(bool enable) { + depth_averaging_enabled_.store(enable, std::memory_order_relaxed); + } + bool getCurrentImages( int32_t n_images_requested, uint8_t** images, diff --git a/src/spot-observer.cpp b/src/spot-observer.cpp index 11280fa..3a34a41 100644 --- a/src/spot-observer.cpp +++ b/src/spot-observer.cpp @@ -207,6 +207,24 @@ static bool getNextImageSetFromVisionPipeline( } } +static bool setPipelineDepthAveraging(int32_t robot_id, int32_t cam_stream_id, bool enable) { + auto it = __robot_connections.find(robot_id); + if (it == __robot_connections.end()) { + LogMessage("SOb_SetDepthAveraging: Robot ID {} not found", robot_id); + return false; + } + + VisionPipeline* pipeline = it->second->getVisionPipeline(cam_stream_id); + if (!pipeline) { + LogMessage("SOb_SetDepthAveraging: Vision pipeline for camera stream ID {} not found for robot ID {}", + cam_stream_id, robot_id); + return false; + } + + pipeline->setDepthAveraging(enable); + return true; +} + static SObModel loadTorchModel(const std::string& modelPath, const std::string& backend) { LogMessage("Loading Torch model: {}", modelPath); LogMessage("Using backend: {}", backend); @@ -637,9 +655,10 @@ bool UNITY_INTERFACE_API SOb_ToggleDepthCompletion(bool enable) { } UNITY_INTERFACE_EXPORT -bool UNITY_INTERFACE_API SOb_ToggleDepthAveraging(bool enable) { - SOb::LogMessage("SOb_ToggleDepthAveraging called with enable: {}", enable); - return true; +bool UNITY_INTERFACE_API SOb_SetDepthAveraging(int32_t robot_id, int32_t cam_stream_id, bool enable) { + SOb::LogMessage("SOb_SetDepthAveraging called for robot ID {} @ stream-ID {} with enable: {}", + robot_id, cam_stream_id, enable); + return SOb::setPipelineDepthAveraging(robot_id, cam_stream_id, enable); } UNITY_INTERFACE_EXPORT diff --git a/src/unity-cuda-interop.cpp b/src/unity-cuda-interop.cpp index e69cd6a..494ba4e 100644 --- a/src/unity-cuda-interop.cpp +++ b/src/unity-cuda-interop.cpp @@ -2,14 +2,33 @@ // Created by fmz on 8/13/2025. // +#ifdef _WIN32 +#include +#include +#include +#include + +#else +#error "Only Windows is supported for D3D12" +#endif + +// This ordering is required unfortunately +#include "spot-observer.h" +#include "IUnityGraphics.h" +#include "IUnityGraphicsD3D12.h" + #include "unity-cuda-interop.h" #include "utils.h" #include "d3dx12.h" + #include #include #include #include +#include +#include + namespace SOb { // Unity interface diff --git a/src/vision-pipeline.cpp b/src/vision-pipeline.cpp index 971ead1..f4c1afc 100644 --- a/src/vision-pipeline.cpp +++ b/src/vision-pipeline.cpp @@ -110,6 +110,8 @@ bool VisionPipeline::allocateCudaBuffers() { checkCudaError(cudaMalloc(&cuda_ws_.d_depth_data_, depth_size), "cudaMalloc for vision pipeline input depth data"); checkCudaError(cudaMalloc(&cuda_ws_.d_preprocessed_depth_data_, depth_size), "cudaMalloc for vision pipeline input depth data preprocessed"); checkCudaError(cudaMalloc(&cuda_ws_.d_depth_cached_, depth_size), "cudaMalloc for vision pipeline input depth cached"); + // Zero = invalid depth, so the EMA cache starts empty and gets seeded on the first update. + checkCudaError(cudaMemset(cuda_ws_.d_depth_cached_, 0, depth_size), "cudaMemset for vision pipeline input depth cached"); size_t depth_workspace_size = depth_preprocessor2_get_workspace_size(depth_shape_.W, depth_shape_.H); checkCudaError(cudaMalloc(&cuda_ws_.d_depth_preprocessor_workspace_, depth_workspace_size), "cudaMalloc for vision pipeline depth preprocessor workspace"); @@ -214,6 +216,21 @@ void VisionPipeline::pipelineWorker(std::stop_token stop_token) { constexpr int32_t depth_scale_factor = 4; constexpr float inv_scale_factor = 1.0f / depth_scale_factor; + constexpr float min_valid_depth = 0.01f; + constexpr float max_valid_depth = 100.0f; + // EMA weights for the depth cache: weight of the new sample when the + // sensor depth is valid, and of the generated depth when it is not. + constexpr float ema_alpha_valid = 0.5f; + constexpr float ema_alpha_invalid = 0.5f; + + // Runtime EMA toggle, sampled once per iteration so a mid-frame + // toggle can't mix behaviors within one image set. When disabled, + // the fused postprocess runs as a pass-through (alpha 1 => output = + // generated depth) but still refreshes the cache, so re-enabling + // resumes blending from the latest frame with no stale data. + const bool ema_enabled = depth_averaging_enabled_.load(std::memory_order_relaxed); + const float alpha_invalid = ema_enabled ? ema_alpha_invalid : 1.0f; + // Copy inputs on the per-connection stream checkCudaError(cudaMemcpyAsync( d_rgb_ptr, @@ -285,15 +302,15 @@ void VisionPipeline::pipelineWorker(std::stop_token stop_token) { LogMessage("Starting pipeline for image {}. cur_rgb_ptr = {:#x}, cur_depth_ptr = {:#x}, cur_depth_output_ptr = {:#x}", i, size_t(cur_rgb_input_ptr), size_t(cur_depth_input_ptr), size_t(cur_depth_output_ptr)); - if (!first_run_) { + if (ema_enabled && !first_run_) { checkCudaError(prefill_invalid_depth( cur_depth_input_ptr, cur_preprocessed_depth_ptr, depth_cache_ptr, depth_shape_.W, depth_shape_.H, - 0.01f, - 100.0f, + min_valid_depth, + max_valid_depth, cuda_stream_ ), "prefill_invalid_depth"); @@ -361,7 +378,6 @@ void VisionPipeline::pipelineWorker(std::stop_token stop_token) { continue; } - auto inference_time = std::chrono::high_resolution_clock::now(); auto inference_duration = std::chrono::duration_cast(inference_time - preprocess_time); LogMessage("VisionPipeline inference time: {} ms", inference_duration.count()); @@ -375,29 +391,25 @@ void VisionPipeline::pipelineWorker(std::stop_token stop_token) { float* cur_depth_output_ptr = d_depth_output_ptr + i * output_shape_.C * output_shape_.H * output_shape_.W; float* depth_cache_ptr = cuda_ws_.d_depth_cached_ + i * depth_shape_.C * depth_shape_.H * depth_shape_.W; - // Postprocess output: rotate back if input was rotated + // Fused postprocess + EMA depth averaging: rotate the model output + // back if the input was rotated, fold it into the running cache, + // and publish the blended result as the pipeline output in a + // single pass — no intermediate copies. float* temp_output_ptr = reinterpret_cast(cuda_ws_.d_depth_preprocessor_workspace_); checkCudaError(postprocess_depth_image( cur_depth_output_ptr, + cur_depth_input_ptr, + depth_cache_ptr, + ema_alpha_valid, + alpha_invalid, output_shape_.W, output_shape_.H, temp_output_ptr, do_rotate_90_cw, + min_valid_depth, + max_valid_depth, cuda_stream_ - ), "postprocess_depth_image"); - // - // checkCudaError(update_depth_cache( - // cur_depth_output_ptr, - // cur_depth_input_ptr, - // depth_cache_ptr, - // 0.5, - // 0.1, - // output_shape_.W, - // output_shape_.H, - // 0.01f, - // 100.0f, - // cuda_stream_ - // ), "update_depth_cache"); + ), "postprocess_update_depth_cache"); // Ensure dumps see completed work (dumpers likely use default stream) checkCudaError(cudaStreamSynchronize(cuda_stream_), "sync before dumps"); @@ -446,7 +458,7 @@ void VisionPipeline::pipelineWorker(std::stop_token stop_token) { LogMessage("VisionPipeline: Updating write index from {} to {}", write_idx_, (write_idx_ + 1) % max_size_); write_idx_ = (write_idx_ + 1) % max_size_; - // first_run_ = false; + first_run_ = false; dump_id++; } catch (const std::exception& e) { From e4454e37819671f35b2f2a85c8c76c5f0073656e Mon Sep 17 00:00:00 2001 From: Faisal Zaghloul Date: Fri, 10 Jul 2026 17:42:48 -0400 Subject: [PATCH 3/6] Rework depth registration to use original u16 vals Avoid converting u16 depth data to floats on the CPU. This is now done on the GPU during the course of registering the depth to the RGB frame. Also addressed review comments. --- PySpotObserver/examples/basic_streaming.py | 1 + PySpotObserver/examples/config_example.yaml | 4 +- PySpotObserver/pyspotobserver/__init__.py | 10 +- .../pyspotobserver/color_correction.py | 1 + .../tests/test_depth_registration.py | 18 +-- src/cuda_kernels.cu | 109 ++++++++---------- src/dumper.cpp | 94 +++++++-------- src/include/cuda_kernels.cuh | 8 +- src/include/dumper.h | 3 +- src/include/spot-connection.h | 8 +- src/spot-connection.cpp | 51 +++++--- tests/integ-test/integ-test.cpp | 12 +- 12 files changed, 168 insertions(+), 151 deletions(-) diff --git a/PySpotObserver/examples/basic_streaming.py b/PySpotObserver/examples/basic_streaming.py index ba4e551..d227a22 100644 --- a/PySpotObserver/examples/basic_streaming.py +++ b/PySpotObserver/examples/basic_streaming.py @@ -26,6 +26,7 @@ build_config_from_args, parse_camera_list, ) + from pyspotobserver import CameraType, SpotConfig, SpotConnection logging.basicConfig( diff --git a/PySpotObserver/examples/config_example.yaml b/PySpotObserver/examples/config_example.yaml index e1f635a..bec26cb 100644 --- a/PySpotObserver/examples/config_example.yaml +++ b/PySpotObserver/examples/config_example.yaml @@ -4,8 +4,8 @@ # TUSKER: "128.148.138.22" # GOUGER: "128.148.138.21" robot_ip: "128.148.138.22" -username: "user" -password: "bigbubbabigbubba" +username: "" +password: "" # Streaming settings image_buffer_size: 5 diff --git a/PySpotObserver/pyspotobserver/__init__.py b/PySpotObserver/pyspotobserver/__init__.py index 1b124ae..5a0108a 100644 --- a/PySpotObserver/pyspotobserver/__init__.py +++ b/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", ] diff --git a/PySpotObserver/pyspotobserver/color_correction.py b/PySpotObserver/pyspotobserver/color_correction.py index 7bfd417..0a79efc 100644 --- a/PySpotObserver/pyspotobserver/color_correction.py +++ b/PySpotObserver/pyspotobserver/color_correction.py @@ -1,4 +1,5 @@ import numpy as np + from .config import CameraType # Known robot IPs diff --git a/PySpotObserver/tests/test_depth_registration.py b/PySpotObserver/tests/test_depth_registration.py index d9489fb..877568e 100644 --- a/PySpotObserver/tests/test_depth_registration.py +++ b/PySpotObserver/tests/test_depth_registration.py @@ -12,7 +12,9 @@ ) -def make_params(M, t, src_shape, dst_shape) -> DepthRegistrationParams: +def make_params( + M: np.ndarray, t: np.ndarray, src_shape: tuple[int, int], dst_shape: tuple[int, int] +) -> DepthRegistrationParams: """Build params directly from a projection matrix, bypassing proto extraction.""" v, u = np.indices(src_shape, dtype=np.float64) uv1 = np.stack([u, v, np.ones_like(u)], axis=-1) @@ -26,7 +28,7 @@ def make_params(M, t, src_shape, dst_shape) -> DepthRegistrationParams: class TestRegisterDepth: - def test_identity_maps_constant_plane(self): + def test_identity_maps_constant_plane(self) -> None: """With an identity projection, a constant plane fills the whole output.""" params = make_params(np.eye(3), np.zeros(3), src_shape=(6, 8), dst_shape=(6, 8)) raw = np.full((6, 8), 2.0, dtype=np.float32) @@ -36,7 +38,7 @@ def test_identity_maps_constant_plane(self): np.testing.assert_allclose(out, 2.0) - def test_nearest_surface_wins(self): + def test_nearest_surface_wins(self) -> None: """Two samples landing on the same pixel keep the nearer depth.""" # Both source pixels project along the optical axis to dst pixel (0, 0). params = make_params(np.zeros((3, 3)), np.zeros(3), src_shape=(1, 2), dst_shape=(2, 2)) @@ -49,7 +51,7 @@ def test_nearest_surface_wins(self): np.testing.assert_allclose(out, 1.0) # 2x2 splat covers the full output - def test_invalid_depth_leaves_holes(self): + def test_invalid_depth_leaves_holes(self) -> None: params = make_params(np.eye(3), np.zeros(3), src_shape=(4, 4), dst_shape=(4, 4)) raw = np.zeros((4, 4), dtype=np.float32) raw[1, 1] = 3.0 @@ -61,7 +63,7 @@ def test_invalid_depth_leaves_holes(self): expected[1:3, 1:3] = 3.0 # the sample plus its 2x2 splat np.testing.assert_allclose(out, expected) - def test_points_behind_camera_are_skipped(self): + def test_points_behind_camera_are_skipped(self) -> None: params = make_params(-np.eye(3), np.zeros(3), src_shape=(4, 4), dst_shape=(4, 4)) raw = np.full((4, 4), 2.0, dtype=np.float32) out = np.empty((4, 4), dtype=np.float32) @@ -70,7 +72,7 @@ def test_points_behind_camera_are_skipped(self): np.testing.assert_allclose(out, 0.0) - def test_shape_mismatch_raises(self): + def test_shape_mismatch_raises(self) -> None: params = make_params(np.eye(3), np.zeros(3), src_shape=(4, 4), dst_shape=(4, 4)) with pytest.raises(ValueError, match="Raw depth shape"): register_depth(np.zeros((2, 2), dtype=np.float32), params, np.zeros((4, 4), np.float32)) @@ -113,7 +115,7 @@ def _make_response( class TestExtractRegistrationParams: - def test_coincident_cameras_reproject_in_place(self): + def test_coincident_cameras_reproject_in_place(self) -> None: """Same pose and intrinsics for both cameras => registration is a no-op.""" rgb = _make_response("rgb", rows=20, cols=30, fx=50.0, fy=50.0, cx=0.0, cy=0.0) depth = _make_response("depth", rows=20, cols=30, fx=50.0, fy=50.0, cx=0.0, cy=0.0) @@ -126,7 +128,7 @@ def test_coincident_cameras_reproject_in_place(self): np.testing.assert_allclose(out, 2.0, rtol=1e-5) - def test_translated_depth_camera_shifts_projection(self): + def test_translated_depth_camera_shifts_projection(self) -> None: """A depth camera offset by baseline b shifts pixels by fx * b / z.""" fx = 100.0 rgb = _make_response("rgb", rows=20, cols=40, fx=fx, fy=fx, cx=0.0, cy=0.0) diff --git a/src/cuda_kernels.cu b/src/cuda_kernels.cu index 47b7db4..d18e7f0 100644 --- a/src/cuda_kernels.cu +++ b/src/cuda_kernels.cu @@ -895,83 +895,86 @@ cudaError_t preprocess_depth_image2( // Reprojects a raw depth image (depth sensor frame) into the RGB camera's image // plane, producing the same layout as Spot's *_depth_in_visual_frame sources. -__global__ void fill_float_kernel(float* __restrict__ buf, int n, float value) { - int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n) buf[i] = value; -} +// Empty-pixel sentinel: 0xFF-memset-able, and the max unsigned value, so any +// real depth sample's atomicMin replaces it. +constexpr unsigned int kRegistrationSentinelBits = 0xFFFFFFFFu; __global__ void register_depth_kernel( - const float* __restrict__ src, + const uint16_t* __restrict__ src, float* __restrict__ dst, DepthRegistrationParams p ) { - int u = blockIdx.x * blockDim.x + threadIdx.x; - int v = blockIdx.y * blockDim.y + threadIdx.y; - if (u >= p.src_width || v >= p.src_height) return; - - float z = src[v * p.src_width + u]; - if (z <= 0.0f) return; // invalid depth sample - - float uf = float(u), vf = float(v); - float px = z * (p.M[0] * uf + p.M[1] * vf + p.M[2]) + p.t[0]; - float py = z * (p.M[3] * uf + p.M[4] * vf + p.M[5]) + p.t[1]; - float pz = z * (p.M[6] * uf + p.M[7] * vf + p.M[8]) + p.t[2]; - if (pz <= 0.0f) return; // behind the RGB camera - - float inv_z = 1.0f / pz; - int u0 = __float2int_rd(px * inv_z); - int v0 = __float2int_rd(py * inv_z); - - // 2x2 splat: the depth image is lower resolution than the RGB target, so - // covering the four nearest pixels reduces holes without a dilation pass. - #pragma unroll - for (int dv = 0; dv <= 1; ++dv) { + const int n_src = p.src_width * p.src_height; + for (int idx = blockIdx.x * blockDim.x + threadIdx.x; idx < n_src; + idx += gridDim.x * blockDim.x) { + const uint16_t raw = src[idx]; + if (raw == 0) continue; // invalid depth sample + + const int u = idx % p.src_width; + const int v = idx / p.src_width; + const float z = float(raw) * p.inv_depth_scale; + + // Depth -> RGB reprojection + float uf = float(u), vf = float(v); + float px = z * (p.M[0] * uf + p.M[1] * vf + p.M[2]) + p.t[0]; + float py = z * (p.M[3] * uf + p.M[4] * vf + p.M[5]) + p.t[1]; + float pz = z * (p.M[6] * uf + p.M[7] * vf + p.M[8]) + p.t[2]; + if (pz <= 0.0f) continue; // behind the RGB camera + + float inv_z = fast_rcp(pz); + int u0 = __float2int_rd(px * inv_z); + int v0 = __float2int_rd(py * inv_z); + + // 2x2 splat: the depth image is lower resolution than the RGB target, so + // covering the four nearest pixels reduces holes without a dilation pass. #pragma unroll - for (int du = 0; du <= 1; ++du) { - int x = u0 + du; - int y = v0 + dv; - if (x < 0 || y < 0 || x >= p.dst_width || y >= p.dst_height) continue; - // atomicMin on the raw bits keeps the nearest surface when several - // depth pixels land on one RGB pixel; ordering is correct because - // positive IEEE-754 floats compare like unsigned ints. - atomicMin(reinterpret_cast(&dst[y * p.dst_width + x]), - __float_as_uint(pz)); + for (int dv = 0; dv <= 1; ++dv) { + #pragma unroll + for (int du = 0; du <= 1; ++du) { + int x = u0 + du; + int y = v0 + dv; + if (x < 0 || y < 0 || x >= p.dst_width || y >= p.dst_height) continue; + // atomicMin on the raw bits keeps the nearest surface when several + // depth pixels land on one RGB pixel; ordering is correct because + // positive IEEE-754 floats compare like unsigned ints. + atomicMin(reinterpret_cast(&dst[y * p.dst_width + x]), + __float_as_uint(pz)); + } } } } -__global__ void finalize_registered_depth_kernel(float* __restrict__ dst, int n, float sentinel) { +__global__ void finalize_registered_depth_kernel(float* __restrict__ dst, int n) { int i = blockIdx.x * blockDim.x + threadIdx.x; - if (i < n && dst[i] == sentinel) dst[i] = 0.0f; + // The sentinel bit pattern is a NaN, so compare bits, not float values. + if (i < n && __float_as_uint(dst[i]) == kRegistrationSentinelBits) dst[i] = 0.0f; } cudaError_t register_depth_to_rgb( - const float* d_depth_in, + const uint16_t* d_depth_in, float* d_depth_out, const DepthRegistrationParams& params, cudaStream_t stream ) { const int n_dst = params.dst_width * params.dst_height; - if (n_dst <= 0 || params.src_width <= 0 || params.src_height <= 0) { + const int n_src = params.src_width * params.src_height; + if (n_dst <= 0 || n_src <= 0) { return cudaErrorInvalidValue; } - constexpr float kSentinel = FLT_MAX; constexpr int threads = 256; - fill_float_kernel<<<(n_dst + threads - 1) / threads, threads, 0, stream>>>( - d_depth_out, n_dst, kSentinel); - cudaError_t err = cudaGetLastError(); + // 0xFF in every byte == kRegistrationSentinelBits; the memset fast path + // replaces a dedicated fill kernel launch. + cudaError_t err = cudaMemsetAsync(d_depth_out, 0xFF, n_dst * sizeof(float), stream); if (err != cudaSuccess) return err; - dim3 block(32, 8); - dim3 grid((params.src_width + block.x - 1) / block.x, - (params.src_height + block.y - 1) / block.y); - register_depth_kernel<<>>(d_depth_in, d_depth_out, params); + register_depth_kernel<<<(n_src + threads - 1) / threads, threads, 0, stream>>>( + d_depth_in, d_depth_out, params); err = cudaGetLastError(); if (err != cudaSuccess) return err; finalize_registered_depth_kernel<<<(n_dst + threads - 1) / threads, threads, 0, stream>>>( - d_depth_out, n_dst, kSentinel); + d_depth_out, n_dst); return cudaGetLastError(); } @@ -1357,21 +1360,11 @@ __global__ void postprocess_depth_image( int idx = y * width + x; float generated_val = generated_depth[idx]; - //float new_val = sparse_depth[idx]; float old_val = cached_depth[idx]; - // Check if new depth value is valid - //bool new_valid = (new_val >= min_valid_depth && new_val <= max_valid_depth); bool old_valid = (old_val >= min_valid_depth && old_val <= max_valid_depth); float result; - // if (new_valid) { - // if (old_valid) { - // result = old_val * (1.f-alpha_valid) + new_val * alpha_valid; - // } else { - // result = new_val; - // } - // } else if (old_valid) { result = old_val * (1.f-alpha_invalid) + generated_val * alpha_invalid; } else { diff --git a/src/dumper.cpp b/src/dumper.cpp index 3a9f110..80e1813 100644 --- a/src/dumper.cpp +++ b/src/dumper.cpp @@ -327,55 +327,49 @@ void DumpCameraTransform( } } -// void DumpDepthImageFromCuda(const uint16_t* depth, int32_t width, int32_t height, const std::string& subdir, int32_t dump_id) { -// if (!m_dumps_enabled) { -// return; -// } -// -// const size_t num_pixels = width * height; -// const size_t byte_size = num_pixels * sizeof(uint16_t); -// -// // Allocate host memory for float depth image -// std::vector h_depth(num_pixels); -// -// // Copy depth image from device to host -// cudaError_t err = cudaMemcpy(h_depth.data(), depth, byte_size, cudaMemcpyDeviceToHost); -// if (err != cudaSuccess) { -// LogMessage("Failed to copy depth image from device to host: {}", cudaGetErrorString(err)); -// return; -// } -// -// // Find min and max depth values for normalization -// uint16_t min_val = h_depth[0]; -// uint16_t max_val = h_depth[0]; -// for (size_t i = 1; i < num_pixels; ++i) { -// if (h_depth[i] < min_val) min_val = h_depth[i]; -// if (h_depth[i] > max_val) max_val = h_depth[i]; -// } -// -// // Allocate host memory for 8-bit grayscale image -// std::vector h_depth_u8(num_pixels); -// uint16_t range = max_val - min_val; -// -// // Normalize and convert to uint8 -// if (range > 0) { -// for (size_t i = 0; i < num_pixels; ++i) { -// h_depth_u8[i] = static_cast(((h_depth[i] - min_val) / range) * 255.0f); -// } -// } else { -// // If range is zero, the image is constant. Set to mid-gray. -// std::fill(h_depth_u8.begin(), h_depth_u8.end(), 128); -// } -// -// // Generate filename -// std::string file_path = m_dump_path + "/" + subdir + "/depth_" + std::to_string(dump_id) + ".png"; -// -// // Write image to file -// const int channels = 1; // Grayscale -// const int stride_in_bytes = width * channels; -// if (!stbi_write_png(file_path.c_str(), width, height, channels, h_depth_u8.data(), stride_in_bytes)) { -// LogMessage("Failed to write depth image to: {}", file_path); -// } -// } +void DumpDepthImageFromCuda( + const uint16_t* depth, + int32_t width, + int32_t height, + const std::string& subdir, + int32_t dump_id, + float scale +) { + if (!m_dumps_enabled) { + return; + } + + const size_t num_pixels = size_t(width) * height; + + std::vector h_depth_u16(num_pixels); + cudaError_t err = cudaMemcpy( + h_depth_u16.data(), depth, num_pixels * sizeof(uint16_t), cudaMemcpyDeviceToHost); + if (err != cudaSuccess) { + LogMessage("Failed to copy depth image from device to host: {}", cudaGetErrorString(err)); + return; + } + + // Convert with scale (meters per raw unit) and write the same reloadable + // count-prefixed float32 format as the float overload. + std::vector h_depth(num_pixels); + for (size_t i = 0; i < num_pixels; ++i) { + h_depth[i] = float(h_depth_u16[i]) * scale; + } + + std::string file_path = get_file_path(m_dump_path, subdir, "", dump_id, ""); + std::ofstream file(file_path, std::ios::binary); + if (!file) { + LogMessage("Failed to write depth image to: {}", file_path); + return; + } + + const uint32_t count = static_cast(h_depth.size()); + file.write(reinterpret_cast(&count), sizeof(count)); + file.write(reinterpret_cast(h_depth.data()), static_cast(h_depth.size() * sizeof(float))); + + if (!file) { + LogMessage("Failed while writing depth image to: {}", file_path); + } +} } // namespace STb diff --git a/src/include/cuda_kernels.cuh b/src/include/cuda_kernels.cuh index 85049d8..f49aa27 100644 --- a/src/include/cuda_kernels.cuh +++ b/src/include/cuda_kernels.cuh @@ -16,16 +16,18 @@ cudaError_t setOutputToOnes_launcher(float* d_input, float* d_output, int size); struct DepthRegistrationParams { float M[9]; // Row-major 3x3 float t[3]; + float inv_depth_scale{0.f}; // Meters per raw uint16 unit (1 / ImageSource.depth_scale) int src_width{0}; int src_height{0}; int dst_width{0}; int dst_height{0}; }; -// Reprojects d_depth_in (src-sized, meters) into d_depth_out (dst-sized, meters, -// fully overwritten; pixels with no depth sample are 0). +// Reprojects d_depth_in (src-sized, raw uint16 as served by Spot, converted to +// meters via inv_depth_scale) into d_depth_out (dst-sized, meters, fully +// overwritten; pixels with no depth sample are 0). cudaError_t register_depth_to_rgb( - const float* d_depth_in, + const uint16_t* d_depth_in, float* d_depth_out, const DepthRegistrationParams& params, cudaStream_t stream = 0 diff --git a/src/include/dumper.h b/src/include/dumper.h index d1bba7e..dbf0521 100644 --- a/src/include/dumper.h +++ b/src/include/dumper.h @@ -11,7 +11,8 @@ void DumpRGBImageFromCuda(const uint8_t* image, int32_t width, int32_t height, i // png_mode=false writes full-precision float32 binary (count-prefixed, reloadable); // png_mode=true writes a min/max-normalized 8-bit grayscale PNG for visualization. void DumpDepthImageFromCuda(const float* depth, int32_t width, int32_t height, const std::string& subdir, int32_t dump_id, bool png_mode = false); -void DumpDepthImageFromCuda(const uint16_t* depth, int32_t width, int32_t height, const std::string& subdir, int32_t dump_id); +// scale converts raw u16 units to meters (e.g. 1 / ImageSource.depth_scale). +void DumpDepthImageFromCuda(const uint16_t* depth, int32_t width, int32_t height, const std::string& subdir, int32_t dump_id, float scale = 1.f); void DumpCameraTransform(const std::string& name, const std::string& subdir, const float* transform, int32_t dump_id); diff --git a/src/include/spot-connection.h b/src/include/spot-connection.h index 567d468..839d9c3 100644 --- a/src/include/spot-connection.h +++ b/src/include/spot-connection.h @@ -42,12 +42,12 @@ class ReaderWriterCBuf { float* depth_data_{nullptr}; float* cached_depth_{nullptr}; - // Client-side depth registration: per-camera reprojection params. Raw depth - // images are staged in raw_depth_scratch_ (CUDA memory) and registered into - // the depth ring at RGB resolution. + // Client-side depth registration: per-camera reprojection params. Raw u16 + // depth images are staged in raw_depth_scratch_ (CUDA memory) and registered + // into the depth ring at RGB resolution (converted to meters on the GPU). std::vector depth_registration_; size_t n_elems_per_raw_depth_{0}; - float* raw_depth_scratch_{nullptr}; + uint16_t* raw_depth_scratch_{nullptr}; bool first_run_{true}; const size_t max_size_; // Maximum size of the queue diff --git a/src/spot-connection.cpp b/src/spot-connection.cpp index acf88d3..2287d70 100644 --- a/src/spot-connection.cpp +++ b/src/spot-connection.cpp @@ -125,6 +125,15 @@ static bool build_depth_registration_params( return false; } + // Raw u16 -> meters conversion happens on the GPU inside the registration kernel. + const double depth_scale = depth_source.depth_scale(); + if (depth_scale <= 0.0) { + LogMessage("build_depth_registration_params: Invalid depth_scale {} for depth '{}'", + depth_scale, depth_source.name()); + return false; + } + out.inv_depth_scale = static_cast(1.0 / depth_scale); + const std::array T = se3_pose_to_row_major_matrix(rgb_T_depth); // A = R * K_depth^-1 @@ -281,7 +290,7 @@ bool ReaderWriterCBuf::initialize( raw_depth_scratch_ = nullptr; } checkCudaError( - cudaMalloc(&raw_depth_scratch_, n_elems_per_raw_depth_ * sizeof(float)), + cudaMalloc(&raw_depth_scratch_, n_elems_per_raw_depth_ * sizeof(uint16_t)), "cudaMalloc for raw depth scratch" ); @@ -324,14 +333,14 @@ void ReaderWriterCBuf::push(const google::protobuf::RepeatedPtrField start_time = high_resolution_clock::now(); - // time_point end_time = high_resolution_clock::now(); - // - // auto duration = duration_cast(end_time - start_time); - // double latency_ms = duration.count() / 1000.0; - // - // LogMessage("ReaderWriterCBuf::push: latency: {:.4f} ms", latency_ms); - // start_time = end_time; // Reset start time for next push + static time_point start_time = high_resolution_clock::now(); + time_point end_time = high_resolution_clock::now(); + + auto duration = duration_cast(end_time - start_time); + double latency_ms = duration.count() / 1000.0; + + LogPerf("ReaderWriterCBuf::push: latency: {:.4f} ms", latency_ms); + start_time = end_time; // Reset start time for next push // Compute write pointer int write_idx = write_idx_.load(std::memory_order_relaxed); @@ -424,12 +433,11 @@ void ReaderWriterCBuf::push(const google::protobuf::RepeatedPtrField= n_images_per_response_) { throw std::runtime_error("ReaderWriterCBuf::push: Too many depth responses"); } - cv::Mat cv_img = convert_image_to_cv_mat(img, response.source().depth_scale()); - int depth_width = cv_img.cols; - int depth_height = cv_img.rows; + int depth_width = img.cols(); + int depth_height = img.rows(); const DepthRegistrationParams& reg = depth_registration_[n_depths_written]; - size_t depth_size = depth_width * depth_height; + size_t depth_size = size_t(depth_width) * depth_height; if (depth_size != n_elems_per_raw_depth_ || depth_width != reg.src_width || depth_height != reg.src_height) { LogMessage("Raw depth size mismatch: expected {} ({}x{}), got {} ({}x{})", @@ -437,18 +445,24 @@ void ReaderWriterCBuf::push(const google::protobuf::RepeatedPtrField Date: Fri, 10 Jul 2026 18:16:35 -0400 Subject: [PATCH 4/6] Update the install --- install | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/install b/install index 440ce5a..d402033 160000 --- a/install +++ b/install @@ -1 +1 @@ -Subproject commit 440ce5acb594f13b3ac758d14e0cd03be5341164 +Subproject commit d40203386a71030a87972f92be0879a4f8028ca0 From 089501933cf56fbf05d351eb21744eaffd6b0bda Mon Sep 17 00:00:00 2001 From: Faisal Zaghloul Date: Fri, 10 Jul 2026 19:21:26 -0400 Subject: [PATCH 5/6] Fixes Run nvcc only for sm86 and sm89. Previously Torch was mucking with CMAKE_CUDA_ARCHITECTURES so it was building for more archs. Fix the DLL copying nightmare in cmake. Fix corner case where depth cameras transmit different-sized depth. Rework the python depth registration to avoid reallocating buffers every iteration. --- CMakeLists.txt | 202 +++++------------- .../pyspotobserver/camera_stream.py | 26 ++- .../pyspotobserver/depth_registration.py | 122 ++++++++--- .../tests/test_depth_registration.py | 19 ++ src/include/spot-connection.h | 5 +- src/include/vision-pipeline.h | 2 +- src/spot-connection.cpp | 57 +++-- tests/integ-test-dx12/CMakeLists.txt | 42 +--- tests/integ-test-dx12/integ-test-dx12.cpp | 2 +- tests/integ-test/CMakeLists.txt | 42 +--- tests/standalone-tests/CMakeLists.txt | 33 +-- 11 files changed, 233 insertions(+), 319 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 606a7b7..f833750 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -5,6 +5,12 @@ option(BUILD_TESTS "Build the test executable" ON) option(INSTALL_TESTS "Install test executables" ON) set(PLUGIN_NAME "SpotObserverLib") +# All runtime artifacts (plugin DLL, test executables, copied dependency DLLs) +# land in the top-level build directory, so dependency DLLs are copied exactly +# once and every executable runs from the same place. Multi-config generators +# (Visual Studio) append a per-config subdir, e.g. build/Release. +set(CMAKE_RUNTIME_OUTPUT_DIRECTORY ${CMAKE_BINARY_DIR}) + # Set default install directory to SpotObserver/install if not specified if(CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT) set(CMAKE_INSTALL_PREFIX "${CMAKE_CURRENT_SOURCE_DIR}/install" CACHE PATH "Default install directory" FORCE) @@ -37,11 +43,20 @@ find_package(Protobuf REQUIRED PATHS "${Protobuf_ROOT}" NO_DEFAULT_PATH) find_package(gRPC REQUIRED) find_package(Threads REQUIRED) -# Support Ampere and Ada architectures - set BEFORE finding Torch +# Support Ampere and Ada architectures - set BEFORE finding Torch. +# NOTE: TorchConfig (Caffe2/public/cuda.cmake) force-sets CMAKE_CUDA_ARCHITECTURES +# to OFF and instead injects -gencode flags derived from TORCH_CUDA_ARCH_LIST into +# CMAKE_CUDA_FLAGS for every .cu in the project. TORCH_CUDA_ARCH_LIST is therefore +# what actually selects the built architectures, so derive it from our arch list +# ("Common" here used to make every .cu build for ~10 archs, sm_35 through sm_90). set(CMAKE_CUDA_ARCHITECTURES 86;89 CACHE STRING "CUDA archs" FORCE) -#set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -gencode arch=compute_86,code=sm_86 --use_fast_math") -set(TORCH_CUDA_ARCH_LIST "Common") +set(TORCH_CUDA_ARCH_LIST "") +foreach(_arch ${CMAKE_CUDA_ARCHITECTURES}) + string(REGEX REPLACE "([0-9])$" ".\\1" _arch_dotted "${_arch}") + list(APPEND TORCH_CUDA_ARCH_LIST "${_arch_dotted}") +endforeach() message(STATUS "CMAKE_CUDA_ARCHITECTURES: ${CMAKE_CUDA_ARCHITECTURES}") +message(STATUS "TORCH_CUDA_ARCH_LIST: ${TORCH_CUDA_ARCH_LIST}") set(USE_SYSTEM_NVTX ON) find_package(Torch REQUIRED) @@ -50,6 +65,9 @@ if(NOT Torch_FOUND) endif() message(STATUS "Found libtorch: ${TORCH_LIBRARIES}") message(STATUS "Torch install prefix: ${TORCH_INSTALL_PREFIX}") +# Torch's config appended its arch flags to CMAKE_CUDA_FLAGS; surface them so a +# wrong -gencode set is visible at configure time. +message(STATUS "CMAKE_CUDA_FLAGS after Torch: ${CMAKE_CUDA_FLAGS}") set(ONNX_DIR ${CMAKE_CURRENT_SOURCE_DIR}/extern/onnxruntime-win-x64-gpu-1.22.0) # Check if ONNX_DIR exists @@ -162,127 +180,46 @@ endif() # Option to disable automatic DLL copying option(COPY_DLLS "Automatically copy required DLLs to executable directory" ON) -# Copy required DLLs to executable directory after build -if(COPY_DLLS) - # Find vcpkg DLL directory dynamically - if(DEFINED CMAKE_TOOLCHAIN_FILE) - get_filename_component(VCPKG_ROOT "${CMAKE_TOOLCHAIN_FILE}" DIRECTORY) - get_filename_component(VCPKG_ROOT "${VCPKG_ROOT}" DIRECTORY) - get_filename_component(VCPKG_ROOT "${VCPKG_ROOT}" DIRECTORY) - set(VCPKG_DLL_DIR "${VCPKG_ROOT}/vcpkg_installed/x64-windows/bin") - else() - # Fallback: try to find vcpkg in common locations - set(VCPKG_DLL_DIR "") - foreach(path "$ENV{VCPKG_ROOT}/vcpkg_installed/x64-windows/bin" - "${CMAKE_SOURCE_DIR}/../vcpkg/vcpkg_installed/x64-windows/bin" - "${CMAKE_SOURCE_DIR}/../../vcpkg/vcpkg_installed/x64-windows/bin") - if(EXISTS "${path}") - set(VCPKG_DLL_DIR "${path}") - break() - endif() - endforeach() - endif() - - set(OPENCV_DLL_DIR "${CMAKE_SOURCE_DIR}/extern/opencv/x64/vc16/bin") - - # List of required DLLs from vcpkg - set(VCPKG_DLLS - "re2.dll" - "libprotobuf.dll" - "libprotobuf-lite.dll" - "abseil_dll.dll" - "cares.dll" - "libcrypto-3-x64.dll" - "libssl-3-x64.dll" - "zlib1.dll" - ) - - # List of required OpenCV DLLs - set(OPENCV_DLLS - "opencv_world4110.dll" - ) - - set(PLUGIN_DLLS - ${CMAKE_CURRENT_BINARY_DIR}/${PLUGIN_NAME}.dll - ) - - set(MLLIB_DLLS) - - # Use GLOB to find all DLLs in the libtorch/lib directory +# Copy required DLLs next to the built binaries after build +if(COPY_DLLS AND WIN32) + # DLLs of imported targets (vcpkg: protobuf, gRPC, OpenSSL, abseil, ...) are + # enumerated by CMake itself via $ — no hardcoded DLL + # names or vcpkg paths. OpenCV / LibTorch / ONNX Runtime are linked as raw + # libs (not imported targets), and torch/ONNX also load some DLLs dynamically + # at runtime (cuDNN, CUDA execution provider), so their DLLs are collected + # from their lib dirs explicitly. + set(OPENCV_DLLS "${CMAKE_SOURCE_DIR}/extern/opencv/x64/vc16/bin/opencv_world4110.dll") file(GLOB LIBTORCH_DLLS "${TORCH_INSTALL_PREFIX}/lib/*.dll") - message(STATUS "Found LibTorch DLLs via GLOB: ${LIBTORCH_DLLS}") - list(APPEND MLLIB_DLLS ${LIBTORCH_DLLS}) - - # Find ONNX Runtime DLLs file(GLOB ONNX_DLLS "${ONNX_LIB_DIR}/*.dll") - message(STATUS "Found ONNX Runtime DLLs via GLOB: ${ONNX_DLLS}") - list(APPEND MLLIB_DLLS ${ONNX_DLLS}) - # Copy vcpkg DLLs (only if vcpkg directory was found) - if(VCPKG_DLL_DIR AND EXISTS "${VCPKG_DLL_DIR}") - message(STATUS "Found vcpkg DLL directory: ${VCPKG_DLL_DIR}") - foreach(dll ${VCPKG_DLLS}) - if(EXISTS "${VCPKG_DLL_DIR}/${dll}") - add_custom_command(TARGET ${PLUGIN_NAME} POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy_if_different - "${VCPKG_DLL_DIR}/${dll}" - $ - COMMENT "Copying ${dll}" - ) - else() - message(WARNING "vcpkg DLL not found: ${VCPKG_DLL_DIR}/${dll}") - endif() - endforeach() - else() - message(WARNING "vcpkg DLL directory not found. DLLs will need to be copied manually.") + # $ misses OpenSSL/ZLIB: they enter through CMake + # find-modules (via gRPC), whose UNKNOWN-type imported targets the genex + # cannot enumerate. Locate the vcpkg bin dir from protobuf's imported DLL + # (no path guessing) and take every DLL vcpkg built for this manifest; + # overlap with the genex is harmless for copy_if_different/install. + get_target_property(_protobuf_dll protobuf::libprotobuf IMPORTED_LOCATION_RELEASE) + if(NOT _protobuf_dll) + get_target_property(_protobuf_dll protobuf::libprotobuf IMPORTED_LOCATION) endif() - - # Copy OpenCV DLLs (only if OpenCV directory exists) - if(EXISTS "${OPENCV_DLL_DIR}") - message(STATUS "Found OpenCV DLL directory: ${OPENCV_DLL_DIR}") - foreach(dll ${OPENCV_DLLS}) - if(EXISTS "${OPENCV_DLL_DIR}/${dll}") - add_custom_command(TARGET ${PLUGIN_NAME} POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy_if_different - "${OPENCV_DLL_DIR}/${dll}" - $ - COMMENT "Copying ${dll}" - ) - else() - message(WARNING "OpenCV DLL not found: ${OPENCV_DLL_DIR}/${dll}") - endif() - endforeach() - else() - message(WARNING "OpenCV DLL directory not found: ${OPENCV_DLL_DIR}") - endif() - - # Copy ML library DLLs (Libtorch, ONNX Runtime, etc.) - foreach(dll ${MLLIB_DLLS}) - if(EXISTS "${dll}") - add_custom_command(TARGET ${PLUGIN_NAME} POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy_if_different - "${dll}" - $ - COMMENT "Copying ${dll}" - ) - else() - message(WARNING "ML library DLL not found: ${dll}") - endif() - endforeach() + get_filename_component(VCPKG_BIN_DIR "${_protobuf_dll}" DIRECTORY) + file(GLOB VCPKG_DLLS "${VCPKG_BIN_DIR}/*.dll") + + set(EXTERN_RUNTIME_DLLS ${VCPKG_DLLS} ${OPENCV_DLLS} ${LIBTORCH_DLLS} ${ONNX_DLLS}) + + add_custom_command(TARGET ${PLUGIN_NAME} POST_BUILD + COMMAND ${CMAKE_COMMAND} -E copy_if_different + $ + ${EXTERN_RUNTIME_DLLS} + $ + COMMAND_EXPAND_LISTS + COMMENT "Copying runtime DLLs next to ${PLUGIN_NAME}" + ) endif() # COPY_DLLS if (BUILD_TESTS) add_subdirectory(tests) endif() -## Output directories -#set_target_properties(${PLUGIN_NAME} -# PROPERTIES -# RUNTIME_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/bin" -# LIBRARY_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/bin" -# ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}/bin" -#) -# # INSTALL include(GNUInstallDirs) @@ -301,38 +238,13 @@ install(DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}/include/ FILES_MATCHING PATTERN "*.h" PATTERN "*.hpp" ) -# Install required DLLs on Windows +# Install required DLLs on Windows (same set as the post-build copy above). if(WIN32 AND COPY_DLLS) - # Install vcpkg DLLs - if(VCPKG_DLL_DIR AND EXISTS "${VCPKG_DLL_DIR}") - foreach(dll ${VCPKG_DLLS}) - if(EXISTS "${VCPKG_DLL_DIR}/${dll}") - install(FILES "${VCPKG_DLL_DIR}/${dll}" - DESTINATION ${CMAKE_INSTALL_BINDIR} - ) - endif() - endforeach() - endif() - - # Install OpenCV DLLs - if(EXISTS "${OPENCV_DLL_DIR}") - foreach(dll ${OPENCV_DLLS}) - if(EXISTS "${OPENCV_DLL_DIR}/${dll}") - install(FILES "${OPENCV_DLL_DIR}/${dll}" - DESTINATION ${CMAKE_INSTALL_BINDIR} - ) - endif() - endforeach() - endif() - - # Install ML library DLLs - foreach(dll ${MLLIB_DLLS}) - if(EXISTS "${dll}") - install(FILES "${dll}" - DESTINATION ${CMAKE_INSTALL_BINDIR} - ) - endif() - endforeach() + install(FILES + $ + ${EXTERN_RUNTIME_DLLS} + DESTINATION ${CMAKE_INSTALL_BINDIR} + ) endif() # Export targets for find_package support diff --git a/PySpotObserver/pyspotobserver/camera_stream.py b/PySpotObserver/pyspotobserver/camera_stream.py index 4ddf94c..bd13a9e 100644 --- a/PySpotObserver/pyspotobserver/camera_stream.py +++ b/PySpotObserver/pyspotobserver/camera_stream.py @@ -22,6 +22,8 @@ from .config import CameraType, SpotConfig from .depth_registration import ( DepthRegistrationParams, + DepthRegistrationWorkspace, + create_registration_workspace, extract_registration_params, register_depth, ) @@ -157,6 +159,7 @@ def __init__( # 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 + self._depth_reg_workspaces: list[DepthRegistrationWorkspace] | None = None # Optional per-stream vision pipeline, imported lazily. self._vision_pipeline: VisionPipeline | None = None @@ -251,6 +254,7 @@ def start_streaming(self, camera_mask: int) -> None: self._frame_pool_index = 0 self._ccm_scratch_by_shape = {} self._depth_reg_params = None + self._depth_reg_workspaces = None self._frame_count = 0 self._error_count = 0 self._last_body_to_worlds = None @@ -839,13 +843,17 @@ def _fill_frame_from_responses( responses[rgb_idx], is_depth=False, out_array=frame.rgb_images[frame_i], ccm=ccm ) - 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) + if self._depth_reg_params is None or self._depth_reg_workspaces is None: + raise SpotCamStreamError("Depth registration state not initialized") + workspace = self._depth_reg_workspaces[i] + self._convert_image_response_inplace( + responses[depth_idx], is_depth=True, out_array=workspace.raw_depth + ) register_depth( - raw_depth, + workspace.raw_depth, self._depth_reg_params[i], out=frame.depth_images[frame_i], + workspace=workspace, ) # Update timestamps and camera extrinsics params @@ -897,6 +905,7 @@ def _decode_initial_responses( decoded: list[np.ndarray] = [] reg_params: list[DepthRegistrationParams] = [] + reg_workspaces: list[DepthRegistrationWorkspace] = [] for i in range(n_cameras): rgb_idx = i * 2 depth_idx = i * 2 + 1 @@ -905,7 +914,6 @@ def _decode_initial_responses( self._apply_ccm_inplace(rgb, self._ccms[self._sdk_camera_order[i]]) decoded.append(rgb) - 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] @@ -916,12 +924,18 @@ def _decode_initial_responses( f"{self._sdk_camera_order[i].name}: {exc}" ) from exc reg_params.append(params) + workspace = create_registration_workspace(params) + reg_workspaces.append(workspace) + self._convert_image_response_inplace( + responses[depth_idx], is_depth=True, out_array=workspace.raw_depth + ) registered = np.zeros(rgb.shape[:2], dtype=np.float32) - register_depth(raw_depth, params, registered) + register_depth(workspace.raw_depth, params, registered, workspace=workspace) decoded.append(registered) self._depth_reg_params = reg_params + self._depth_reg_workspaces = reg_workspaces return decoded def _cache_stitch_params(self, responses: list[image_pb2.ImageResponse]) -> None: diff --git a/PySpotObserver/pyspotobserver/depth_registration.py b/PySpotObserver/pyspotobserver/depth_registration.py index 7c40396..3d6c0b6 100644 --- a/PySpotObserver/pyspotobserver/depth_registration.py +++ b/PySpotObserver/pyspotobserver/depth_registration.py @@ -40,6 +40,46 @@ class DepthRegistrationParams: dst_shape: tuple[int, int] # (H_dst, W_dst) +@dataclass +class DepthRegistrationWorkspace: + """Reusable scratch buffers for :func:`register_depth`.""" + + raw_depth: np.ndarray + points: np.ndarray + projected_x: np.ndarray + projected_y: np.ndarray + pixel_x: np.ndarray + pixel_y: np.ndarray + indices: np.ndarray + valid: np.ndarray + splat_valid: np.ndarray + bool_scratch: np.ndarray + z_buffer: np.ndarray + dst_empty: np.ndarray + + +def create_registration_workspace( + params: DepthRegistrationParams, +) -> DepthRegistrationWorkspace: + """Allocate scratch storage once for repeated registration with ``params``.""" + src_shape = params.src_shape + dst_h, dst_w = params.dst_shape + return DepthRegistrationWorkspace( + raw_depth=np.empty(src_shape, dtype=np.float32), + points=np.empty((*src_shape, 3), dtype=np.float32), + projected_x=np.empty(src_shape, dtype=np.float32), + projected_y=np.empty(src_shape, dtype=np.float32), + pixel_x=np.empty(src_shape, dtype=np.int32), + pixel_y=np.empty(src_shape, dtype=np.int32), + indices=np.empty(src_shape, dtype=np.intp), + valid=np.empty(src_shape, dtype=np.bool_), + splat_valid=np.empty(src_shape, dtype=np.bool_), + bool_scratch=np.empty(src_shape, dtype=np.bool_), + z_buffer=np.empty(dst_h * dst_w + 1, dtype=np.float32), + dst_empty=np.empty((dst_h, dst_w), dtype=np.bool_), + ) + + def _pinhole_intrinsics(source: image_pb2.ImageSource) -> np.ndarray: """Return the 3x3 pinhole camera matrix of an ImageSource.""" which = source.WhichOneof("camera_models") @@ -109,6 +149,7 @@ def register_depth( raw_depth: np.ndarray, params: DepthRegistrationParams, out: np.ndarray, + workspace: DepthRegistrationWorkspace | None = None, ) -> None: """Reproject a raw depth image into the RGB frame, in-place into ``out``. @@ -117,40 +158,71 @@ def register_depth( params: Registration params for this camera pair. out: (H_dst, W_dst) float32 output, fully overwritten. Pixels that receive no depth sample are 0. + workspace: Optional reusable scratch storage. Supplying one avoids + per-frame allocations in streaming callers. """ 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 workspace is None: + workspace = create_registration_workspace(params) + if ( + workspace.raw_depth.shape != params.src_shape + or workspace.dst_empty.shape != params.dst_shape + ): + raise ValueError("Registration workspace shape does not match params") + + valid = workspace.valid + np.greater(raw_depth, 0.0, out=valid) if not valid.any(): + out.fill(0.0) 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: + points = workspace.points + np.multiply(params.rays, raw_depth[..., None], out=points) + np.add(points, params.t, out=points) + pz = points[..., 2] + np.greater(pz, 0.0, out=workspace.bool_scratch) + np.logical_and(valid, workspace.bool_scratch, out=valid) + if not valid.any(): + out.fill(0.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)) + np.divide(points[..., 0], pz, out=workspace.projected_x, where=valid) + np.divide(points[..., 1], pz, out=workspace.projected_y, where=valid) + np.floor(workspace.projected_x, out=workspace.projected_x) + np.floor(workspace.projected_y, out=workspace.projected_y) + workspace.pixel_x.fill(0) + workspace.pixel_y.fill(0) + np.copyto(workspace.pixel_x, workspace.projected_x, where=valid, casting="unsafe") + np.copyto(workspace.pixel_y, workspace.projected_y, where=valid, casting="unsafe") 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] + sentinel_index = dst_h * dst_w + workspace.z_buffer.fill(np.inf) + + # A z-buffered 2x2 splat. Invalid samples all target the extra sentinel + # element, avoiding boolean-indexed temporaries and a global depth sort. + for du, dv in ((0, 0), (1, 0), (0, 1), (1, 1)): + splat_valid = workspace.splat_valid + np.copyto(splat_valid, valid) + np.greater_equal(workspace.pixel_x, -du, out=workspace.bool_scratch) + np.logical_and(splat_valid, workspace.bool_scratch, out=splat_valid) + np.less(workspace.pixel_x, dst_w - du, out=workspace.bool_scratch) + np.logical_and(splat_valid, workspace.bool_scratch, out=splat_valid) + np.greater_equal(workspace.pixel_y, -dv, out=workspace.bool_scratch) + np.logical_and(splat_valid, workspace.bool_scratch, out=splat_valid) + np.less(workspace.pixel_y, dst_h - dv, out=workspace.bool_scratch) + np.logical_and(splat_valid, workspace.bool_scratch, out=splat_valid) + + np.multiply(workspace.pixel_y, dst_w, out=workspace.indices, casting="unsafe") + np.add(workspace.indices, workspace.pixel_x, out=workspace.indices) + np.add(workspace.indices, dv * dst_w + du, out=workspace.indices) + np.logical_not(splat_valid, out=workspace.bool_scratch) + np.copyto(workspace.indices, sentinel_index, where=workspace.bool_scratch) + np.minimum.at(workspace.z_buffer, workspace.indices, pz) + + np.copyto(out, workspace.z_buffer[:-1].reshape(params.dst_shape)) + np.isinf(out, out=workspace.dst_empty) + np.copyto(out, 0.0, where=workspace.dst_empty) diff --git a/PySpotObserver/tests/test_depth_registration.py b/PySpotObserver/tests/test_depth_registration.py index 877568e..70785d9 100644 --- a/PySpotObserver/tests/test_depth_registration.py +++ b/PySpotObserver/tests/test_depth_registration.py @@ -7,6 +7,7 @@ from bosdyn.api import image_pb2 # type: ignore[import-untyped] from pyspotobserver.depth_registration import ( DepthRegistrationParams, + create_registration_workspace, extract_registration_params, register_depth, ) @@ -79,6 +80,24 @@ def test_shape_mismatch_raises(self) -> None: with pytest.raises(ValueError, match="Output shape"): register_depth(np.zeros((4, 4), dtype=np.float32), params, np.zeros((2, 2), np.float32)) + def test_workspace_can_be_reused(self) -> None: + params = make_params( + np.eye(3), np.zeros(3), src_shape=(4, 4), dst_shape=(4, 4) + ) + workspace = create_registration_workspace(params) + out = np.empty((4, 4), dtype=np.float32) + + register_depth(np.full((4, 4), 2.0, np.float32), params, out, workspace=workspace) + np.testing.assert_allclose(out, 2.0) + + raw = np.zeros((4, 4), dtype=np.float32) + raw[1, 1] = 3.0 + register_depth(raw, params, out, workspace=workspace) + + expected = np.zeros((4, 4), dtype=np.float32) + expected[1:3, 1:3] = 3.0 + np.testing.assert_allclose(out, expected) + def _make_response( frame_name: str, diff --git a/src/include/spot-connection.h b/src/include/spot-connection.h index 839d9c3..09769c7 100644 --- a/src/include/spot-connection.h +++ b/src/include/spot-connection.h @@ -46,7 +46,8 @@ class ReaderWriterCBuf { // depth images are staged in raw_depth_scratch_ (CUDA memory) and registered // into the depth ring at RGB resolution (converted to meters on the GPU). std::vector depth_registration_; - size_t n_elems_per_raw_depth_{0}; + std::vector n_elems_per_raw_depth_; + size_t raw_depth_scratch_size_{0}; uint16_t* raw_depth_scratch_{nullptr}; bool first_run_{true}; @@ -70,7 +71,7 @@ class ReaderWriterCBuf { std::vector cameras, std::vector camera_dump_subdirs, std::vector depth_registration, - size_t n_elems_per_raw_depth + std::vector n_elems_per_raw_depth ); /** diff --git a/src/include/vision-pipeline.h b/src/include/vision-pipeline.h index c1dafb2..952c12a 100644 --- a/src/include/vision-pipeline.h +++ b/src/include/vision-pipeline.h @@ -28,7 +28,7 @@ class VisionPipeline { MLModel& model_; const SpotCamStream& spot_cam_stream_; bool first_run_{true}; - std::atomic depth_averaging_enabled_{true}; + std::atomic depth_averaging_enabled_{false}; // Threading std::unique_ptr pipeline_thread_; diff --git a/src/spot-connection.cpp b/src/spot-connection.cpp index 2287d70..71df593 100644 --- a/src/spot-connection.cpp +++ b/src/spot-connection.cpp @@ -14,6 +14,7 @@ #include #include +#include #include #include #include @@ -245,7 +246,7 @@ bool ReaderWriterCBuf::initialize( std::vector cameras, std::vector camera_dump_subdirs, std::vector depth_registration, - size_t n_elems_per_raw_depth + std::vector n_elems_per_raw_depth ) { n_elems_per_rgb_ = n_elems_per_rgb; n_elems_per_depth_ = n_elems_per_depth; @@ -253,19 +254,22 @@ bool ReaderWriterCBuf::initialize( cameras_ = std::move(cameras); camera_dump_subdirs_ = std::move(camera_dump_subdirs); depth_registration_ = std::move(depth_registration); - n_elems_per_raw_depth_ = n_elems_per_raw_depth; + n_elems_per_raw_depth_ = std::move(n_elems_per_raw_depth); if (camera_dump_subdirs_.size() != cameras_.size()) { LogMessage("ReaderWriterCBuf::initialize: Expected {} camera dump subdirs, got {}", cameras_.size(), camera_dump_subdirs_.size()); return false; } - if (depth_registration_.size() != cameras_.size() || n_elems_per_raw_depth_ == 0) { - LogMessage("ReaderWriterCBuf::initialize: Expected {} depth registration params " - "(got {}) and a nonzero raw depth size (got {})", - cameras_.size(), depth_registration_.size(), n_elems_per_raw_depth_); + if (depth_registration_.size() != cameras_.size() + || n_elems_per_raw_depth_.size() != cameras_.size() + || std::ranges::any_of(n_elems_per_raw_depth_, [](size_t size) { return size == 0; })) { + LogMessage("ReaderWriterCBuf::initialize: Expected {} depth registration params and " + "raw depth sizes (got {} params and {} sizes)", + cameras_.size(), depth_registration_.size(), n_elems_per_raw_depth_.size()); return false; } + raw_depth_scratch_size_ = *std::ranges::max_element(n_elems_per_raw_depth_); if (rgb_data_ != nullptr || depth_data_ != nullptr) { if (!rgb_data_) checkCudaError(cudaFree(rgb_data_), "cudaFree for RGB data"); @@ -290,7 +294,7 @@ bool ReaderWriterCBuf::initialize( raw_depth_scratch_ = nullptr; } checkCudaError( - cudaMalloc(&raw_depth_scratch_, n_elems_per_raw_depth_ * sizeof(uint16_t)), + cudaMalloc(&raw_depth_scratch_, raw_depth_scratch_size_ * sizeof(uint16_t)), "cudaMalloc for raw depth scratch" ); @@ -333,14 +337,14 @@ void ReaderWriterCBuf::push(const google::protobuf::RepeatedPtrField start_time = high_resolution_clock::now(); - time_point end_time = high_resolution_clock::now(); - - auto duration = duration_cast(end_time - start_time); - double latency_ms = duration.count() / 1000.0; - - LogPerf("ReaderWriterCBuf::push: latency: {:.4f} ms", latency_ms); - start_time = end_time; // Reset start time for next push + // static thread_local time_point start_time = high_resolution_clock::now(); + // time_point end_time = high_resolution_clock::now(); + // + // auto duration = duration_cast(end_time - start_time); + // double latency_ms = duration.count() / 1000.0; + // + // LogPerf("ReaderWriterCBuf::push: latency: {:.4f} ms", latency_ms); + // start_time = end_time; // Reset start time for next push // Compute write pointer int write_idx = write_idx_.load(std::memory_order_relaxed); @@ -437,11 +441,12 @@ void ReaderWriterCBuf::push(const google::protobuf::RepeatedPtrField depth_registration; depth_registration.reserve(num_cams_requested); + std::vector raw_depth_sizes; + raw_depth_sizes.reserve(num_cams_requested); for (int32_t j = 0; j < num_cams_requested; j++) { const auto& rgb_response = image_responses[j]; @@ -935,6 +929,7 @@ bool SpotCamStream::streamCameras(uint32_t cam_mask) { reg.dst_width = static_cast(current_rgb_shape_.W); reg.dst_height = static_cast(current_rgb_shape_.H); depth_registration.push_back(reg); + raw_depth_sizes.push_back(size_t(reg.src_width) * reg.src_height); LogMessage("Depth registration for camera {}: {}x{} ({}) -> {}x{} ({})", get_dump_name_for_camera(camera_order_[j]), @@ -955,7 +950,7 @@ bool SpotCamStream::streamCameras(uint32_t cam_mask) { camera_order_, std::move(camera_dump_subdirs), std::move(depth_registration), - raw_depth_ref_size + std::move(raw_depth_sizes) )) { LogMessage("SpotCamStream::streamCameras: Failed to initialize image buffer"); return false; diff --git a/tests/integ-test-dx12/CMakeLists.txt b/tests/integ-test-dx12/CMakeLists.txt index 07f90a1..da278f4 100644 --- a/tests/integ-test-dx12/CMakeLists.txt +++ b/tests/integ-test-dx12/CMakeLists.txt @@ -30,45 +30,9 @@ target_link_libraries(integ-test-dx12 CUDA::cudart ) -# Copy only the DLLs needed for this test -if(COPY_DLLS AND WIN32) - # Copy vcpkg DLLs (for gRPC/protobuf) - if(VCPKG_DLL_DIR AND EXISTS "${VCPKG_DLL_DIR}") - foreach(dll ${VCPKG_DLLS}) - if(EXISTS "${VCPKG_DLL_DIR}/${dll}") - add_custom_command(TARGET integ-test-dx12 POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy_if_different - "${VCPKG_DLL_DIR}/${dll}" - $ - COMMENT "Copying ${dll} for test" - ) - endif() - endforeach() - endif() - - # Copy OpenCV DLLs - if(EXISTS "${OPENCV_DLL_DIR}") - foreach(dll ${OPENCV_DLLS}) - if(EXISTS "${OPENCV_DLL_DIR}/${dll}") - add_custom_command(TARGET integ-test-dx12 POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy_if_different - "${OPENCV_DLL_DIR}/${dll}" - $ - COMMENT "Copying ${dll} for test" - ) - endif() - endforeach() - endif() - # Copy the main library DLL - if(EXISTS ${PLUGIN_DLLS}) - add_custom_command(TARGET integ-test-dx12 POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy_if_different - "${PLUGIN_DLLS}" - $ - COMMENT "Copying ${dll} for test" - ) - endif() -endif() +# Runtime DLLs are copied once into the shared top-level runtime directory +# (see CMAKE_RUNTIME_OUTPUT_DIRECTORY in the top-level CMakeLists.txt), where +# this executable also lands — no per-test copying needed. if(INSTALL_TESTS) install(TARGETS integ-test-dx12 diff --git a/tests/integ-test-dx12/integ-test-dx12.cpp b/tests/integ-test-dx12/integ-test-dx12.cpp index f14ce6d..84bd3fd 100644 --- a/tests/integ-test-dx12/integ-test-dx12.cpp +++ b/tests/integ-test-dx12/integ-test-dx12.cpp @@ -114,7 +114,7 @@ int main(int argc, char* argv[]) { // --------------------------------------------------------------------------------------- std::cout << "Registering test log callback...\n"; SOb_SetUnityLogCallback(TestLogCallback); - SOb_ToggleLogging(true); + SOb_SetLogLevel(1); //SOb_ToggleDebugDumps("./spot_dump_dx12"); diff --git a/tests/integ-test/CMakeLists.txt b/tests/integ-test/CMakeLists.txt index 716e960..fe88e55 100644 --- a/tests/integ-test/CMakeLists.txt +++ b/tests/integ-test/CMakeLists.txt @@ -30,45 +30,9 @@ target_link_libraries(integ-test CUDA::cudart ) -# Copy only the DLLs needed for this test -if(COPY_DLLS AND WIN32) - # Copy vcpkg DLLs (for gRPC/protobuf) - if(VCPKG_DLL_DIR AND EXISTS "${VCPKG_DLL_DIR}") - foreach(dll ${VCPKG_DLLS}) - if(EXISTS "${VCPKG_DLL_DIR}/${dll}") - add_custom_command(TARGET integ-test POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy_if_different - "${VCPKG_DLL_DIR}/${dll}" - $ - COMMENT "Copying ${dll} for test" - ) - endif() - endforeach() - endif() - - # Copy OpenCV DLLs - if(EXISTS "${OPENCV_DLL_DIR}") - foreach(dll ${OPENCV_DLLS}) - if(EXISTS "${OPENCV_DLL_DIR}/${dll}") - add_custom_command(TARGET integ-test POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy_if_different - "${OPENCV_DLL_DIR}/${dll}" - $ - COMMENT "Copying ${dll} for test" - ) - endif() - endforeach() - endif() - # Copy the main library DLL - if(EXISTS ${PLUGIN_DLLS}) - add_custom_command(TARGET integ-test POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy_if_different - "${PLUGIN_DLLS}" - $ - COMMENT "Copying ${dll} for test" - ) - endif() -endif() +# Runtime DLLs are copied once into the shared top-level runtime directory +# (see CMAKE_RUNTIME_OUTPUT_DIRECTORY in the top-level CMakeLists.txt), where +# this executable also lands — no per-test copying needed. if(INSTALL_TESTS) install(TARGETS integ-test diff --git a/tests/standalone-tests/CMakeLists.txt b/tests/standalone-tests/CMakeLists.txt index 7a9110a..5ccf95e 100644 --- a/tests/standalone-tests/CMakeLists.txt +++ b/tests/standalone-tests/CMakeLists.txt @@ -45,36 +45,9 @@ target_link_libraries(spot_camera_stream_test gRPC::grpc ) -# Copy only the DLLs needed for this test -if(COPY_DLLS AND WIN32) - # Copy vcpkg DLLs (for gRPC/protobuf) - if(VCPKG_DLL_DIR AND EXISTS "${VCPKG_DLL_DIR}") - foreach(dll ${VCPKG_DLLS}) - if(EXISTS "${VCPKG_DLL_DIR}/${dll}") - add_custom_command(TARGET spot_camera_stream_test POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy_if_different - "${VCPKG_DLL_DIR}/${dll}" - $ - COMMENT "Copying ${dll} for test" - ) - endif() - endforeach() - endif() - - # Copy OpenCV DLLs - if(EXISTS "${OPENCV_DLL_DIR}") - foreach(dll ${OPENCV_DLLS}) - if(EXISTS "${OPENCV_DLL_DIR}/${dll}") - add_custom_command(TARGET spot_camera_stream_test POST_BUILD - COMMAND ${CMAKE_COMMAND} -E copy_if_different - "${OPENCV_DLL_DIR}/${dll}" - $ - COMMENT "Copying ${dll} for test" - ) - endif() - endforeach() - endif() -endif() +# Runtime DLLs are copied once into the shared top-level runtime directory +# (see CMAKE_RUNTIME_OUTPUT_DIRECTORY in the top-level CMakeLists.txt), where +# this executable also lands — no per-test copying needed. if(INSTALL_TESTS) install(TARGETS spot_camera_stream_test From 7cfc928b8b087ca27e445a69746ff5e3c8bb1df9 Mon Sep 17 00:00:00 2001 From: Faisal Zaghloul Date: Fri, 10 Jul 2026 19:31:30 -0400 Subject: [PATCH 6/6] Update release --- install | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/install b/install index d402033..04c66a1 160000 --- a/install +++ b/install @@ -1 +1 @@ -Subproject commit d40203386a71030a87972f92be0879a4f8028ca0 +Subproject commit 04c66a1d9d1dd36a281b841e16c507ceefc133e6