mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-23 12:40:08 +00:00
feat(depth): persist depth metadata
This commit is contained in:
@@ -338,6 +338,19 @@ class LeRobotDatasetMetadata:
|
|||||||
"""Keys to access visual modalities stored as videos."""
|
"""Keys to access visual modalities stored as videos."""
|
||||||
return [key for key, ft in self.features.items() if ft["dtype"] == "video"]
|
return [key for key, ft in self.features.items() if ft["dtype"] == "video"]
|
||||||
|
|
||||||
|
@property
|
||||||
|
def depth_keys(self) -> list[str]:
|
||||||
|
"""Keys to access depth-map modalities stored as videos or images.
|
||||||
|
|
||||||
|
A depth key is a feature whose ``info`` dict carries ``"<dtype>.is_depth_map": True``.
|
||||||
|
"""
|
||||||
|
return [
|
||||||
|
key
|
||||||
|
for key, ft in self.features.items()
|
||||||
|
# TODO(CarolinePascal): Make sure the legacy video_info works here as well.
|
||||||
|
if (ft.get("info") or {}).get(ft["dtype"] + ".is_depth_map", False)
|
||||||
|
]
|
||||||
|
|
||||||
@property
|
@property
|
||||||
def camera_keys(self) -> list[str]:
|
def camera_keys(self) -> list[str]:
|
||||||
"""Keys to access visual modalities (regardless of their storage method)."""
|
"""Keys to access visual modalities (regardless of their storage method)."""
|
||||||
|
|||||||
@@ -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:
|
def _make_dummy_stats(features: dict) -> dict:
|
||||||
"""Create minimal episode stats matching the given features."""
|
"""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 == []
|
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):
|
def test_create_raises_on_existing_directory(tmp_path):
|
||||||
"""create() raises if root directory already exists."""
|
"""create() raises if root directory already exists."""
|
||||||
root = tmp_path / "existing"
|
root = tmp_path / "existing"
|
||||||
|
|||||||
Reference in New Issue
Block a user