diff --git a/src/lerobot/scripts/lerobot_dataset_viz.py b/src/lerobot/scripts/lerobot_dataset_viz.py index d705ef420..b02a670c1 100644 --- a/src/lerobot/scripts/lerobot_dataset_viz.py +++ b/src/lerobot/scripts/lerobot_dataset_viz.py @@ -123,6 +123,12 @@ def to_hwc_uint8_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray: return hwc_uint8_numpy +def to_hwc_float32_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray: + check_chw_float32(chw_float32_torch) + hwc_float32_numpy = chw_float32_torch.permute(1, 2, 0).numpy() + return hwc_float32_numpy + + def build_blueprint_from_dataset(dataset: LeRobotDataset): """Build a Rerun blueprint laying out camera images and time series for the given dataset. @@ -148,12 +154,6 @@ def build_blueprint_from_dataset(dataset: LeRobotDataset): return rrb.Blueprint(rrb.Grid(*views)) -def to_hwc_uint16_numpy(chw_float32_torch: torch.Tensor) -> np.ndarray: - check_chw_float32(chw_float32_torch) - hwc_uint16_numpy = chw_float32_torch.round().type(torch.uint16).permute(1, 2, 0).numpy() - return hwc_uint16_numpy - - def visualize_dataset( dataset: LeRobotDataset, episode_index: int, @@ -249,7 +249,7 @@ def visualize_dataset( # display each camera image (or depth map) for key in dataset.meta.camera_keys: if key in dataset.meta.depth_keys: - depth = to_hwc_uint16_numpy(batch[key][i]) + depth = to_hwc_float32_numpy(batch[key][i]) depth_entity = rr.DepthImage( depth, colormap=rr.components.Colormap.Viridis, diff --git a/src/lerobot/utils/foxglove_visualization.py b/src/lerobot/utils/foxglove_visualization.py index b4315203d..2b0bccd2f 100644 --- a/src/lerobot/utils/foxglove_visualization.py +++ b/src/lerobot/utils/foxglove_visualization.py @@ -182,15 +182,32 @@ def _log_foxglove_image( compress_images: bool, channels: dict | None = None, log_time: int | None = None, + depth_range: tuple[float, float] | None = None, ) -> None: """Log an image on a cached per-topic channel. - ``arr`` may be HWC or CHW (CHW is transposed to HWC) and any dtype; floating-point images are - assumed normalized to [0, 1] and scaled to uint8. With ``compress_images`` set, grayscale (1ch) - and color (3ch) frames are JPEG-encoded, while 4-channel (RGBA) frames are always sent raw. - ``channels`` is the per-topic channel cache to reuse (see :func:`_log_foxglove_scalars`). - ``log_time`` is the message time in nanoseconds; when ``None`` the server's receive time is used. - It is also written to the message header timestamp. + Encoding is chosen from the frame's channel count and dtype: + + - Single channel => depth map, logged uncompressed as ``16UC1`` (integer) or ``32FC1`` (float). + Plain ``uint8`` grayscale (with no ``depth_range``) is the exception and is sent as ``mono8``. + - Three channels => ``rgb8``. + - Four channels => ``rgba8``. + + When ``compress_images`` is set, ``mono8`` and ``rgb8`` frames are JPEG-encoded instead; depth + and ``rgba8`` frames are always sent raw. + + Args: + topic: Foxglove topic to log on. + frame_id: Frame id stamped on the message. + arr: Image as HWC or CHW (CHW is transposed to HWC), any dtype. Non-depth floating-point + images are assumed normalized to [0, 1] and scaled to uint8. + compress_images: JPEG-encode ``mono8`` and ``rgb8`` frames; ignored for depth and ``rgba8``. + channels: Per-topic channel cache to reuse (see :func:`_log_foxglove_scalars`). + log_time: Message time in nanoseconds, also written to the header timestamp; when ``None`` + the server's receive time is used. + depth_range: ``(lo, hi)`` bounds that normalize depth to ``[0, 1]`` ``32FC1`` so Foxglove's + default panel scaling matches rerun's ``depth_range`` contrast (a ``RawImage`` carries + no value range itself). """ from foxglove.channels import CompressedImageChannel, RawImageChannel @@ -205,11 +222,44 @@ def _log_foxglove_image( # Convert CHW -> HWC when needed (mirrors log_rerun_data). if arr.ndim == 3 and arr.shape[0] in (1, 3, 4) and arr.shape[-1] not in (1, 3, 4): arr = np.transpose(arr, (1, 2, 0)) + height, width = arr.shape[0], arr.shape[1] + n_channels = 1 if arr.ndim == 2 else arr.shape[2] + + # Single-channel => depth, unless it's plain uint8 grayscale (kept as mono8 below). Foxglove + # renders both 16UC1 (uint16) and 32FC1 (float32) depth with a colormap. + if n_channels == 1 and (depth_range is not None or arr.dtype != np.uint8): + depth = arr if arr.ndim == 2 else arr[..., 0] + if depth_range is not None: + lo, hi = depth_range + depth = depth.astype(np.float32) + depth = np.clip((depth - lo) / (hi - lo), 0.0, 1.0) if hi > lo else np.zeros_like(depth) + encoding, step = "32FC1", width * 4 + elif np.issubdtype(depth.dtype, np.floating): + encoding, step = "32FC1", width * 4 + else: + depth = np.clip(np.rint(depth), 0, 65535).astype("