diff --git a/CMakeLists.txt b/CMakeLists.txt index 6992205..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 @@ -90,8 +108,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 @@ -146,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.") - 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}") + # $ 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 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) @@ -285,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/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/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/camera_stream.py b/PySpotObserver/pyspotobserver/camera_stream.py index 538c95b..bd13a9e 100644 --- a/PySpotObserver/pyspotobserver/camera_stream.py +++ b/PySpotObserver/pyspotobserver/camera_stream.py @@ -20,6 +20,13 @@ from .color_correction import _ROBOT_CCMS from .config import CameraType, SpotConfig +from .depth_registration import ( + DepthRegistrationParams, + DepthRegistrationWorkspace, + create_registration_workspace, + extract_registration_params, + register_depth, +) from .stitch import ( STITCH_OUT_H, STITCH_OUT_W, @@ -53,7 +60,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 +156,11 @@ 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 + self._depth_reg_workspaces: list[DepthRegistrationWorkspace] | None = None + # Optional per-stream vision pipeline, imported lazily. self._vision_pipeline: VisionPipeline | None = None @@ -240,6 +253,8 @@ 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._depth_reg_workspaces = None self._frame_count = 0 self._error_count = 0 self._last_body_to_worlds = None @@ -828,10 +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 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=frame.depth_images[frame_i], + responses[depth_idx], is_depth=True, out_array=workspace.raw_depth + ) + register_depth( + workspace.raw_depth, + self._depth_reg_params[i], + out=frame.depth_images[frame_i], + workspace=workspace, ) # Update timestamps and camera extrinsics params @@ -870,6 +892,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 +904,8 @@ def _decode_initial_responses( raise SpotCamStreamError(f"Expected {expected_count} responses, got {len(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 @@ -885,7 +913,29 @@ 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)) + + 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) + 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(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/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/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..3d6c0b6 --- /dev/null +++ b/PySpotObserver/pyspotobserver/depth_registration.py @@ -0,0 +1,228 @@ +"""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) + + +@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") + 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, + workspace: DepthRegistrationWorkspace | None = None, +) -> 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. + 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}") + 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 + + 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 + + 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 + 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_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..70785d9 --- /dev/null +++ b/PySpotObserver/tests/test_depth_registration.py @@ -0,0 +1,175 @@ +""" +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, + create_registration_workspace, + extract_registration_params, + register_depth, +) + + +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) + 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) -> 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) + 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) -> 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)) + 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) -> 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 + 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) -> 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) + + register_depth(raw, params, out) + + np.testing.assert_allclose(out, 0.0) + + 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)) + 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, + *, + 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) -> 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) + + 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) -> 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) + 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/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/install b/install index 440ce5a..04c66a1 160000 --- a/install +++ b/install @@ -1 +1 @@ -Subproject commit 440ce5acb594f13b3ac758d14e0cd03be5341164 +Subproject commit 04c66a1d9d1dd36a281b841e16c507ceefc133e6 diff --git a/src/cuda_kernels.cu b/src/cuda_kernels.cu index abbdccb..d18e7f0 100644 --- a/src/cuda_kernels.cu +++ b/src/cuda_kernels.cu @@ -2,10 +2,11 @@ #include "cuda_kernels.cuh" #include +#include #include #include -#include -#include +//#include +//#include #include #include @@ -83,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); @@ -648,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__ @@ -923,6 +891,93 @@ 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. + +// 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 uint16_t* __restrict__ src, + float* __restrict__ dst, + DepthRegistrationParams p +) { + 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 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) { + int i = blockIdx.x * blockDim.x + threadIdx.x; + // 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 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; + const int n_src = params.src_width * params.src_height; + if (n_dst <= 0 || n_src <= 0) { + return cudaErrorInvalidValue; + } + constexpr int threads = 256; + + // 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; + + 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); + 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; @@ -1281,15 +1336,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, @@ -1299,55 +1355,69 @@ __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 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); - - 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 (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/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 8021ccc..f49aa27 100644 --- a/src/include/cuda_kernels.cuh +++ b/src/include/cuda_kernels.cuh @@ -8,6 +8,31 @@ 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]; + 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, 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 uint16_t* 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, @@ -31,15 +56,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 @@ -62,14 +78,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/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 625bae7..09769c7 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,14 @@ class ReaderWriterCBuf { float* depth_data_{nullptr}; float* cached_depth_{nullptr}; + // 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_; + std::vector n_elems_per_raw_depth_; + size_t raw_depth_scratch_size_{0}; + uint16_t* raw_depth_scratch_{nullptr}; + bool first_run_{true}; const size_t max_size_; // Maximum size of the queue @@ -60,7 +69,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, + std::vector n_elems_per_raw_depth ); /** @@ -113,6 +124,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/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 756f648..952c12a 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 @@ -32,6 +28,7 @@ class VisionPipeline { MLModel& model_; const SpotCamStream& spot_cam_stream_; bool first_run_{true}; + std::atomic depth_averaging_enabled_{false}; // Threading std::unique_ptr pipeline_thread_; @@ -80,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-connection.cpp b/src/spot-connection.cpp index d9811cd..71df593 100644 --- a/src/spot-connection.cpp +++ b/src/spot-connection.cpp @@ -12,7 +12,9 @@ #include #include +#include +#include #include #include #include @@ -89,6 +91,77 @@ 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; + } + + // 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 + 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 +232,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 +244,32 @@ 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, + std::vector 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_ = 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_.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"); @@ -199,6 +289,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_, raw_depth_scratch_size_ * sizeof(uint16_t)), + "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_); @@ -238,13 +337,13 @@ void ReaderWriterCBuf::push(const google::protobuf::RepeatedPtrField start_time = high_resolution_clock::now(); + // 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; // - // LogMessage("ReaderWriterCBuf::push: latency: {:.4f} ms", latency_ms); + // LogPerf("ReaderWriterCBuf::push: latency: {:.4f} ms", latency_ms); // start_time = end_time; // Reset start time for next push // Compute write pointer @@ -338,34 +437,61 @@ 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; - - size_t depth_size = depth_width * depth_height; - if (depth_size != n_elems_per_depth_) { - LogMessage("Depth size mismatch: expected {}, got {}", n_elems_per_depth_, depth_size); + int depth_width = img.cols(); + int depth_height = img.rows(); + + const DepthRegistrationParams& reg = depth_registration_[n_depths_written]; + const size_t expected_depth_size = n_elems_per_raw_depth_[n_depths_written]; + size_t depth_size = size_t(depth_width) * depth_height; + if (depth_size != expected_depth_size + || depth_width != reg.src_width || depth_height != reg.src_height) { + LogMessage("Raw depth size mismatch: expected {} ({}x{}), got {} ({}x{})", + expected_depth_size, reg.src_width, reg.src_height, + depth_size, depth_width, depth_height); throw std::runtime_error("ReaderWriterCBuf::push: Depth size mismatch"); } + if (img.data().size() != depth_size * sizeof(uint16_t)) { + LogMessage("Raw depth is not uncompressed u16: expected {} bytes, got {}", + depth_size * sizeof(uint16_t), img.data().size()); + throw std::runtime_error("ReaderWriterCBuf::push: Unexpected raw depth payload"); + } - LogMessage("Copying depth image of size {} bytes to write pointer at index {}, {:#x} rgb_write_ptr", - depth_size, write_idx, size_t(depth_write_ptr)); + LogMessage("Registering {}x{} raw depth into {}x{} RGB frame at index {}, {:#x} depth_write_ptr", + depth_width, depth_height, reg.dst_width, reg.dst_height, + write_idx, size_t(depth_write_ptr)); + // Stage the raw u16 depth on the GPU; the scale to meters and the + // reprojection into the RGB camera's image plane both happen in the + // registration kernel (client-side *_depth_in_visual_frame). checkCudaError( cudaMemcpyAsync( - depth_write_ptr, - cv_img.data, - depth_size * sizeof(float), + raw_depth_scratch_, + img.data().data(), + depth_size * sizeof(uint16_t), cudaMemcpyHostToDevice, cuda_stream_ ), "cudaMemcpyAsync DEPTH" ); + DumpDepthImageFromCuda( + raw_depth_scratch_, + depth_width, + depth_height, + camera_dump_subdirs_[n_depths_written] + "/raw-depth", + dump_frame_id, + reg.inv_depth_scale + ); + + checkCudaError( + register_depth_to_rgb(raw_depth_scratch_, depth_write_ptr, reg, cuda_stream_), + "register_depth_to_rgb" + ); + DumpDepthImageFromCuda( depth_write_ptr, - cv_img.cols, - cv_img.rows, + reg.dst_width, + reg.dst_height, camera_dump_subdirs_[n_depths_written] + "/input-depth", dump_frame_id ); @@ -454,12 +580,15 @@ static std::vector 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 +690,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 +881,61 @@ bool SpotCamStream::streamCameras(uint32_t cam_mask) { return false; } } - // Same thing for depth images + // 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); + 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]; + 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); + 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]), + 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 +948,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), + std::move(raw_depth_sizes) )) { LogMessage("SpotCamStream::streamCameras: Failed to initialize image buffer"); return false; 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) { 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/integ-test/integ-test.cpp b/tests/integ-test/integ-test.cpp index d346676..fbef4e0 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; @@ -165,13 +165,21 @@ int main(int argc, char* argv[]) { for (int32_t spot = 0; spot < 2; spot++) { int32_t spot_id = spot_ids[spot]; if (spot_id < 0) { - std::cout << "Skipping Spot " << spot << " as it is not connected." << std::endl; + static bool printed_once = false; + if (!printed_once) { + std::cout << "Skipping Spot " << spot << " as it is not connected." << std::endl; + printed_once = true; + } continue; } for (int32_t stream = 0; stream < cam_stream_ids[spot_id].size(); stream++) { int32_t cam_stream_id = cam_stream_ids[spot_id][stream]; if (cam_stream_id < 0) { - std::cout << "Skipping Spot " << spot << " Stream " << stream << " as it is not valid." << std::endl; + static bool printed_once = false; + if (!printed_once) { + std::cout << "Skipping Spot " << spot << " Stream " << stream << " as it is not valid." << std::endl; + printed_once = true; + } continue; } 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