mirror of
https://github.com/huggingface/lerobot.git
synced 2026-07-29 12:39:41 +00:00
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:
@@ -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(
|
||||
|
||||
Reference in New Issue
Block a user