admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-06-17 08:17:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

test(cleaning): cleaning up tests

Browse Source
This commit is contained in:
CarolinePascal
2026-05-26 17:25:13 +02:00
parent ec6264d768
commit ca7168f413
4 changed files with 51 additions and 12 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/datasets/test_aggregate.py
+8 -7
View File
@@ -33,6 +33,7 @@ from tests.fixtures.constants import (
DUMMY_CAMERA_FEATURES,
DUMMY_DEPTH_CAMERA_FEATURES,
DUMMY_REPO_ID,
DUMMY_CAMERA_FEATURES_WITH_DEPTH,
)
CAMERA_FEATURES_WITH_DEPTH = {**DUMMY_CAMERA_FEATURES, **DUMMY_DEPTH_CAMERA_FEATURES}
@@ -304,14 +305,14 @@ def test_aggregate_datasets(tmp_path, lerobot_dataset_factory):
repo_id=f"{DUMMY_REPO_ID}_0",
total_episodes=ds_0_num_episodes,
total_frames=ds_0_num_frames,
camera_features=CAMERA_FEATURES_WITH_DEPTH,
camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
)
ds_1 = lerobot_dataset_factory(
root=tmp_path / "test_1",
repo_id=f"{DUMMY_REPO_ID}_1",
total_episodes=ds_1_num_episodes,
total_frames=ds_1_num_frames,
camera_features=CAMERA_FEATURES_WITH_DEPTH,
camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
)
# Confirm depth was actually wired into the source datasets so the
@@ -476,14 +477,14 @@ def test_aggregate_with_low_threshold(tmp_path, lerobot_dataset_factory):
repo_id=f"{DUMMY_REPO_ID}_small_0",
total_episodes=ds_0_num_episodes,
total_frames=ds_0_num_frames,
camera_features=CAMERA_FEATURES_WITH_DEPTH,
camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
)
ds_1 = lerobot_dataset_factory(
root=tmp_path / "small_1",
repo_id=f"{DUMMY_REPO_ID}_small_1",
total_episodes=ds_1_num_episodes,
total_frames=ds_1_num_frames,
camera_features=CAMERA_FEATURES_WITH_DEPTH,
camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
)
assert len(ds_0.meta.depth_keys) > 0, "ds_0 should expose at least one depth key"
@@ -549,7 +550,7 @@ def test_video_timestamps_regression(tmp_path, lerobot_dataset_factory):
repo_id=f"{DUMMY_REPO_ID}_regression_{i}",
total_episodes=2,
total_frames=100,
camera_features=CAMERA_FEATURES_WITH_DEPTH,
camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
)
datasets.append(ds)
@@ -711,7 +712,7 @@ def test_aggregate_image_datasets(tmp_path, lerobot_dataset_factory):
total_episodes=ds_0_num_episodes,
total_frames=ds_0_num_frames,
use_videos=False,
camera_features=CAMERA_FEATURES_WITH_DEPTH,
camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
)
ds_1 = lerobot_dataset_factory(
root=tmp_path / "image_1",
@@ -719,7 +720,7 @@ def test_aggregate_image_datasets(tmp_path, lerobot_dataset_factory):
total_episodes=ds_1_num_episodes,
total_frames=ds_1_num_frames,
use_videos=False,
camera_features=CAMERA_FEATURES_WITH_DEPTH,
camera_features=DUMMY_CAMERA_FEATURES_WITH_DEPTH,
)
# Verify source datasets have image keys
tests/datasets/test_dataset_tools.py
+4 -3
View File
@@ -37,7 +37,8 @@ from lerobot.datasets.dataset_tools import (
split_dataset,
)
from lerobot.datasets.io_utils import load_info
from tests.datasets.test_video_encoding import _add_frames, require_h264, require_libsvtav1
from tests.datasets.test_video_encoding import require_h264, require_libsvtav1
from tests.fixtures.dataset_factories import add_frames
@pytest.fixture
@@ -1350,9 +1351,9 @@ def test_reencode_dataset_multi_key_multiprocessing(
camera_encoder=initial_cfg,
)
_add_frames(dataset, num_frames=4)
add_frames(dataset, num_frames=4)
dataset.save_episode()
_add_frames(dataset, num_frames=4)
add_frames(dataset, num_frames=4)
dataset.save_episode()
dataset.finalize()
tests/fixtures/constants.py
Vendored
+1
View File
@@ -57,5 +57,6 @@ DUMMY_DEPTH_CAMERA_FEATURES = {
"info": DUMMY_DEPTH_VIDEO_INFO,
},
}
DUMMY_CAMERA_FEATURES_WITH_DEPTH = {**DUMMY_CAMERA_FEATURES, **DUMMY_DEPTH_CAMERA_FEATURES}
DUMMY_CHW = (3, 96, 128)
DUMMY_HWC = (96, 128, 3)
tests/fixtures/dataset_factories.py
Vendored
+38 -2
View File
@@ -38,8 +38,6 @@ from lerobot.datasets.utils import (
DEFAULT_VIDEO_PATH,
DatasetInfo,
)
from lerobot.datasets.video_utils import encode_video_frames
from lerobot.utils.constants import DEFAULT_FEATURES
from tests.fixtures.constants import (
DEFAULT_FPS,
DUMMY_CAMERA_FEATURES,
@@ -47,6 +45,44 @@ from tests.fixtures.constants import (
DUMMY_REPO_ID,
DUMMY_ROBOT_TYPE,
)
from lerobot.datasets.video_utils import encode_video_frames
from lerobot.utils.constants import DEFAULT_FEATURES
def add_frames(
dataset: LeRobotDataset, num_frames: int
) -> None:
"""Append ``num_frames`` synthetic frames to ``dataset``.
Generates per-feature payloads from ``dataset.meta``: uint16 depth ramps for
keys in ``dataset.meta.depth_keys``, uint8 random noise for video/image keys,
and float32 zeros for everything else. ``DEFAULT_FEATURES`` (timestamp,
frame_index, ...) are auto-populated by ``add_frame`` and skipped here.
"""
if video_keys is None:
video_keys = dataset.meta.video_keys
depth_keys = set(dataset.meta.depth_keys)
# Smooth gradient base reused per (H, W) to keep depth frames cheap to
# encode (HEVC Main 12 hates white noise).
_depth_base_cache: dict[tuple[int, int], np.ndarray] = {}
for i in range(num_frames):
frame: dict = {"task": "test"}
for key, ft in dataset.meta.features.items():
if key in DEFAULT_FEATURES:
continue
shape = ft["shape"]
if key in depth_keys:
h, w, _ = shape
base = _depth_base_cache.setdefault(
(h, w),
np.linspace(100.0, 10_000.0, h * w, dtype=np.float32).reshape(h, w, 1),
)
frame[key] = (base + 50.0 * i).clip(0, 65535).astype(np.uint16)
elif key in video_keys:
frame[key] = np.random.randint(0, 256, shape, dtype=np.uint8)
else:
frame[key] = np.zeros(shape, dtype=np.float32)
dataset.add_frame(frame)
class LeRobotDatasetFactory(Protocol):