fix(depth unit): adding input depth unit storage in the dataset metadata (#3899)

* fix(depth unit): storing raw depth units in the dataset metadata for correct depth statistics and depth raw frames handling. The unit is stored as a string ("m","mm") under "depth_unit" at the same level as "is_depth_map". Unit is inferred from the depth frame type.

* feat(raw frame unit): adapting dataset reader so that raw depth frames are scaled according to the requested unit

* feat(stats units): rescaling stats when loading a dataset so that the stats are given in the requested unit

* tests(unit): adapting and extending depth tests to units manipulations

* chore(format): formating code

* feat(warning): adding a warning when depth unit is not specified in the dataset

* chore(infer_depth_unit): moving the depth unit inference utility in a more accessible location

* feat(rerun unit): adding correct depth unit display for rerun (foxglove does not support units yet)

* feat(unit getter): adding a proper output_depth_unit getter to LeRobotDataset for cleaner integration

* fix(streaming dataset): extending support for depth units to streaming datasets

* test(rerun): fixing rerun tests
This commit is contained in:
Caroline Pascal
2026-07-02 11:53:13 +02:00
committed by GitHub
parent b961d2a8c5
commit c746ca2df2
14 changed files with 229 additions and 19 deletions
+8
View File
@@ -26,6 +26,7 @@ import pytest
import torch
from datasets import Dataset
from lerobot.configs.video import infer_depth_unit
from lerobot.datasets.dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
from lerobot.datasets.feature_utils import get_hf_features_from_features
from lerobot.datasets.io_utils import flatten_dict, hf_transform_to_torch
@@ -535,6 +536,13 @@ def lerobot_dataset_factory(
chunks_size=chunks_size,
**info_kwargs,
)
# This synthetic path skips add_frame, so record the depth unit the writer would
# have stored (dummy depth is uint16) to keep ``depth_unit`` present in info.json.
# Reassign a fresh info dict to avoid mutating the shared feature constants.
for ft in info.features.values():
ft_info = ft.get("info")
if ft_info is not None and ft_info.get("is_depth_map") and "depth_unit" not in ft_info:
ft["info"] = {**ft_info, "depth_unit": infer_depth_unit(np.uint16)}
if stats is None:
stats = stats_factory(features=info.features)
if tasks is None: