admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-05-20 11:09:59 +00:00
Code Issues Packages Projects Releases Wiki Activity

test(reencode dataset): adding missing test for reencode dataset

Browse Source
This commit is contained in:
CarolinePascal
2026-05-17 23:23:10 +02:00
parent 5547757cea
commit 6e01006d94
2 changed files with 57 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/datasets/test_video_encoding.py
+15 -9
View File
@@ -348,16 +348,22 @@ def _read_feature_info(dataset: LeRobotDataset) -> dict:
return info["features"][VIDEO_KEY]["info"]
def _add_frames(dataset: LeRobotDataset, num_frames: int) -> None:
shape = dataset.meta.features[VIDEO_KEY]["shape"]
def _add_frames(dataset: LeRobotDataset, num_frames: int, video_keys: list[str] | None = None) -> None:
from lerobot.utils.constants import DEFAULT_FEATURES
if video_keys is None:
video_keys = dataset.meta.video_keys
for _ in range(num_frames):
dataset.add_frame(
{
VIDEO_KEY: np.random.randint(0, 256, shape, dtype=np.uint8),
"action": np.zeros(2, dtype=np.float32),
"task": "test",
}
)
frame: dict = {"task": "test"}
for key, ft in dataset.meta.features.items():
if key in DEFAULT_FEATURES:
continue
shape = ft["shape"]
if 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 TestGetVideoInfo: