mirror of
https://github.com/huggingface/lerobot.git
synced 2026-08-08 17:39:44 +00:00
Compare commits
15 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 77a76ed24b | |||
| ced63c41d8 | |||
| 5e0421c1e6 | |||
| daa808cb42 | |||
| caec113349 | |||
| e346b1573f | |||
| 6eaea07046 | |||
| ef71bc8ccb | |||
| 96d5c3258a | |||
| d7a460d80c | |||
| 112eb1ed1b | |||
| f84495d206 | |||
| c01f3ffb28 | |||
| e2804c9fbd | |||
| e0226b23c8 |
@@ -39,7 +39,6 @@ from .io_utils import (
|
||||
hf_transform_to_torch,
|
||||
load_nested_dataset,
|
||||
)
|
||||
from .utils import resolve_episode_indices
|
||||
from .video_utils import decode_video_frames
|
||||
|
||||
|
||||
@@ -69,8 +68,9 @@ class DatasetReader:
|
||||
Args:
|
||||
meta: Dataset metadata instance.
|
||||
root: Local dataset root directory.
|
||||
episodes: Optional list of episode indices to select. ``None``
|
||||
means all episodes.
|
||||
episodes: Optional list of episode indices to select, assumed
|
||||
already validated by the caller. ``None`` means
|
||||
all episodes.
|
||||
tolerance_s: Timestamp synchronization tolerance in seconds.
|
||||
video_backend: Video decoding backend identifier.
|
||||
delta_timestamps: Optional dict mapping feature keys to lists of
|
||||
@@ -84,7 +84,7 @@ class DatasetReader:
|
||||
"""
|
||||
self._meta = meta
|
||||
self.root = root
|
||||
self.episodes = resolve_episode_indices(episodes, meta.total_episodes)
|
||||
self.episodes = episodes
|
||||
self._tolerance_s = tolerance_s
|
||||
self._video_backend = video_backend
|
||||
if image_transforms is not None and not callable(image_transforms):
|
||||
@@ -152,9 +152,9 @@ class DatasetReader:
|
||||
@property
|
||||
def num_frames(self) -> int:
|
||||
"""Number of frames in selected episodes."""
|
||||
if self.episodes is not None and self.hf_dataset is not None:
|
||||
return len(self.hf_dataset)
|
||||
return self._meta.total_frames
|
||||
if self.episodes is None:
|
||||
return self._meta.total_frames
|
||||
return sum(self._meta.episodes[ep]["length"] for ep in self.episodes)
|
||||
|
||||
@property
|
||||
def num_episodes(self) -> int:
|
||||
@@ -240,7 +240,12 @@ class DatasetReader:
|
||||
def _get_query_indices(
|
||||
self, abs_idx: int, ep_idx: int
|
||||
) -> tuple[dict[str, list[int]], dict[str, torch.Tensor]]:
|
||||
"""Compute query indices for delta timestamps."""
|
||||
"""Compute query indices for delta timestamps.
|
||||
|
||||
A delta is padding when ``abs_idx + delta`` falls outside the episode's
|
||||
``[dataset_from_index, dataset_to_index)`` range, and is clamped back into it
|
||||
otherwise.
|
||||
"""
|
||||
ep = self._meta.episodes[ep_idx]
|
||||
ep_start = ep["dataset_from_index"]
|
||||
ep_end = ep["dataset_to_index"]
|
||||
@@ -250,7 +255,7 @@ class DatasetReader:
|
||||
}
|
||||
padding = {
|
||||
f"{key}_is_pad": torch.BoolTensor(
|
||||
[(abs_idx + delta < ep_start) | (abs_idx + delta >= ep_end) for delta in delta_idx]
|
||||
[(abs_idx + delta < ep_start) or (abs_idx + delta >= ep_end) for delta in delta_idx]
|
||||
)
|
||||
for key, delta_idx in self.delta_indices.items()
|
||||
}
|
||||
|
||||
@@ -91,6 +91,7 @@ def make_dataset(cfg: TrainPipelineConfig) -> LeRobotDataset | MultiLeRobotDatas
|
||||
repo_type=cfg.dataset.repo_type,
|
||||
)
|
||||
delta_timestamps = resolve_delta_timestamps(cfg.trainable_config, ds_meta)
|
||||
# Resolved here because dataset reader does not apply exclude_episodes.
|
||||
episodes = resolve_episode_indices(
|
||||
cfg.dataset.episodes, ds_meta.total_episodes, cfg.dataset.exclude_episodes
|
||||
)
|
||||
|
||||
@@ -34,6 +34,7 @@ from .utils import (
|
||||
create_lerobot_dataset_card,
|
||||
get_safe_version,
|
||||
is_valid_version,
|
||||
resolve_episode_indices,
|
||||
)
|
||||
from .video_utils import (
|
||||
StreamingVideoEncoder,
|
||||
@@ -237,13 +238,12 @@ class LeRobotDataset(torch.utils.data.Dataset):
|
||||
self.revision = self.meta.revision
|
||||
self.meta.rescale_depth_stats(self._depth_output_unit)
|
||||
|
||||
if episodes is not None and any(
|
||||
episode >= self.meta.total_episodes or episode < 0 for episode in episodes
|
||||
):
|
||||
logger.warning(
|
||||
f"Some episodes in the provided episodes list are out of range for this dataset ({self.meta.total_episodes})."
|
||||
episodes = resolve_episode_indices(episodes, self.meta.total_episodes)
|
||||
if episodes is not None and not episodes:
|
||||
raise ValueError(
|
||||
"No valid episodes: the requested episode selection is empty after resolving "
|
||||
f"against the dataset range [0, {self.meta.total_episodes})."
|
||||
)
|
||||
|
||||
if episode_filter is not None:
|
||||
resolved = self.meta.filter_episodes(episode_filter, candidates=episodes)
|
||||
if not resolved:
|
||||
|
||||
@@ -13,8 +13,10 @@
|
||||
# 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.
|
||||
import logging
|
||||
from collections import deque
|
||||
from collections.abc import Callable, Generator, Iterable, Iterator
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
@@ -25,6 +27,7 @@ from datasets import load_dataset
|
||||
|
||||
from lerobot.configs import DEFAULT_DEPTH_UNIT, DEPTH_METER_UNIT, DepthEncoderConfig
|
||||
from lerobot.utils.constants import HF_LEROBOT_HOME, LOOKAHEAD_BACKTRACKTABLE, LOOKBACK_BACKTRACKTABLE
|
||||
from lerobot.utils.import_utils import get_safe_default_video_backend
|
||||
|
||||
from .dataset_metadata import CODEBASE_VERSION, LeRobotDatasetMetadata
|
||||
from .depth_utils import MM_PER_METRE, dequantize_depth
|
||||
@@ -32,8 +35,7 @@ from .feature_utils import get_delta_indices
|
||||
from .io_utils import item_to_torch
|
||||
from .utils import (
|
||||
check_version_compatibility,
|
||||
find_float_index,
|
||||
is_float_in_list,
|
||||
resolve_episode_indices,
|
||||
safe_shard,
|
||||
)
|
||||
from .video_utils import (
|
||||
@@ -42,6 +44,8 @@ from .video_utils import (
|
||||
decode_video_frames_torchcodec,
|
||||
)
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class LookBackError(Exception):
|
||||
"""
|
||||
@@ -248,6 +252,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
||||
repo_id: str,
|
||||
root: str | Path | None = None,
|
||||
episodes: list[int] | None = None,
|
||||
episode_filter: Callable[[dict], bool] | None = None,
|
||||
image_transforms: Callable | None = None,
|
||||
delta_timestamps: dict[list[float]] | None = None,
|
||||
tolerance_s: float = 1e-4,
|
||||
@@ -274,6 +279,9 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
||||
When omitted, Hub metadata is resolved through the cache under ``$HF_LEROBOT_HOME/hub``.
|
||||
episodes (list[int] | None, optional): If specified, this will only load episodes specified by
|
||||
their episode_index in this list.
|
||||
episode_filter (Callable[[dict], bool] | None, optional): Predicate over per-episode
|
||||
metadata rows used to select episodes (e.g. ``lambda ep: ep["length"] >= 100``).
|
||||
Intersected with ``episodes`` when both are set. Defaults to None.
|
||||
image_transforms (Callable | None, optional): Transform to apply to image data.
|
||||
tolerance_s (float, optional): Tolerance in seconds for timestamp matching.
|
||||
revision (str, optional): Git revision id (branch name, tag, or commit hash).
|
||||
@@ -305,7 +313,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
||||
self.streaming_from_local = root is not None and self.repo_type == "dataset"
|
||||
|
||||
self.image_transforms = image_transforms
|
||||
self.episodes = episodes
|
||||
self.tolerance_s = tolerance_s
|
||||
self.revision = revision if revision else CODEBASE_VERSION
|
||||
self.seed = seed
|
||||
@@ -338,6 +345,40 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
||||
# Check version
|
||||
check_version_compatibility(self.repo_id, self.meta._version, CODEBASE_VERSION)
|
||||
|
||||
self.episodes = resolve_episode_indices(episodes, self.meta.total_episodes)
|
||||
if self.episodes is not None and not self.episodes:
|
||||
raise ValueError(
|
||||
"No valid episodes: the requested episode selection is empty after resolving "
|
||||
f"against the dataset range [0, {self.meta.total_episodes})."
|
||||
)
|
||||
if episode_filter is not None:
|
||||
resolved = self.meta.filter_episodes(episode_filter, candidates=self.episodes)
|
||||
if not resolved:
|
||||
raise ValueError(
|
||||
"The episode filter did not match any episode. Make sure the filter and episodes list are valid and compatible."
|
||||
)
|
||||
logger.info(f"The episode filter matched {len(resolved)} episode(s).")
|
||||
self.episodes = resolved
|
||||
|
||||
# RGB video decodes with torchcodec when available, otherwise pyav (local paths only).
|
||||
# Remote streaming (hf://) needs torchcodec's fsspec-backed reader and can't decode depth
|
||||
# at all (pyav can't read hf:// URLs).
|
||||
is_remote_stream = self.streaming and not self.streaming_from_local
|
||||
self._video_backend = get_safe_default_video_backend()
|
||||
rgb_video_keys = [key for key in self.meta.video_keys if key not in self.meta.depth_keys]
|
||||
if rgb_video_keys and is_remote_stream and self._video_backend != "torchcodec":
|
||||
raise RuntimeError(
|
||||
"Remote StreamingLeRobotDataset requires the 'torchcodec' backend to decode RGB video, "
|
||||
f"but it is not available on this platform (affected keys: {rgb_video_keys}). "
|
||||
"Stream from a local root or use the non-streaming LeRobotDataset instead."
|
||||
)
|
||||
if self.meta.depth_keys and is_remote_stream:
|
||||
raise NotImplementedError(
|
||||
f"Remote streaming of depth video ({self.meta.depth_keys}) is not supported: depth is "
|
||||
"decoded with pyav, which cannot read hf:// URLs. Stream from a local root or use the "
|
||||
"non-streaming LeRobotDataset."
|
||||
)
|
||||
|
||||
self._depth_encoder_configs: dict[str, DepthEncoderConfig] = {
|
||||
vid_key: DepthEncoderConfig.from_video_info(self.meta.features[vid_key].get("info"))
|
||||
for vid_key in self.meta.depth_keys
|
||||
@@ -379,15 +420,22 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
||||
**token_kwargs,
|
||||
)
|
||||
|
||||
# Streaming loads all shards, so episode selection is applied as a lazy stream filter.
|
||||
if self.episodes is not None:
|
||||
selected = set(self.episodes)
|
||||
self.hf_dataset = self.hf_dataset.filter(lambda x: x["episode_index"] in selected)
|
||||
|
||||
self.num_shards = min(self.hf_dataset.num_shards, max_num_shards)
|
||||
|
||||
@property
|
||||
def num_frames(self):
|
||||
return self.meta.total_frames
|
||||
if self.episodes is None:
|
||||
return self.meta.total_frames
|
||||
return sum(self.meta.episodes[ep]["length"] for ep in self.episodes)
|
||||
|
||||
@property
|
||||
def num_episodes(self):
|
||||
return self.meta.total_episodes
|
||||
return len(self.episodes) if self.episodes is not None else self.meta.total_episodes
|
||||
|
||||
@property
|
||||
def fps(self):
|
||||
@@ -478,49 +526,28 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
||||
lookback, lookahead = self._get_window_steps(self.delta_timestamps)
|
||||
return Backtrackable(dataset, history=lookback, lookahead=lookahead)
|
||||
|
||||
def _make_timestamps_from_indices(
|
||||
self, start_ts: float, indices: dict[str, list[int]] | None = None
|
||||
) -> dict[str, list[float]]:
|
||||
if indices is not None:
|
||||
return {
|
||||
key: (
|
||||
start_ts + torch.tensor(indices[key]) / self.fps
|
||||
).tolist() # NOTE: why not delta_timestamps directly?
|
||||
for key in self.delta_timestamps
|
||||
}
|
||||
else:
|
||||
return dict.fromkeys(self.meta.video_keys, [start_ts])
|
||||
def _get_query_indices(
|
||||
self, abs_idx: int, ep_idx: int
|
||||
) -> tuple[dict[str, list[int]], dict[str, torch.BoolTensor]]:
|
||||
"""Video-key query indices and padding from integer episode boundaries.
|
||||
|
||||
def _make_padding_camera_frame(self, camera_key: str):
|
||||
"""Variable-shape padding frame for given camera keys, given in (H, W, C)"""
|
||||
return torch.zeros(self.meta.info.features[camera_key]["shape"]).permute(-1, 0, 1)
|
||||
|
||||
def _get_video_frame_padding_mask(
|
||||
self,
|
||||
video_frames: dict[str, torch.Tensor],
|
||||
query_timestamps: dict[str, list[float]],
|
||||
original_timestamps: dict[str, list[float]],
|
||||
) -> dict[str, torch.BoolTensor]:
|
||||
padding_mask = {}
|
||||
|
||||
for video_key, timestamps in original_timestamps.items():
|
||||
if video_key not in video_frames:
|
||||
continue # only padding on video keys that are available
|
||||
frames = []
|
||||
mask = []
|
||||
padding_frame = self._make_padding_camera_frame(video_key)
|
||||
for ts in timestamps:
|
||||
if is_float_in_list(ts, query_timestamps[video_key]):
|
||||
idx = find_float_index(ts, query_timestamps[video_key])
|
||||
frames.append(video_frames[video_key][idx, :])
|
||||
mask.append(False)
|
||||
else:
|
||||
frames.append(padding_frame)
|
||||
mask.append(True)
|
||||
|
||||
padding_mask[f"{video_key}_is_pad"] = torch.BoolTensor(mask)
|
||||
|
||||
return padding_mask
|
||||
Mirrors ``DatasetReader._get_query_indices`` but only for video keys.
|
||||
"""
|
||||
ep = self.meta.episodes[ep_idx]
|
||||
ep_start, ep_end = ep["dataset_from_index"], ep["dataset_to_index"]
|
||||
query_indices = {
|
||||
key: [max(ep_start, min(ep_end - 1, abs_idx + delta)) for delta in delta_idx]
|
||||
for key, delta_idx in self.delta_indices.items()
|
||||
if key in self.meta.video_keys
|
||||
}
|
||||
padding = {
|
||||
f"{key}_is_pad": torch.BoolTensor(
|
||||
[(abs_idx + delta < ep_start) or (abs_idx + delta >= ep_end) for delta in delta_idx]
|
||||
)
|
||||
for key, delta_idx in self.delta_indices.items()
|
||||
if key in self.meta.video_keys
|
||||
}
|
||||
return query_indices, padding
|
||||
|
||||
def make_frame(self, dataset_iterator: Backtrackable) -> Generator:
|
||||
"""Makes a frame starting from a dataset iterator"""
|
||||
@@ -533,19 +560,11 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
||||
|
||||
updates = [] # list of "updates" to apply to the item retrieved from hf_dataset (w/o camera features)
|
||||
|
||||
# Get episode index from the item
|
||||
# Get episode index and absolute frame index from the item
|
||||
ep_idx = item["episode_index"]
|
||||
|
||||
# "timestamp" restarts from 0 for each episode, whereas we need a global timestep within the single .mp4 file (given by index/fps)
|
||||
current_ts = item["index"] / self.fps
|
||||
|
||||
episode_boundaries_ts = {
|
||||
key: (
|
||||
self.meta.episodes[ep_idx][f"videos/{key}/from_timestamp"],
|
||||
self.meta.episodes[ep_idx][f"videos/{key}/to_timestamp"],
|
||||
)
|
||||
for key in self.meta.video_keys
|
||||
}
|
||||
abs_idx = int(item["index"])
|
||||
ep_start = self.meta.episodes[ep_idx]["dataset_from_index"]
|
||||
current_ts = float(item["timestamp"])
|
||||
|
||||
# Apply delta querying logic if necessary
|
||||
if self.delta_indices is not None:
|
||||
@@ -555,12 +574,19 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
||||
|
||||
# Load video frames, when needed
|
||||
if len(self.meta.video_keys) > 0:
|
||||
original_timestamps = self._make_timestamps_from_indices(current_ts, self.delta_indices)
|
||||
query_indices = None
|
||||
if self.delta_indices is not None:
|
||||
query_indices, video_padding = self._get_query_indices(abs_idx, ep_idx)
|
||||
|
||||
# Some timestamps might not result available considering the episode's boundaries
|
||||
query_timestamps = self._get_query_timestamps(
|
||||
current_ts, self.delta_indices, episode_boundaries_ts
|
||||
)
|
||||
# Episode-local timestamps; `_query_videos` shifts them by the per-key `from_timestamp` at decode.
|
||||
query_timestamps = {
|
||||
key: (
|
||||
[(idx - ep_start) / self.fps for idx in query_indices[key]]
|
||||
if query_indices is not None and key in query_indices
|
||||
else [current_ts]
|
||||
)
|
||||
for key in self.meta.video_keys
|
||||
}
|
||||
video_frames = self._query_videos(query_timestamps, ep_idx)
|
||||
|
||||
if self.image_transforms is not None:
|
||||
@@ -572,10 +598,7 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
||||
|
||||
if self.delta_indices is not None:
|
||||
# We always return the same number of frames. Unavailable frames are padded.
|
||||
padding_mask = self._get_video_frame_padding_mask(
|
||||
video_frames, query_timestamps, original_timestamps
|
||||
)
|
||||
updates.append(padding_mask)
|
||||
updates.append(video_padding)
|
||||
|
||||
result = item.copy()
|
||||
for update in updates:
|
||||
@@ -594,27 +617,6 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
||||
|
||||
yield result
|
||||
|
||||
def _get_query_timestamps(
|
||||
self,
|
||||
current_ts: float,
|
||||
query_indices: dict[str, list[int]] | None = None,
|
||||
episode_boundaries_ts: dict[str, tuple[float, float]] | None = None,
|
||||
) -> dict[str, list[float]]:
|
||||
query_timestamps = {}
|
||||
keys_to_timestamps = self._make_timestamps_from_indices(current_ts, query_indices)
|
||||
for key in self.meta.video_keys:
|
||||
if query_indices is not None and key in query_indices:
|
||||
timestamps = keys_to_timestamps[key]
|
||||
# Clamp out timesteps outside of episode boundaries
|
||||
query_timestamps[key] = torch.clamp(
|
||||
torch.tensor(timestamps), *episode_boundaries_ts[key]
|
||||
).tolist()
|
||||
|
||||
else:
|
||||
query_timestamps[key] = [current_ts]
|
||||
|
||||
return query_timestamps
|
||||
|
||||
def _query_videos(self, query_timestamps: dict[str, list[float]], ep_idx: int) -> dict:
|
||||
"""Note: When using data workers (e.g. DataLoader with num_workers>0), do not call this function
|
||||
in the main process (e.g. by using a second Dataloader with num_workers=0). It will result in a
|
||||
@@ -622,22 +624,28 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
||||
the main process and a subprocess fails to access it.
|
||||
"""
|
||||
|
||||
item = {}
|
||||
for video_key, query_ts in query_timestamps.items():
|
||||
root = self.meta.url_root if self.streaming and not self.streaming_from_local else self.root
|
||||
video_path = f"{root}/{self.meta.get_video_file_path(ep_idx, video_key)}"
|
||||
if video_key in self.meta.depth_keys:
|
||||
ep = self.meta.episodes[ep_idx]
|
||||
root = self.meta.url_root if self.streaming and not self.streaming_from_local else self.root
|
||||
|
||||
def _decode_single(vid_key: str, query_ts: list[float]) -> tuple[str, torch.Tensor]:
|
||||
# Episode-local timestamps restart from 0 each episode; shift by the per-key
|
||||
# `from_timestamp` to reach the episode's segment within its video file, matching
|
||||
# `DatasetReader._decode_single`.
|
||||
from_timestamp = ep[f"videos/{vid_key}/from_timestamp"]
|
||||
shifted_query_ts = [from_timestamp + ts for ts in query_ts]
|
||||
video_path = f"{root}/{self.meta.get_video_file_path(ep_idx, vid_key)}"
|
||||
if vid_key in self.meta.depth_keys:
|
||||
# Depth maps are 12-bit quantized and only decodable via pyav; dequantize back
|
||||
# to physical units to match the non-streaming reader.
|
||||
frames = decode_video_frames(
|
||||
video_path,
|
||||
query_ts,
|
||||
shifted_query_ts,
|
||||
self.tolerance_s,
|
||||
backend="pyav",
|
||||
return_uint8=False,
|
||||
is_depth=True,
|
||||
)
|
||||
depth_encoder = self._depth_encoder_configs[video_key]
|
||||
depth_encoder = self._depth_encoder_configs[vid_key]
|
||||
frames = dequantize_depth(
|
||||
frames,
|
||||
depth_min=depth_encoder.depth_min,
|
||||
@@ -646,18 +654,36 @@ class StreamingLeRobotDataset(torch.utils.data.IterableDataset):
|
||||
use_log=depth_encoder.use_log,
|
||||
output_unit=self._depth_output_unit,
|
||||
)
|
||||
else:
|
||||
elif self._video_backend == "torchcodec":
|
||||
frames = decode_video_frames_torchcodec(
|
||||
video_path,
|
||||
query_ts,
|
||||
shifted_query_ts,
|
||||
self.tolerance_s,
|
||||
decoder_cache=self.video_decoder_cache,
|
||||
return_uint8=self._return_uint8,
|
||||
)
|
||||
else:
|
||||
# torchcodec unavailable: only reachable for local streaming
|
||||
frames = decode_video_frames(
|
||||
video_path,
|
||||
shifted_query_ts,
|
||||
self.tolerance_s,
|
||||
backend="pyav",
|
||||
return_uint8=self._return_uint8,
|
||||
)
|
||||
|
||||
item[video_key] = frames.squeeze(0) if len(query_ts) == 1 else frames
|
||||
return vid_key, frames.squeeze(0) if len(shifted_query_ts) == 1 else frames
|
||||
|
||||
return item
|
||||
items = list(query_timestamps.items())
|
||||
|
||||
# Single camera: no threading overhead
|
||||
if len(items) <= 1:
|
||||
return {vid_key: _decode_single(vid_key, query_ts)[1] for vid_key, query_ts in items}
|
||||
|
||||
# Multi-camera: decode in parallel (video decoding releases the GIL)
|
||||
with ThreadPoolExecutor(max_workers=len(items)) as pool:
|
||||
futures = [pool.submit(_decode_single, k, ts) for k, ts in items]
|
||||
return dict(f.result() for f in futures)
|
||||
|
||||
def _get_delta_frames(self, dataset_iterator: Backtrackable, current_item: dict):
|
||||
# TODO(fracapuano): Modularize this function, refactor the code
|
||||
|
||||
@@ -523,17 +523,6 @@ def create_lerobot_dataset_card(
|
||||
)
|
||||
|
||||
|
||||
def is_float_in_list(target, float_list, threshold=1e-6):
|
||||
return any(abs(target - x) <= threshold for x in float_list)
|
||||
|
||||
|
||||
def find_float_index(target, float_list, threshold=1e-6):
|
||||
for i, x in enumerate(float_list):
|
||||
if abs(target - x) <= threshold:
|
||||
return i
|
||||
return -1
|
||||
|
||||
|
||||
def safe_shard(dataset: datasets.IterableDataset, index: int, num_shards: int) -> datasets.Dataset:
|
||||
"""
|
||||
Safe shards the dataset.
|
||||
|
||||
@@ -1828,25 +1828,21 @@ def test_episode_filter_intersects_with_episodes(tmp_path, lerobot_dataset_facto
|
||||
assert seen_eps == set(expected_eps)
|
||||
|
||||
|
||||
def test_episode_filter_no_match_raises(tmp_path, lerobot_dataset_factory):
|
||||
"""An empty match in LeRobotDataset's episode_filter raises a ValueError rather than silently returning an empty dataset."""
|
||||
@pytest.mark.parametrize(
|
||||
"kwargs, exc, match",
|
||||
[
|
||||
(
|
||||
{"episode_filter": lambda ep: ep["length"] < 0},
|
||||
ValueError,
|
||||
r"The episode filter did not match any episode",
|
||||
),
|
||||
({"episode_filter": lambda ep: ep["not_a_real_field"] > 0}, KeyError, "not_a_real_field"),
|
||||
({"episodes": [99]}, ValueError, "No valid episodes"),
|
||||
],
|
||||
)
|
||||
def test_episode_selection_invalid_raises(tmp_path, lerobot_dataset_factory, kwargs, exc, match):
|
||||
"""Invalid selections fail loudly: non-matching filter, unknown filter key, or out-of-range episodes."""
|
||||
dataset = lerobot_dataset_factory(root=tmp_path / "test", total_episodes=4, total_frames=100)
|
||||
|
||||
with pytest.raises(ValueError, match=r"The episode filter did not match any episode"):
|
||||
LeRobotDataset(
|
||||
dataset.repo_id,
|
||||
root=dataset.root,
|
||||
episode_filter=lambda ep: ep["length"] < 0,
|
||||
)
|
||||
|
||||
|
||||
def test_episode_filter_unknown_key_raises(tmp_path, lerobot_dataset_factory):
|
||||
"""A predicate referencing a column absent from meta.episodes surfaces a clear KeyError."""
|
||||
dataset = lerobot_dataset_factory(root=tmp_path / "test", total_episodes=4, total_frames=100)
|
||||
|
||||
with pytest.raises(KeyError, match="not_a_real_field"):
|
||||
LeRobotDataset(
|
||||
dataset.repo_id,
|
||||
root=dataset.root,
|
||||
episode_filter=lambda ep: ep["not_a_real_field"] > 0,
|
||||
)
|
||||
with pytest.raises(exc, match=match):
|
||||
LeRobotDataset(dataset.repo_id, root=dataset.root, **kwargs)
|
||||
|
||||
+202
-181
@@ -24,11 +24,17 @@ pytest.importorskip("datasets", reason="datasets is required (install lerobot[da
|
||||
|
||||
import lerobot.datasets.streaming_dataset as streaming_dataset_module
|
||||
from lerobot.datasets.dataset_metadata import LeRobotDatasetMetadata
|
||||
from lerobot.datasets.lerobot_dataset import LeRobotDataset
|
||||
from lerobot.datasets.streaming_dataset import StreamingLeRobotDataset
|
||||
from lerobot.datasets.utils import safe_shard
|
||||
from lerobot.utils.constants import ACTION
|
||||
from tests.fixtures.constants import DUMMY_REPO_ID
|
||||
|
||||
# A dataset whose videos roll over into a second file, as v3.0 does past
|
||||
# DEFAULT_VIDEO_FILE_SIZE_IN_MB: episodes 4-7 live in file-001, whose timeline restarts at 0.
|
||||
MULTI_FILE_EPISODES = 8
|
||||
MULTI_FILE_FRAMES = 200
|
||||
MULTI_FILE_EPISODES_PER_VIDEO_FILE = 4
|
||||
|
||||
|
||||
def get_frames_expected_order(streaming_ds: StreamingLeRobotDataset) -> list[int]:
|
||||
"""Replicates the shuffling logic of StreamingLeRobotDataset to get the expected order of indices."""
|
||||
@@ -80,14 +86,10 @@ def get_frames_expected_order(streaming_ds: StreamingLeRobotDataset) -> list[int
|
||||
@pytest.mark.parametrize("from_local", [False, True])
|
||||
def test_streaming_dataset_forwards_hub_token_only_for_remote_data(tmp_path, monkeypatch, token, from_local):
|
||||
requested_root = tmp_path / "local" if from_local else None
|
||||
metadata = SimpleNamespace(
|
||||
metadata = _fake_meta(
|
||||
root=requested_root or tmp_path / "snapshot",
|
||||
revision=streaming_dataset_module.CODEBASE_VERSION,
|
||||
_version=streaming_dataset_module.CODEBASE_VERSION,
|
||||
features={},
|
||||
depth_keys=[],
|
||||
image_keys=[],
|
||||
rescale_depth_stats=Mock(),
|
||||
total_episodes=10,
|
||||
)
|
||||
metadata_cls = Mock(return_value=metadata)
|
||||
load_dataset = Mock(return_value=SimpleNamespace(num_shards=1))
|
||||
@@ -111,8 +113,68 @@ def test_streaming_dataset_forwards_hub_token_only_for_remote_data(tmp_path, mon
|
||||
assert not hasattr(dataset, "_token")
|
||||
|
||||
|
||||
def assert_videos_roll_over(ds: LeRobotDataset) -> None:
|
||||
"""Videos spanning several files must decode each frame from its own file (v3.0 rollover)."""
|
||||
for key in ds.meta.video_keys:
|
||||
episodes = [ds.meta.episodes[ep_idx] for ep_idx in range(ds.meta.total_episodes)]
|
||||
file_indices = {ep[f"videos/{key}/file_index"] for ep in episodes}
|
||||
assert len(file_indices) > 1, f"{key} is not split across video files (file_index: {file_indices})"
|
||||
|
||||
# The property under test: a later file's timeline starts back at 0, so an episode's
|
||||
# `from_timestamp` is no longer its global position in the dataset.
|
||||
assert any(
|
||||
ep[f"videos/{key}/file_index"] > 0 and ep[f"videos/{key}/from_timestamp"] == 0.0
|
||||
for ep in episodes
|
||||
), f"No episode of {key} restarts a video file's timeline at 0"
|
||||
|
||||
|
||||
def assert_frame_matches(streaming_frame: dict, target_frame: dict, ds: LeRobotDataset, context: str) -> None:
|
||||
"""Assert a streamed frame equals the same frame read by the non-streaming reader."""
|
||||
assert set(streaming_frame.keys()) == set(target_frame.keys()), (
|
||||
f"Keys differ between streaming frame and target one ({context}). "
|
||||
f"Differ at: {set(streaming_frame.keys()) ^ set(target_frame.keys())}"
|
||||
)
|
||||
|
||||
mismatched = []
|
||||
for key in streaming_frame:
|
||||
left, right = streaming_frame[key], target_frame[key]
|
||||
|
||||
if isinstance(left, str):
|
||||
check = left == right
|
||||
|
||||
elif isinstance(left, float):
|
||||
check = left == right.item() # right is a torch.Tensor
|
||||
|
||||
elif isinstance(left, torch.Tensor):
|
||||
if key not in ds.meta.camera_keys and "is_pad" not in key and f"{key}_is_pad" in streaming_frame:
|
||||
# comparing frames only on non-padded regions. Padding is applied to last-valid broadcasting
|
||||
left = left[~streaming_frame[f"{key}_is_pad"]]
|
||||
right = right[~target_frame[f"{key}_is_pad"]]
|
||||
|
||||
check = left.shape == right.shape and torch.allclose(left, right)
|
||||
|
||||
else:
|
||||
check = left == right
|
||||
|
||||
if not check:
|
||||
mismatched.append(key)
|
||||
|
||||
assert not mismatched, f"Streaming and target frames differ on {mismatched} ({context})"
|
||||
|
||||
|
||||
def assert_stream_matches_reference(
|
||||
streaming_ds: StreamingLeRobotDataset, ds: LeRobotDataset, num_frames: int
|
||||
) -> None:
|
||||
"""Stream ``num_frames`` frames and assert each equals the same frame from the non-streaming reader."""
|
||||
stream = iter(streaming_ds)
|
||||
for i in range(num_frames):
|
||||
streaming_frame = next(stream)
|
||||
frame_idx = streaming_frame["index"]
|
||||
assert_frame_matches(streaming_frame, ds[frame_idx], ds, context=f"i: {i}, frame_idx: {frame_idx}")
|
||||
|
||||
|
||||
def test_single_frame_consistency(tmp_path, lerobot_dataset_factory):
|
||||
"""Test if are correctly accessed"""
|
||||
"""Streaming without deltas returns the same frames as the non-streaming reader."""
|
||||
ds_num_frames = 400
|
||||
ds_num_episodes = 10
|
||||
buffer_size = 100
|
||||
@@ -127,32 +189,83 @@ def test_single_frame_consistency(tmp_path, lerobot_dataset_factory):
|
||||
total_frames=ds_num_frames,
|
||||
)
|
||||
|
||||
streaming_ds = iter(StreamingLeRobotDataset(repo_id=repo_id, root=local_path, buffer_size=buffer_size))
|
||||
streaming_ds = StreamingLeRobotDataset(repo_id=repo_id, root=local_path, buffer_size=buffer_size)
|
||||
assert_stream_matches_reference(streaming_ds, ds, ds_num_frames)
|
||||
|
||||
key_checks = []
|
||||
for _ in range(ds_num_frames):
|
||||
streaming_frame = next(streaming_ds)
|
||||
frame_idx = streaming_frame["index"]
|
||||
target_frame = ds[frame_idx]
|
||||
|
||||
for key in streaming_frame:
|
||||
left = streaming_frame[key]
|
||||
right = target_frame[key]
|
||||
def test_streaming_episode_selection(tmp_path, lerobot_dataset_factory):
|
||||
"""episodes=[...] restricts both the streamed frames and the count properties to the selection."""
|
||||
ds_num_frames = 200
|
||||
ds_num_episodes = 10
|
||||
selected = [2, 5, 7]
|
||||
|
||||
if isinstance(left, str):
|
||||
check = left == right
|
||||
local_path = tmp_path / "test"
|
||||
repo_id = DUMMY_REPO_ID
|
||||
|
||||
elif isinstance(left, torch.Tensor):
|
||||
check = torch.allclose(left, right) and left.shape == right.shape
|
||||
ds = lerobot_dataset_factory(
|
||||
root=local_path,
|
||||
repo_id=repo_id,
|
||||
total_episodes=ds_num_episodes,
|
||||
total_frames=ds_num_frames,
|
||||
)
|
||||
|
||||
elif isinstance(left, float):
|
||||
check = left == right.item() # right is a torch.Tensor
|
||||
streaming_ds = StreamingLeRobotDataset(
|
||||
repo_id=repo_id, root=local_path, episodes=selected, buffer_size=50, shuffle=False
|
||||
)
|
||||
|
||||
key_checks.append((key, check))
|
||||
assert streaming_ds.num_episodes == len(selected)
|
||||
assert streaming_ds.num_frames == sum(ds.meta.episodes[ep]["length"] for ep in selected)
|
||||
|
||||
assert all(t[1] for t in key_checks), (
|
||||
f"Checking {list(filter(lambda t: not t[1], key_checks))[0][0]} left and right were found different (frame_idx: {frame_idx})"
|
||||
)
|
||||
episode_indices = [int(frame["episode_index"]) for frame in streaming_ds]
|
||||
assert set(episode_indices) == set(selected)
|
||||
assert len(episode_indices) == streaming_ds.num_frames
|
||||
|
||||
|
||||
def test_streaming_episode_filter(tmp_path, lerobot_dataset_factory):
|
||||
"""episode_filter restricts the stream to episodes whose metadata matches the predicate."""
|
||||
ds_num_episodes = 6
|
||||
ds_num_frames = 120
|
||||
|
||||
local_path = tmp_path / "test"
|
||||
repo_id = DUMMY_REPO_ID
|
||||
|
||||
lerobot_dataset_factory(
|
||||
root=local_path,
|
||||
repo_id=repo_id,
|
||||
total_episodes=ds_num_episodes,
|
||||
total_frames=ds_num_frames,
|
||||
)
|
||||
|
||||
keep = {1, 4}
|
||||
streaming_ds = StreamingLeRobotDataset(
|
||||
repo_id=repo_id,
|
||||
root=local_path,
|
||||
episode_filter=lambda ep: ep["episode_index"] in keep,
|
||||
buffer_size=50,
|
||||
shuffle=False,
|
||||
)
|
||||
|
||||
assert set(streaming_ds.episodes) == keep
|
||||
assert streaming_ds.num_episodes == len(keep)
|
||||
assert {int(frame["episode_index"]) for frame in streaming_ds} == keep
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"kwargs, match",
|
||||
[
|
||||
({"episodes": [99]}, "No valid episodes"),
|
||||
({"episode_filter": lambda ep: False}, "episode filter did not match"),
|
||||
],
|
||||
)
|
||||
def test_streaming_empty_selection_raises(tmp_path, lerobot_dataset_factory, kwargs, match):
|
||||
"""An empty selection (out-of-range episodes or a non-matching filter) must fail loudly."""
|
||||
local_path = tmp_path / "test"
|
||||
repo_id = DUMMY_REPO_ID
|
||||
|
||||
lerobot_dataset_factory(root=local_path, repo_id=repo_id, total_episodes=4, total_frames=40)
|
||||
|
||||
with pytest.raises(ValueError, match=match):
|
||||
StreamingLeRobotDataset(repo_id=repo_id, root=local_path, **kwargs)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
@@ -302,87 +415,7 @@ def test_iter_raises_on_nested_generator_error(tmp_path, lerobot_dataset_factory
|
||||
next(iter(streaming_ds))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"state_deltas, action_deltas",
|
||||
[
|
||||
([-1, -0.5, -0.20, 0], [0, 1, 2, 3]),
|
||||
([-1, -0.5, -0.20, 0], [-1.5, -1, -0.5, -0.20, -0.10, 0]),
|
||||
([-2, -1, -0.5, 0], [0, 1, 2, 3]),
|
||||
([-2, -1, -0.5, 0], [-1.5, -1, -0.5, -0.20, -0.10, 0]),
|
||||
],
|
||||
)
|
||||
def test_frames_with_delta_consistency(tmp_path, lerobot_dataset_factory, state_deltas, action_deltas):
|
||||
ds_num_frames = 500
|
||||
ds_num_episodes = 10
|
||||
buffer_size = 100
|
||||
|
||||
seed = 42
|
||||
|
||||
local_path = tmp_path / "test"
|
||||
repo_id = f"{DUMMY_REPO_ID}-ciao"
|
||||
camera_key = "phone"
|
||||
|
||||
delta_timestamps = {
|
||||
camera_key: state_deltas,
|
||||
"state": state_deltas,
|
||||
ACTION: action_deltas,
|
||||
}
|
||||
|
||||
ds = lerobot_dataset_factory(
|
||||
root=local_path,
|
||||
repo_id=repo_id,
|
||||
total_episodes=ds_num_episodes,
|
||||
total_frames=ds_num_frames,
|
||||
delta_timestamps=delta_timestamps,
|
||||
)
|
||||
|
||||
streaming_ds = iter(
|
||||
StreamingLeRobotDataset(
|
||||
repo_id=repo_id,
|
||||
root=local_path,
|
||||
buffer_size=buffer_size,
|
||||
seed=seed,
|
||||
shuffle=False,
|
||||
delta_timestamps=delta_timestamps,
|
||||
)
|
||||
)
|
||||
|
||||
for i in range(ds_num_frames):
|
||||
streaming_frame = next(streaming_ds)
|
||||
frame_idx = streaming_frame["index"]
|
||||
target_frame = ds[frame_idx]
|
||||
|
||||
assert set(streaming_frame.keys()) == set(target_frame.keys()), (
|
||||
f"Keys differ between streaming frame and target one. Differ at: {set(streaming_frame.keys()) - set(target_frame.keys())}"
|
||||
)
|
||||
|
||||
key_checks = []
|
||||
for key in streaming_frame:
|
||||
left = streaming_frame[key]
|
||||
right = target_frame[key]
|
||||
|
||||
if isinstance(left, str):
|
||||
check = left == right
|
||||
|
||||
elif isinstance(left, torch.Tensor):
|
||||
if (
|
||||
key not in ds.meta.camera_keys
|
||||
and "is_pad" not in key
|
||||
and f"{key}_is_pad" in streaming_frame
|
||||
):
|
||||
# comparing frames only on non-padded regions. Padding is applied to last-valid broadcasting
|
||||
left = left[~streaming_frame[f"{key}_is_pad"]]
|
||||
right = right[~target_frame[f"{key}_is_pad"]]
|
||||
|
||||
check = torch.allclose(left, right) and left.shape == right.shape
|
||||
|
||||
key_checks.append((key, check))
|
||||
|
||||
assert all(t[1] for t in key_checks), (
|
||||
f"Checking {list(filter(lambda t: not t[1], key_checks))[0][0]} left and right were found different (i: {i}, frame_idx: {frame_idx})"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sharded", [False, True])
|
||||
@pytest.mark.parametrize(
|
||||
"state_deltas, action_deltas",
|
||||
[
|
||||
@@ -392,94 +425,40 @@ def test_frames_with_delta_consistency(tmp_path, lerobot_dataset_factory, state_
|
||||
([-2, -1, -0.5, 0], [-20, -1.5, -1, -0.5, -0.20, -0.10, 0]),
|
||||
],
|
||||
)
|
||||
def test_frames_with_delta_consistency_with_shards(
|
||||
tmp_path, lerobot_dataset_factory, state_deltas, action_deltas
|
||||
def test_frames_with_delta_consistency(
|
||||
tmp_path, lerobot_dataset_factory, sharded, state_deltas, action_deltas
|
||||
):
|
||||
ds_num_frames = 100
|
||||
ds_num_episodes = 10
|
||||
buffer_size = 10
|
||||
data_file_size_mb = 0.001
|
||||
chunks_size = 1
|
||||
|
||||
seed = 42
|
||||
|
||||
"""Delta-window frames streamed match the non-streaming reader, with and without sharding."""
|
||||
local_path = tmp_path / "test"
|
||||
repo_id = f"{DUMMY_REPO_ID}-ciao"
|
||||
camera_key = "phone"
|
||||
delta_timestamps = {"phone": state_deltas, "state": state_deltas, ACTION: action_deltas}
|
||||
|
||||
delta_timestamps = {
|
||||
camera_key: state_deltas,
|
||||
"state": state_deltas,
|
||||
ACTION: action_deltas,
|
||||
}
|
||||
if sharded:
|
||||
num_frames, buffer_size = 100, 10
|
||||
factory_extra = {"data_files_size_in_mb": 0.001, "chunks_size": 1}
|
||||
stream_extra = {"max_num_shards": 4}
|
||||
else:
|
||||
num_frames, buffer_size = 500, 100
|
||||
factory_extra, stream_extra = {}, {}
|
||||
|
||||
ds = lerobot_dataset_factory(
|
||||
root=local_path,
|
||||
repo_id=repo_id,
|
||||
total_episodes=ds_num_episodes,
|
||||
total_frames=ds_num_frames,
|
||||
total_episodes=10,
|
||||
total_frames=num_frames,
|
||||
delta_timestamps=delta_timestamps,
|
||||
data_files_size_in_mb=data_file_size_mb,
|
||||
chunks_size=chunks_size,
|
||||
**factory_extra,
|
||||
)
|
||||
streaming_ds = StreamingLeRobotDataset(
|
||||
repo_id=repo_id,
|
||||
root=local_path,
|
||||
buffer_size=buffer_size,
|
||||
seed=seed,
|
||||
seed=42,
|
||||
shuffle=False,
|
||||
delta_timestamps=delta_timestamps,
|
||||
max_num_shards=4,
|
||||
**stream_extra,
|
||||
)
|
||||
|
||||
iter(streaming_ds)
|
||||
|
||||
num_shards = 4
|
||||
shards_indices = []
|
||||
for shard_idx in range(num_shards):
|
||||
shard = safe_shard(streaming_ds.hf_dataset, shard_idx, num_shards)
|
||||
shard_indices = [item["index"] for item in shard]
|
||||
shards_indices.append(shard_indices)
|
||||
|
||||
streaming_ds = iter(streaming_ds)
|
||||
|
||||
for i in range(ds_num_frames):
|
||||
streaming_frame = next(streaming_ds)
|
||||
frame_idx = streaming_frame["index"]
|
||||
target_frame = ds[frame_idx]
|
||||
|
||||
assert set(streaming_frame.keys()) == set(target_frame.keys()), (
|
||||
f"Keys differ between streaming frame and target one. Differ at: {set(streaming_frame.keys()) - set(target_frame.keys())}"
|
||||
)
|
||||
|
||||
key_checks = []
|
||||
for key in streaming_frame:
|
||||
left = streaming_frame[key]
|
||||
right = target_frame[key]
|
||||
|
||||
if isinstance(left, str):
|
||||
check = left == right
|
||||
|
||||
elif isinstance(left, torch.Tensor):
|
||||
if (
|
||||
key not in ds.meta.camera_keys
|
||||
and "is_pad" not in key
|
||||
and f"{key}_is_pad" in streaming_frame
|
||||
):
|
||||
# comparing frames only on non-padded regions. Padding is applied to last-valid broadcasting
|
||||
left = left[~streaming_frame[f"{key}_is_pad"]]
|
||||
right = right[~target_frame[f"{key}_is_pad"]]
|
||||
|
||||
check = torch.allclose(left, right) and left.shape == right.shape
|
||||
|
||||
elif isinstance(left, float):
|
||||
check = left == right.item() # right is a torch.Tensor
|
||||
|
||||
key_checks.append((key, check))
|
||||
|
||||
assert all(t[1] for t in key_checks), (
|
||||
f"Checking {list(filter(lambda t: not t[1], key_checks))[0][0]} left and right were found different (i: {i}, frame_idx: {frame_idx})"
|
||||
)
|
||||
assert_stream_matches_reference(streaming_ds, ds, num_frames)
|
||||
|
||||
|
||||
class _StopConstructionError(Exception):
|
||||
@@ -493,7 +472,10 @@ def _fake_meta(*args, **kwargs):
|
||||
revision = kwargs.get("revision", args[2] if len(args) > 2 else None)
|
||||
meta.root = root or "/tmp/_streaming_meta"
|
||||
meta.revision = revision or "v0"
|
||||
meta._version = "v3.0"
|
||||
meta._version = streaming_dataset_module.CODEBASE_VERSION
|
||||
meta.total_episodes = 10
|
||||
meta.features = {}
|
||||
meta.video_keys = []
|
||||
meta.depth_keys = []
|
||||
meta.image_keys = []
|
||||
meta.rescale_depth_stats = lambda *_a, **_k: None
|
||||
@@ -616,7 +598,7 @@ def test_repo_type_is_keyword_only_and_preserves_positional_episodes():
|
||||
patch("lerobot.datasets.streaming_dataset.check_version_compatibility"),
|
||||
patch(
|
||||
"lerobot.datasets.streaming_dataset.load_dataset",
|
||||
return_value=SimpleNamespace(num_shards=1),
|
||||
return_value=SimpleNamespace(num_shards=1, filter=lambda *a, **k: SimpleNamespace(num_shards=1)),
|
||||
),
|
||||
):
|
||||
dataset = StreamingLeRobotDataset(DUMMY_REPO_ID, None, episodes)
|
||||
@@ -645,3 +627,42 @@ def test_bucket_root_caches_metadata_without_switching_to_local_streaming(tmp_pa
|
||||
def test_invalid_repo_type_fails_before_io():
|
||||
with pytest.raises(ValueError, match="repo_type must be 'dataset' or 'bucket'"):
|
||||
StreamingLeRobotDataset(DUMMY_REPO_ID, repo_type="space")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"state_deltas, action_deltas",
|
||||
[
|
||||
(None, None),
|
||||
([-1, -0.5, -0.20, 0], [0, 1, 2, 3]),
|
||||
([-2, -1, -0.5, 0], [-1.5, -1, -0.5, -0.20, -0.10, 0]),
|
||||
],
|
||||
)
|
||||
def test_consistency_across_video_files(tmp_path, lerobot_dataset_factory, state_deltas, action_deltas):
|
||||
"""Videos spanning several files must decode each frame from its own file (v3.0 rollover)."""
|
||||
local_path = tmp_path / "test"
|
||||
repo_id = f"{DUMMY_REPO_ID}-video-rollover"
|
||||
delta_timestamps = (
|
||||
None
|
||||
if state_deltas is None
|
||||
else {"phone": state_deltas, "state": state_deltas, ACTION: action_deltas}
|
||||
)
|
||||
|
||||
ds = lerobot_dataset_factory(
|
||||
root=local_path,
|
||||
repo_id=repo_id,
|
||||
total_episodes=MULTI_FILE_EPISODES,
|
||||
total_frames=MULTI_FILE_FRAMES,
|
||||
episodes_per_video_file=MULTI_FILE_EPISODES_PER_VIDEO_FILE,
|
||||
delta_timestamps=delta_timestamps,
|
||||
)
|
||||
assert_videos_roll_over(ds)
|
||||
|
||||
streaming_ds = StreamingLeRobotDataset(
|
||||
repo_id=repo_id,
|
||||
root=local_path,
|
||||
buffer_size=100,
|
||||
seed=42,
|
||||
shuffle=False,
|
||||
delta_timestamps=delta_timestamps,
|
||||
)
|
||||
assert_stream_matches_reference(streaming_ds, ds, MULTI_FILE_FRAMES)
|
||||
|
||||
Vendored
+67
-17
@@ -268,7 +268,17 @@ def episodes_factory(tasks_factory, stats_factory):
|
||||
video_keys: list[str] | None = None,
|
||||
tasks: pd.DataFrame | None = None,
|
||||
multi_task: bool = False,
|
||||
episodes_per_video_file: int | None = None,
|
||||
):
|
||||
"""Build episode metadata.
|
||||
|
||||
``episodes_per_video_file`` splits the video keys across several files, as v3.0 does
|
||||
once a video grows past ``DEFAULT_VIDEO_FILE_SIZE_IN_MB``: episode ``i`` lands in
|
||||
``file_index = i // episodes_per_video_file`` and each file's timeline restarts at 0,
|
||||
so ``from_timestamp`` is relative to the file the episode lives in — not to the
|
||||
dataset. Left ``None``, everything goes to ``file-000`` with a cumulative
|
||||
``from_timestamp`` (the single-file layout, where the two happen to coincide).
|
||||
"""
|
||||
if total_episodes <= 0 or total_frames <= 0:
|
||||
raise ValueError("num_episodes and total_length must be positive integers.")
|
||||
if total_frames < total_episodes:
|
||||
@@ -310,8 +320,16 @@ def episodes_factory(tasks_factory, stats_factory):
|
||||
d[stats_key] = []
|
||||
|
||||
num_frames = 0
|
||||
# Frames written to the current video file. Resets on every file rollover, since each
|
||||
# .mp4 carries its own timeline.
|
||||
num_frames_in_video_file = 0
|
||||
video_file_index = 0
|
||||
remaining_tasks = list(tasks.index)
|
||||
for ep_idx in range(total_episodes):
|
||||
if episodes_per_video_file is not None and ep_idx // episodes_per_video_file != video_file_index:
|
||||
video_file_index = ep_idx // episodes_per_video_file
|
||||
num_frames_in_video_file = 0
|
||||
|
||||
num_tasks_in_episode = random.randint(1, min(3, num_tasks_available)) if multi_task else 1
|
||||
tasks_to_sample = remaining_tasks if len(remaining_tasks) > 0 else list(tasks.index)
|
||||
episode_tasks = random.sample(tasks_to_sample, min(num_tasks_in_episode, len(tasks_to_sample)))
|
||||
@@ -333,21 +351,44 @@ def episodes_factory(tasks_factory, stats_factory):
|
||||
if video_keys is not None:
|
||||
for video_key in video_keys:
|
||||
d[f"videos/{video_key}/chunk_index"].append(0)
|
||||
d[f"videos/{video_key}/file_index"].append(0)
|
||||
d[f"videos/{video_key}/from_timestamp"].append(num_frames / fps)
|
||||
d[f"videos/{video_key}/to_timestamp"].append((num_frames + lengths[ep_idx]) / fps)
|
||||
d[f"videos/{video_key}/file_index"].append(video_file_index)
|
||||
d[f"videos/{video_key}/from_timestamp"].append(num_frames_in_video_file / fps)
|
||||
d[f"videos/{video_key}/to_timestamp"].append(
|
||||
(num_frames_in_video_file + lengths[ep_idx]) / fps
|
||||
)
|
||||
|
||||
# Add stats columns like "stats/action/max"
|
||||
for stats_key, stats in flatten_dict({"stats": stats_factory(features)}).items():
|
||||
d[stats_key].append(stats)
|
||||
|
||||
num_frames += lengths[ep_idx]
|
||||
num_frames_in_video_file += lengths[ep_idx]
|
||||
|
||||
return Dataset.from_dict(d)
|
||||
|
||||
return _create_episodes
|
||||
|
||||
|
||||
def video_file_frames(
|
||||
episodes: datasets.Dataset | None, video_key: str, total_frames: int
|
||||
) -> dict[int, list[int]]:
|
||||
"""Map each of ``video_key``'s files to the global frame indices it holds, in file order.
|
||||
|
||||
``episodes`` is the source of truth for how a video key is split across files. Without it
|
||||
(or without the videos columns), the whole key is one file — the pre-v3.0 layout.
|
||||
"""
|
||||
if episodes is None or f"videos/{video_key}/file_index" not in episodes.column_names:
|
||||
return {0: list(range(total_frames))}
|
||||
|
||||
frames_per_file: dict[int, list[int]] = {}
|
||||
for ep in episodes:
|
||||
file_index = ep[f"videos/{video_key}/file_index"]
|
||||
frames = range(ep["dataset_from_index"], ep["dataset_to_index"])
|
||||
frames_per_file.setdefault(file_index, []).extend(frames)
|
||||
|
||||
return frames_per_file
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def create_videos(info_factory, img_array_factory):
|
||||
def _create_video_directory(
|
||||
@@ -356,6 +397,7 @@ def create_videos(info_factory, img_array_factory):
|
||||
total_episodes: int = 3,
|
||||
total_frames: int = 150,
|
||||
total_tasks: int = 1,
|
||||
episodes: datasets.Dataset | None = None,
|
||||
):
|
||||
if info is None:
|
||||
info = info_factory(
|
||||
@@ -364,21 +406,27 @@ def create_videos(info_factory, img_array_factory):
|
||||
|
||||
video_feats = {key: feats for key, feats in info.features.items() if feats["dtype"] == "video"}
|
||||
for key, ft in video_feats.items():
|
||||
# create and save images with identifiable content
|
||||
tmp_dir = root / "tmp_images"
|
||||
tmp_dir.mkdir(parents=True, exist_ok=True)
|
||||
for frame_index in range(info.total_frames):
|
||||
content = f"{key}-{frame_index}"
|
||||
img = img_array_factory(height=ft["shape"][0], width=ft["shape"][1], content=content)
|
||||
pil_img = PIL.Image.fromarray(img)
|
||||
path = tmp_dir / f"frame-{frame_index:06d}.png"
|
||||
pil_img.save(path)
|
||||
for file_index, frame_indices in video_file_frames(episodes, key, info.total_frames).items():
|
||||
# create and save images with identifiable content
|
||||
tmp_dir = root / "tmp_images"
|
||||
tmp_dir.mkdir(parents=True, exist_ok=True)
|
||||
for position, frame_index in enumerate(frame_indices):
|
||||
# Content stays keyed on the *global* frame index so a frame remains
|
||||
# identifiable across files, but its position in the file is what the
|
||||
# file's timeline addresses.
|
||||
content = f"{key}-{frame_index}"
|
||||
img = img_array_factory(height=ft["shape"][0], width=ft["shape"][1], content=content)
|
||||
pil_img = PIL.Image.fromarray(img)
|
||||
path = tmp_dir / f"frame-{position:06d}.png"
|
||||
pil_img.save(path)
|
||||
|
||||
video_path = root / DEFAULT_VIDEO_PATH.format(video_key=key, chunk_index=0, file_index=0)
|
||||
video_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
# Use the global fps from info, not video-specific fps which might not exist
|
||||
encode_video_frames(tmp_dir, video_path, fps=info.fps)
|
||||
shutil.rmtree(tmp_dir)
|
||||
video_path = root / DEFAULT_VIDEO_PATH.format(
|
||||
video_key=key, chunk_index=0, file_index=file_index
|
||||
)
|
||||
video_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
# Use the global fps from info, not video-specific fps which might not exist
|
||||
encode_video_frames(tmp_dir, video_path, fps=info.fps)
|
||||
shutil.rmtree(tmp_dir)
|
||||
|
||||
return _create_video_directory
|
||||
|
||||
@@ -520,6 +568,7 @@ def lerobot_dataset_factory(
|
||||
data_files_size_in_mb: float = DEFAULT_DATA_FILE_SIZE_IN_MB,
|
||||
chunks_size: int = DEFAULT_CHUNK_SIZE,
|
||||
camera_features: dict | None = None,
|
||||
episodes_per_video_file: int | None = None,
|
||||
**kwargs,
|
||||
) -> LeRobotDataset:
|
||||
# Instantiate objects
|
||||
@@ -557,6 +606,7 @@ def lerobot_dataset_factory(
|
||||
video_keys=video_keys,
|
||||
tasks=tasks,
|
||||
multi_task=multi_task,
|
||||
episodes_per_video_file=episodes_per_video_file,
|
||||
)
|
||||
if hf_dataset is None:
|
||||
hf_dataset = hf_dataset_factory(
|
||||
|
||||
Vendored
+8
-2
@@ -29,6 +29,7 @@ from lerobot.datasets.utils import (
|
||||
STATS_PATH,
|
||||
)
|
||||
from tests.fixtures.constants import LEROBOT_TEST_DIR
|
||||
from tests.fixtures.dataset_factories import video_file_frames
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
@@ -99,7 +100,12 @@ def mock_snapshot_download_factory(
|
||||
|
||||
video_keys = [key for key, feats in info.features.items() if feats["dtype"] == "video"]
|
||||
for key in video_keys:
|
||||
all_files.append(DEFAULT_VIDEO_PATH.format(video_key=key, chunk_index=0, file_index=0))
|
||||
# A video key spans one file per `videos/<key>/file_index` in the episodes
|
||||
# metadata — several of them once v3.0 rolls a video over.
|
||||
for file_index in video_file_frames(episodes, key, info.total_frames):
|
||||
all_files.append(
|
||||
DEFAULT_VIDEO_PATH.format(video_key=key, chunk_index=0, file_index=file_index)
|
||||
)
|
||||
|
||||
allowed_files = filter_repo_objects(
|
||||
all_files, allow_patterns=allow_patterns, ignore_patterns=ignore_patterns
|
||||
@@ -138,7 +144,7 @@ def mock_snapshot_download_factory(
|
||||
if request_data:
|
||||
create_hf_dataset(local_dir, hf_dataset, data_files_size_in_mb, chunks_size)
|
||||
if request_videos:
|
||||
create_videos(root=local_dir, info=info)
|
||||
create_videos(root=local_dir, info=info, episodes=episodes)
|
||||
|
||||
return str(local_dir)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user