From 0d52d371be819711d26b406c7b5d9de7feb6909d Mon Sep 17 00:00:00 2001 From: CarolinePascal Date: Wed, 1 Jul 2026 19:39:43 +0200 Subject: [PATCH] feat(rerun unit): adding correct depth unit display for rerun (foxglove does not support units yet) --- src/lerobot/scripts/lerobot_dataset_viz.py | 5 +++++ src/lerobot/utils/rerun_visualization.py | 9 ++++++++- 2 files changed, 13 insertions(+), 1 deletion(-) diff --git a/src/lerobot/scripts/lerobot_dataset_viz.py b/src/lerobot/scripts/lerobot_dataset_viz.py index ee25583a0..3957ca032 100644 --- a/src/lerobot/scripts/lerobot_dataset_viz.py +++ b/src/lerobot/scripts/lerobot_dataset_viz.py @@ -84,6 +84,7 @@ import torch import torch.utils.data import tqdm +from lerobot.configs import DEFAULT_DEPTH_UNIT, DEPTH_MILLIMETER_UNIT from lerobot.datasets import LeRobotDataset from lerobot.utils.constants import ACTION, DONE, OBS_STATE, REWARD, SUCCESS from lerobot.utils.utils import init_logging @@ -228,6 +229,9 @@ def visualize_dataset( logging.info("Logging to Rerun") + # In this script, depth frames and stats are dequantized to the default depth_output_unit on load. + depth_meter = 1000.0 if DEFAULT_DEPTH_UNIT == DEPTH_MILLIMETER_UNIT else 1.0 + # Use the dataset's q01/q99 depth statistics for robust depth range bounds depth_ranges = {} for key in dataset.meta.depth_keys: @@ -254,6 +258,7 @@ def visualize_dataset( depth = to_hwc_float32_numpy(batch[key][i]) depth_entity = rr.DepthImage( depth, + meter=depth_meter, colormap=rr.components.Colormap.Viridis, depth_range=depth_ranges.get(key), ) diff --git a/src/lerobot/utils/rerun_visualization.py b/src/lerobot/utils/rerun_visualization.py index af04b18f7..46f2c0b4b 100644 --- a/src/lerobot/utils/rerun_visualization.py +++ b/src/lerobot/utils/rerun_visualization.py @@ -24,6 +24,7 @@ import os import numpy as np +from lerobot.configs import DEPTH_MILLIMETER_UNIT, infer_depth_unit from lerobot.types import RobotAction, RobotObservation from .constants import ACTION, ACTION_PREFIX, OBS_PREFIX, OBS_STR @@ -161,7 +162,13 @@ def log_rerun_data( observation_paths.add(key) else: if arr.shape[-1] == 1: - img_entity = rr.DepthImage(arr, colormap=rr.components.Colormap.Viridis) + # At record time, the depth unit is inferred from the frame type. + depth_unit = infer_depth_unit(arr.dtype) + img_entity = rr.DepthImage( + arr, + meter=1000.0 if depth_unit == DEPTH_MILLIMETER_UNIT else 1.0, + colormap=rr.components.Colormap.Viridis, + ) else: img_entity = rr.Image(arr).compress() if compress_images else rr.Image(arr) rr.log(key, entity=img_entity, static=True)