File size: 9,085 Bytes
0d80452 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 | #!/usr/bin/env python
# Copyright 2024 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Contract tests for DatasetWriter."""
from pathlib import Path
from unittest.mock import patch
import numpy as np
import pytest
import torch
from PIL import Image
pytest.importorskip("datasets", reason="datasets is required (install lerobot[dataset])")
from lerobot.configs import VideoEncoderConfig
from lerobot.datasets.dataset_writer import _encode_video_worker
from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.datasets.utils import DEFAULT_IMAGE_PATH
from tests.fixtures.constants import DEFAULT_FPS, DUMMY_REPO_ID
SIMPLE_FEATURES = {
"state": {"dtype": "float32", "shape": (6,), "names": None},
"action": {"dtype": "float32", "shape": (6,), "names": None},
}
def _make_frame(features: dict, task: str = "Dummy task") -> dict:
"""Build a valid frame dict for the given features."""
frame = {"task": task}
for key, ft in features.items():
if ft["dtype"] in ("image", "video"):
frame[key] = np.random.randint(0, 256, size=ft["shape"], dtype=np.uint8)
elif ft["dtype"] in ("float32", "float64"):
frame[key] = torch.randn(ft["shape"])
elif ft["dtype"] == "int64":
frame[key] = torch.zeros(ft["shape"], dtype=torch.int64)
return frame
# ── Existing encode_video_worker tests ───────────────────────────────
def test_encode_video_worker_forwards_video_encoder(tmp_path):
"""_encode_video_worker forwards video_encoder to encode_video_frames."""
video_key = "observation.images.laptop"
fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0)
img_dir = tmp_path / Path(fpath).parent
img_dir.mkdir(parents=True, exist_ok=True)
Image.new("RGB", (64, 64), color="red").save(img_dir / "frame-000000.png")
captured_kwargs = {}
def mock_encode(imgs_dir, video_path, fps, **kwargs):
captured_kwargs.update(kwargs)
Path(video_path).parent.mkdir(parents=True, exist_ok=True)
Path(video_path).touch()
with patch("lerobot.datasets.dataset_writer.encode_video_frames", side_effect=mock_encode):
_encode_video_worker(
video_key,
0,
tmp_path,
fps=30,
video_encoder=VideoEncoderConfig(vcodec="h264", preset=None),
encoder_threads=4,
)
assert captured_kwargs["video_encoder"].vcodec == "h264"
assert captured_kwargs["encoder_threads"] == 4
def test_encode_video_worker_default_video_encoder(tmp_path):
"""_encode_video_worker passes None video_encoder which encode_video_frames defaults."""
video_key = "observation.images.laptop"
fpath = DEFAULT_IMAGE_PATH.format(image_key=video_key, episode_index=0, frame_index=0)
img_dir = tmp_path / Path(fpath).parent
img_dir.mkdir(parents=True, exist_ok=True)
Image.new("RGB", (64, 64), color="red").save(img_dir / "frame-000000.png")
captured_kwargs = {}
def mock_encode(imgs_dir, video_path, fps, **kwargs):
captured_kwargs.update(kwargs)
Path(video_path).parent.mkdir(parents=True, exist_ok=True)
Path(video_path).touch()
with patch("lerobot.datasets.dataset_writer.encode_video_frames", side_effect=mock_encode):
_encode_video_worker(video_key, 0, tmp_path, fps=30)
assert captured_kwargs["video_encoder"] is None
assert captured_kwargs["encoder_threads"] is None
# ── add_frame contracts ──────────────────────────────────────────────
def test_add_frame_increments_buffer_size(tmp_path):
"""Each add_frame() call increases episode_buffer['size'] by 1."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
assert dataset.writer.episode_buffer["size"] == 0
dataset.add_frame(_make_frame(SIMPLE_FEATURES))
assert dataset.writer.episode_buffer["size"] == 1
dataset.add_frame(_make_frame(SIMPLE_FEATURES))
assert dataset.writer.episode_buffer["size"] == 2
def test_add_frame_rejects_missing_feature(tmp_path):
"""add_frame() raises ValueError when a required feature is missing."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
with pytest.raises(ValueError, match="Missing features"):
dataset.add_frame({"task": "Dummy task", "state": torch.randn(6)})
# missing 'action'
# ── save_episode contracts ───────────────────────────────────────────
def test_save_episode_writes_parquet(tmp_path):
"""After save_episode(), at least one .parquet file exists under data/."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
for _ in range(3):
dataset.add_frame(_make_frame(SIMPLE_FEATURES))
dataset.save_episode()
parquet_files = list((tmp_path / "ds" / "data").rglob("*.parquet"))
assert len(parquet_files) > 0
def test_save_episode_updates_counters(tmp_path):
"""After save_episode(), metadata counters are updated."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
for _ in range(5):
dataset.add_frame(_make_frame(SIMPLE_FEATURES))
dataset.save_episode()
assert dataset.meta.total_episodes == 1
assert dataset.meta.total_frames == 5
def test_save_episode_resets_buffer(tmp_path):
"""After save_episode(), the episode buffer is reset."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
for _ in range(3):
dataset.add_frame(_make_frame(SIMPLE_FEATURES))
dataset.save_episode()
assert dataset.writer.episode_buffer["size"] == 0
def test_save_multiple_episodes(tmp_path):
"""Recording 3 episodes results in correct total counts."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
total_frames = 0
for ep in range(3):
n_frames = ep + 2 # 2, 3, 4
for _ in range(n_frames):
dataset.add_frame(_make_frame(SIMPLE_FEATURES))
dataset.save_episode()
total_frames += n_frames
assert dataset.meta.total_episodes == 3
assert dataset.meta.total_frames == total_frames
# ── clear / lifecycle ────────────────────────────────────────────────
def test_clear_resets_buffer(tmp_path):
"""clear_episode_buffer() resets the buffer size to 0."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
dataset.add_frame(_make_frame(SIMPLE_FEATURES))
assert dataset.writer.episode_buffer["size"] == 1
dataset.clear_episode_buffer()
assert dataset.writer.episode_buffer["size"] == 0
def test_finalize_is_idempotent(tmp_path):
"""Calling finalize() twice does not raise."""
dataset = LeRobotDataset.create(
repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=SIMPLE_FEATURES, root=tmp_path / "ds"
)
for _ in range(3):
dataset.add_frame(_make_frame(SIMPLE_FEATURES))
dataset.save_episode()
dataset.finalize()
dataset.finalize() # second call should not raise
def test_finalize_then_read_roundtrip(tmp_path):
"""Write data, finalize, re-open, and verify data matches."""
root = tmp_path / "roundtrip"
features = {"state": {"dtype": "float32", "shape": (2,), "names": None}}
dataset = LeRobotDataset.create(repo_id=DUMMY_REPO_ID, fps=DEFAULT_FPS, features=features, root=root)
# Record known values
known_states = []
for i in range(5):
state = torch.tensor([float(i), float(i * 10)])
known_states.append(state)
dataset.add_frame({"task": "Test task", "state": state})
dataset.save_episode()
dataset.finalize()
# Read back
for i in range(5):
item = dataset[i]
assert torch.allclose(item["state"], known_states[i], atol=1e-5)
|