From df93cd130c77dbeedee08dbbb0f9fc051229cab6 Mon Sep 17 00:00:00 2001 From: rikunosuke Date: Wed, 5 Aug 2026 16:56:47 +0900 Subject: [PATCH 1/2] =?UTF-8?q?fix:=20v3=20=E3=81=AE=E5=8B=95=E7=94=BB?= =?UTF-8?q?=E4=BD=8D=E7=BD=AE=E3=82=92=20data=20=E3=83=AC=E3=82=A4?= =?UTF-8?q?=E3=82=A2=E3=82=A6=E3=83=88=E3=81=A7=E3=81=AF=E3=81=AA=E3=81=8F?= =?UTF-8?q?=20meta/episodes=20=E3=81=8B=E3=82=89=E8=A7=A3=E6=B1=BA?= =?UTF-8?q?=E3=81=99=E3=82=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit _build_episode_map は各エピソードの動画 chunk/file を data parquet の レイアウトから導出していたが、v3 では data と動画が独立に集約されるため、 存在しない mp4 を参照(黙ってスキップ)するか、同名ファイルが偶然存在する 場合は data ファイルの行オフセットを動画フレームオフセットとして使い 誤った区間を切り出していた。 - episode map を meta/episodes/chunk-*/file-*.parquet から構築し、 カメラごとの動画位置(chunk/file index, from/to_timestamp)を保持 - 動画フレームオフセットを from_timestamp × fps(meta/info.json)で算出 - 動画ファイル欠落時は黙って落とさず warning を出してスキップ。 切り出しが必要なのに fps が無い場合は例外 Co-Authored-By: Claude Fable 5 --- fastlabel/lerobot/__init__.py | 10 ++- fastlabel/lerobot/v3.py | 126 +++++++++++++++++++++-------- tests/test_lerobot_v3_parquet.py | 94 ++++++++++++++++++--- tests/test_lerobot_v3_video.py | 135 +++++++++++++++++++++++++++++++ 4 files changed, 318 insertions(+), 47 deletions(-) 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..ba3c8f8 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,49 @@ 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 camera segment is extracted. 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"] 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 + if fps is None: + raise FastLabelInvalidException( + "'fps' not found in meta/info.json " + "(required to locate episode video segments).", + 422, + ) 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, + ) From f2851abf9f05a5821c5d0a16c053ba81a0a7af69 Mon Sep 17 00:00:00 2001 From: rikunosuke Date: Wed, 5 Aug 2026 17:08:35 +0900 Subject: [PATCH 2/2] =?UTF-8?q?refactor:=20fps=20=E6=AC=A0=E8=90=BD?= =?UTF-8?q?=E3=83=81=E3=82=A7=E3=83=83=E3=82=AF=E3=82=92=20=5Fassemble=5Fe?= =?UTF-8?q?pisode=5Fzip=20=E3=81=AE=E5=86=92=E9=A0=AD=E3=81=AB=E7=A7=BB?= =?UTF-8?q?=E5=8B=95?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit カメラループ内でのチェックだと、ステージング作成や先行カメラの 切り出し後に失敗しうるため、選択カメラに動画セグメントがあるのに fps が無い場合は処理開始前に例外を出す。 Co-Authored-By: Claude Fable 5 --- fastlabel/lerobot/v3.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/fastlabel/lerobot/v3.py b/fastlabel/lerobot/v3.py index ba3c8f8..95d0a6f 100644 --- a/fastlabel/lerobot/v3.py +++ b/fastlabel/lerobot/v3.py @@ -207,11 +207,19 @@ def _assemble_episode_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 camera segment is extracted. A camera whose video file is missing - is skipped with a warning (the ZIP still ships the telemetry JSON). + 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. """ + 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: @@ -234,12 +242,6 @@ def _assemble_episode_zip( video_path, ) continue - if fps is None: - raise FastLabelInvalidException( - "'fps' not found in meta/info.json " - "(required to locate episode video segments).", - 422, - ) output_path = content_dir / f"{camera.content_name}.mp4" start_frame = round(video_info["from_timestamp"] * fps) _extract_video_segment(video_path, start_frame, length, output_path)