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 = {}