diff --git a/fastlabel/lerobot/__init__.py b/fastlabel/lerobot/__init__.py index 287c011..3709073 100644 --- a/fastlabel/lerobot/__init__.py +++ b/fastlabel/lerobot/__init__.py @@ -58,7 +58,10 @@ def get_episode_raw_frames( """ ep_info = v3.resolve_episode(episode_index, episode_map) return v3.get_episode_raw_frames( - lerobot_data_path, episode_index, ep_info["chunk"], ep_info["file_stem"] + lerobot_data_path, + episode_index, + ep_info["data_chunk"], + ep_info["data_file_stem"], ) @@ -100,6 +103,9 @@ def create_episode_zip( ] cameras = converter.select_cameras(v3.get_camera_dirs(lerobot_data_path)) ep_info = v3.resolve_episode(episode_index, episode_map) + # fps converts each camera's from_timestamp (meta/episodes) into a frame + # offset within its consolidated video file. + fps = load_info(lerobot_data_path).get("fps") return v3._assemble_episode_zip( - ep_info, episode_name, cameras, telemetry_frames, output_dir + ep_info, episode_name, cameras, telemetry_frames, output_dir, fps ) diff --git a/fastlabel/lerobot/v3.py b/fastlabel/lerobot/v3.py index fdb0807..95d0a6f 100644 --- a/fastlabel/lerobot/v3.py +++ b/fastlabel/lerobot/v3.py @@ -1,4 +1,5 @@ import json +import logging import shutil import tempfile from pathlib import Path @@ -9,14 +10,31 @@ from fastlabel.exceptions import FastLabelInvalidException from fastlabel.lerobot.common import Camera +logger = logging.getLogger(__name__) -class EpisodeInfo(TypedDict): - """Per-episode location within the v3 layout.""" + +class VideoInfo(TypedDict): + """Per-episode location of one camera's segment within a consolidated + video file.""" chunk: str file_stem: str - frame_offset: int + from_timestamp: float + to_timestamp: float + + +class EpisodeInfo(TypedDict): + """Per-episode location within the v3 layout (from meta/episodes). + + Data and video files are consolidated independently in v3, so each camera + carries its own chunk/file location under ``videos`` (keyed by the video + feature key, e.g. ``observation.images.top``). + """ + + data_chunk: str + data_file_stem: str length: int + videos: dict[str, VideoInfo] # episode_index -> EpisodeInfo @@ -25,39 +43,61 @@ class EpisodeInfo(TypedDict): Frame = dict[str, Any] +def _chunk_name(chunk_index: int) -> str: + return f"chunk-{chunk_index:03d}" + + +def _file_stem(file_index: int) -> str: + return f"file-{file_index:03d}" + + def _build_episode_map(lerobot_data_path: Path) -> EpisodeMap: """Build a mapping of episode_index -> EpisodeInfo. - Reads all data parquet files across all chunks and computes per-episode - frame offsets within each file (needed for video segment extraction). + Reads meta/episodes/chunk-XXX/file-YYY.parquet, which holds each episode's + location: ``data/chunk_index``, ``data/file_index``, ``length`` and, per + camera, ``videos/{key}/chunk_index``, ``videos/{key}/file_index``, + ``videos/{key}/from_timestamp``, ``videos/{key}/to_timestamp``. - v3 layout: data/chunk-XXX/file-YYY.parquet + Video locations cannot be derived from the data file layout: data and + video files are consolidated independently (different size limits), so + their chunk/file indices generally differ. """ import pandas as pd - data_dir = lerobot_data_path / "data" - episode_map: EpisodeMap = {} + episodes_dir = lerobot_data_path / "meta" / "episodes" + parquet_files = sorted(episodes_dir.glob("chunk-*/file-*.parquet")) + if not parquet_files: + raise FastLabelInvalidException( + f"Episode metadata not found: {episodes_dir}/chunk-*/file-*.parquet", + 422, + ) - for chunk_dir in sorted(data_dir.iterdir()): - if not chunk_dir.is_dir() or not chunk_dir.name.startswith("chunk-"): - continue - chunk_name = chunk_dir.name - - for parquet_file in sorted(chunk_dir.glob("file-*.parquet")): - file_stem = parquet_file.stem - df = pd.read_parquet(parquet_file) - - frame_offset = 0 - for ep_idx in sorted(df["episode_index"].unique()): - ep_df = df[df["episode_index"] == ep_idx] - length = len(ep_df) - episode_map[int(ep_idx)] = { - "chunk": chunk_name, - "file_stem": file_stem, - "frame_offset": frame_offset, - "length": length, + episode_map: EpisodeMap = {} + for parquet_file in parquet_files: + df = pd.read_parquet(parquet_file) + video_keys = [ + column.split("/")[1] + for column in df.columns + if column.startswith("videos/") and column.endswith("/chunk_index") + ] + for _, row in df.iterrows(): + videos: dict[str, VideoInfo] = {} + for key in video_keys: + if pd.isna(row[f"videos/{key}/chunk_index"]): + continue + videos[key] = { + "chunk": _chunk_name(int(row[f"videos/{key}/chunk_index"])), + "file_stem": _file_stem(int(row[f"videos/{key}/file_index"])), + "from_timestamp": float(row[f"videos/{key}/from_timestamp"]), + "to_timestamp": float(row[f"videos/{key}/to_timestamp"]), } - frame_offset += length + episode_map[int(row["episode_index"])] = { + "data_chunk": _chunk_name(int(row["data/chunk_index"])), + "data_file_stem": _file_stem(int(row["data/file_index"])), + "length": int(row["length"]), + "videos": videos, + } return episode_map @@ -160,29 +200,51 @@ def _assemble_episode_zip( cameras: list[Camera], telemetry_frames: list[dict[str, Any]], output_dir: Path, + fps: float | None, ) -> str: """Stage selected video segments + telemetry JSON, then archive as a ZIP. ``telemetry_frames`` is the list of frame dicts written to the episode JSON. + ``fps`` (from meta/info.json) converts each camera's ``from_timestamp`` + into a frame offset within its consolidated video file; it may be None only + when no selected camera has a video segment (otherwise raises before any + staging). A camera whose video file is missing is skipped with a warning + (the ZIP still ships the telemetry JSON). The staging directory is removed automatically; the ZIP is written under ``output_dir`` (owned by the caller) and its path returned. """ - chunk = ep_info["chunk"] - file_stem = ep_info["file_stem"] - frame_offset = ep_info["frame_offset"] + if fps is None and any(camera.key in ep_info["videos"] for camera in cameras): + raise FastLabelInvalidException( + "'fps' not found in meta/info.json " + "(required to locate episode video segments).", + 422, + ) + length = ep_info["length"] with tempfile.TemporaryDirectory() as staging: content_dir = Path(staging) # Extract video segments - # v3: videos/{key}/chunk-XXX/file-YYY.mp4 + # v3: videos/{key}/chunk-XXX/file-YYY.mp4, located per camera via + # meta/episodes (video files are consolidated independently of data). for camera in cameras: - video_path = camera.path / chunk / f"{file_stem}.mp4" + video_info = ep_info["videos"].get(camera.key) + if video_info is None: + continue + video_path = ( + camera.path / video_info["chunk"] / f"{video_info['file_stem']}.mp4" + ) if not video_path.exists(): + logger.warning( + "Video file not found, skipping camera %s: %s", + camera.key, + video_path, + ) continue output_path = content_dir / f"{camera.content_name}.mp4" - _extract_video_segment(video_path, frame_offset, length, output_path) + start_frame = round(video_info["from_timestamp"] * fps) + _extract_video_segment(video_path, start_frame, length, output_path) json_path = content_dir / f"{episode_name}.json" json_path.write_text(json.dumps(telemetry_frames, ensure_ascii=False)) diff --git a/tests/test_lerobot_v3_parquet.py b/tests/test_lerobot_v3_parquet.py index e52c7ed..56579e3 100644 --- a/tests/test_lerobot_v3_parquet.py +++ b/tests/test_lerobot_v3_parquet.py @@ -59,41 +59,108 @@ def v3_dataset(tmp_path): (meta_dir / "info.json").write_text( json.dumps( { + "fps": 10, "features": { "observation.state": {"names": ["s0", "s1"]}, "action": {"names": ["a0", "a1"]}, - } + }, } ) ) + episodes_dir = meta_dir / "episodes" / "chunk-000" + episodes_dir.mkdir(parents=True) + _write_parquet( + episodes_dir / "file-000.parquet", + [ + { + "episode_index": 0, + "data/chunk_index": 0, + "data/file_index": 0, + "length": 3, + }, + { + "episode_index": 1, + "data/chunk_index": 0, + "data/file_index": 0, + "length": 3, + }, + { + "episode_index": 2, + "data/chunk_index": 1, + "data/file_index": 0, + "length": 2, + }, + ], + ) + return tmp_path class TestBuildEpisodeMap: - def test_returns_offsets_per_episode(self, v3_dataset): + def test_returns_locations_per_episode(self, v3_dataset): result = v3._build_episode_map(v3_dataset) assert set(result.keys()) == {0, 1, 2} assert result[0] == { - "chunk": "chunk-000", - "file_stem": "file-000", - "frame_offset": 0, + "data_chunk": "chunk-000", + "data_file_stem": "file-000", "length": 3, + "videos": {}, } assert result[1] == { - "chunk": "chunk-000", - "file_stem": "file-000", - "frame_offset": 3, + "data_chunk": "chunk-000", + "data_file_stem": "file-000", "length": 3, + "videos": {}, } assert result[2] == { - "chunk": "chunk-001", - "file_stem": "file-000", - "frame_offset": 0, + "data_chunk": "chunk-001", + "data_file_stem": "file-000", "length": 2, + "videos": {}, } + def test_reads_per_camera_video_locations(self, v3_dataset): + # Video files are consolidated independently of data files, so the + # video chunk/file indices may differ from the data ones. + _write_parquet( + v3_dataset / "meta" / "episodes" / "chunk-000" / "file-000.parquet", + [ + { + "episode_index": 0, + "data/chunk_index": 0, + "data/file_index": 0, + "length": 3, + "videos/observation.images.top/chunk_index": 1, + "videos/observation.images.top/file_index": 2, + "videos/observation.images.top/from_timestamp": 1.5, + "videos/observation.images.top/to_timestamp": 1.8, + }, + ], + ) + + result = v3._build_episode_map(v3_dataset) + + assert result[0]["videos"] == { + "observation.images.top": { + "chunk": "chunk-001", + "file_stem": "file-002", + "from_timestamp": 1.5, + "to_timestamp": 1.8, + } + } + + def test_missing_episode_metadata_raises(self, v3_dataset): + import shutil + + from fastlabel.exceptions import FastLabelInvalidException + + shutil.rmtree(v3_dataset / "meta" / "episodes") + + with pytest.raises(FastLabelInvalidException): + v3._build_episode_map(v3_dataset) + def test_get_episode_indices_sorted(self, v3_dataset): assert v3.get_episode_indices(v3_dataset) == [0, 1, 2] @@ -131,7 +198,10 @@ def test_writes_zip_into_output_dir_without_leaking_staging( out.mkdir() episode_map = v3._build_episode_map(v3_dataset) raw_frames = v3.get_episode_raw_frames( - v3_dataset, 1, episode_map[1]["chunk"], episode_map[1]["file_stem"] + v3_dataset, + 1, + episode_map[1]["data_chunk"], + episode_map[1]["data_file_stem"], ) zip_path = create_episode_zip( v3_dataset, diff --git a/tests/test_lerobot_v3_video.py b/tests/test_lerobot_v3_video.py index b3be063..912544d 100644 --- a/tests/test_lerobot_v3_video.py +++ b/tests/test_lerobot_v3_video.py @@ -1,8 +1,11 @@ +import zipfile + import cv2 import pytest from fastlabel.exceptions import FastLabelInvalidException from fastlabel.lerobot import v3 +from fastlabel.lerobot.common import Camera class TestExtractVideoSegment: @@ -59,3 +62,135 @@ def test_unopenable_file_raises(self, tmp_path): num_frames=1, output_path=tmp_path / "out.mp4", ) + + +class TestAssembleEpisodeZip: + """Video location comes from meta/episodes (per camera), not from the + data file layout: the video chunk/file indices may differ from the data + ones, and the frame offset is from_timestamp * fps.""" + + CAMERA_KEY = "observation.images.top" + + def _episode_info(self, videos): + return { + "data_chunk": "chunk-000", + "data_file_stem": "file-000", + "length": 4, + "videos": videos, + } + + def _camera(self, tmp_path): + return Camera( + path=tmp_path / "videos" / self.CAMERA_KEY, + key=self.CAMERA_KEY, + content_name="images_top", + ) + + def test_resolves_video_location_and_timestamp_offset( + self, synthetic_video, tmp_path + ): + # The episode's video lives in file-001 even though its data lives in + # file-000, and starts 0.5s (= frame 5 at fps 10) into that file. + (tmp_path / "videos" / self.CAMERA_KEY / "chunk-000").mkdir(parents=True) + synthetic_video( + name=f"videos/{self.CAMERA_KEY}/chunk-000/file-001.mp4", + num_frames=10, + fps=10, + ) + ep_info = self._episode_info( + { + self.CAMERA_KEY: { + "chunk": "chunk-000", + "file_stem": "file-001", + "from_timestamp": 0.5, + "to_timestamp": 0.9, + } + } + ) + out = tmp_path / "out" + out.mkdir() + + zip_path = v3._assemble_episode_zip( + ep_info, + "episode_000001", + [self._camera(tmp_path)], + telemetry_frames=[], + output_dir=out, + fps=10.0, + ) + + with zipfile.ZipFile(zip_path) as zf: + assert sorted(zf.namelist()) == [ + "episode_000001.json", + "images_top.mp4", + ] + extract_dir = tmp_path / "extracted" + zf.extract("images_top.mp4", extract_dir) + + cap = cv2.VideoCapture(str(extract_dir / "images_top.mp4")) + try: + count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + finally: + cap.release() + assert count == 4 + + def test_missing_video_file_warns_and_skips(self, tmp_path, caplog): + (tmp_path / "videos" / self.CAMERA_KEY / "chunk-000").mkdir(parents=True) + ep_info = self._episode_info( + { + self.CAMERA_KEY: { + "chunk": "chunk-000", + "file_stem": "file-009", + "from_timestamp": 0.0, + "to_timestamp": 0.4, + } + } + ) + out = tmp_path / "out" + out.mkdir() + + with caplog.at_level("WARNING", logger="fastlabel.lerobot.v3"): + zip_path = v3._assemble_episode_zip( + ep_info, + "episode_000001", + [self._camera(tmp_path)], + telemetry_frames=[], + output_dir=out, + fps=10.0, + ) + + # The missing camera is skipped with a warning; the ZIP still ships + # the telemetry JSON. + assert any("file-009.mp4" in message for message in caplog.messages) + with zipfile.ZipFile(zip_path) as zf: + assert zf.namelist() == ["episode_000001.json"] + + def test_missing_fps_raises_when_video_needed(self, synthetic_video, tmp_path): + (tmp_path / "videos" / self.CAMERA_KEY / "chunk-000").mkdir(parents=True) + synthetic_video( + name=f"videos/{self.CAMERA_KEY}/chunk-000/file-000.mp4", + num_frames=10, + fps=10, + ) + ep_info = self._episode_info( + { + self.CAMERA_KEY: { + "chunk": "chunk-000", + "file_stem": "file-000", + "from_timestamp": 0.0, + "to_timestamp": 0.4, + } + } + ) + out = tmp_path / "out" + out.mkdir() + + with pytest.raises(FastLabelInvalidException): + v3._assemble_episode_zip( + ep_info, + "episode_000001", + [self._camera(tmp_path)], + telemetry_frames=[], + output_dir=out, + fps=None, + )