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
+10
View File
@@ -36,6 +36,7 @@ from lerobot.configs import (
RGBEncoderConfig,
VideoEncoderConfig,
depth_encoder_defaults,
infer_depth_unit,
rgb_encoder_defaults,
)
@@ -209,6 +210,15 @@ class DatasetWriter:
self.episode_buffer["timestamp"].append(timestamp)
self.episode_buffer["task"].append(frame.pop("task"))
# Record each depth feature's input unit once, inferred from the first frame's dtype.
if frame_index == 0:
for depth_key in self._meta.depth_keys:
if depth_key not in frame:
continue
info = self._meta.features[depth_key].setdefault("info", {})
if info.get("depth_unit") is None:
info["depth_unit"] = infer_depth_unit(np.asarray(frame[depth_key]).dtype)
# Start streaming encoder on first frame of episode
if frame_index == 0 and self._streaming_encoder is not None:
self._streaming_encoder.start_episode(