diff --git a/tests/datasets/test_aggregate.py b/tests/datasets/test_aggregate.py index dfb594276..63d63e789 100644 --- a/tests/datasets/test_aggregate.py +++ b/tests/datasets/test_aggregate.py @@ -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 diff --git a/tests/datasets/test_dataset_tools.py b/tests/datasets/test_dataset_tools.py index d36312920..c9f8c7d91 100644 --- a/tests/datasets/test_dataset_tools.py +++ b/tests/datasets/test_dataset_tools.py @@ -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() diff --git a/tests/fixtures/constants.py b/tests/fixtures/constants.py index 7835aa834..816deef32 100644 --- a/tests/fixtures/constants.py +++ b/tests/fixtures/constants.py @@ -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) diff --git a/tests/fixtures/dataset_factories.py b/tests/fixtures/dataset_factories.py index 5cdf473d9..2083cc2d6 100644 --- a/tests/fixtures/dataset_factories.py +++ b/tests/fixtures/dataset_factories.py @@ -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):