feat(depth): persist depth metadata

This commit is contained in:
CarolinePascal
2026-05-19 16:11:21 +02:00
parent 088352383d
commit b960524d93
2 changed files with 44 additions and 0 deletions
+31
View File
@@ -53,6 +53,16 @@ IMAGE_FEATURES = {
},
}
DEPTH_FEATURES = {
**SIMPLE_FEATURES,
"observation.depth.laptop": {
"dtype": "video",
"shape": (64, 96),
"names": ["height", "width"],
"info": {"video.is_depth_map": True},
},
}
def _make_dummy_stats(features: dict) -> dict:
"""Create minimal episode stats matching the given features."""
@@ -142,6 +152,27 @@ def test_create_without_videos_has_no_video_path(tmp_path):
assert meta.video_keys == []
def test_depth_keys_property_filters_by_marker(tmp_path):
"""``depth_keys`` selects only video features carrying ``video.is_depth_map=True``."""
features = {
**VIDEO_FEATURES,
**DEPTH_FEATURES,
}
meta = LeRobotDatasetMetadata.create(
repo_id="test/depth_keys", fps=DEFAULT_FPS, features=features, root=tmp_path / "depth_keys"
)
assert set(meta.video_keys) == {"observation.images.laptop", "observation.depth.laptop"}
assert meta.depth_keys == ["observation.depth.laptop"]
def test_depth_keys_empty_when_no_marker(tmp_path):
meta = LeRobotDatasetMetadata.create(
repo_id="test/no_depth", fps=DEFAULT_FPS, features=VIDEO_FEATURES, root=tmp_path / "no_depth"
)
assert meta.depth_keys == []
def test_create_raises_on_existing_directory(tmp_path):
"""create() raises if root directory already exists."""
root = tmp_path / "existing"