From 15d94e6108b0d7a1a6b26e96b817a2177f6d93aa Mon Sep 17 00:00:00 2001 From: CarolinePascal Date: Wed, 1 Jul 2026 20:00:20 +0200 Subject: [PATCH] feat(unit getter): adding a proper output_depth_unit getter to LeRobotDataset for cleaner integration --- src/lerobot/datasets/lerobot_dataset.py | 5 +++++ src/lerobot/scripts/lerobot_dataset_viz.py | 6 +++--- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/src/lerobot/datasets/lerobot_dataset.py b/src/lerobot/datasets/lerobot_dataset.py index 672b6958b..aba86efe3 100644 --- a/src/lerobot/datasets/lerobot_dataset.py +++ b/src/lerobot/datasets/lerobot_dataset.py @@ -351,6 +351,11 @@ class LeRobotDataset(torch.utils.data.Dataset): """Frames per second used during data collection.""" return self.meta.fps + @property + def depth_output_unit(self) -> str: + """Physical unit (``"m"`` or ``"mm"``) depth maps and statistics are returned in on read.""" + return self._depth_output_unit + @property def num_frames(self) -> int: """Number of frames in selected episodes.""" diff --git a/src/lerobot/scripts/lerobot_dataset_viz.py b/src/lerobot/scripts/lerobot_dataset_viz.py index 3957ca032..f4be878ad 100644 --- a/src/lerobot/scripts/lerobot_dataset_viz.py +++ b/src/lerobot/scripts/lerobot_dataset_viz.py @@ -84,7 +84,7 @@ import torch import torch.utils.data import tqdm -from lerobot.configs import DEFAULT_DEPTH_UNIT, DEPTH_MILLIMETER_UNIT +from lerobot.configs import 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 @@ -229,8 +229,8 @@ 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 + # Depth frames and stats are dequantized to the dataset's depth_output_unit on load. + depth_meter = 1000.0 if dataset.depth_output_unit == DEPTH_MILLIMETER_UNIT else 1.0 # Use the dataset's q01/q99 depth statistics for robust depth range bounds depth_ranges = {}