From e0f78bad9b3357184236f98524172bdb3330f3b6 Mon Sep 17 00:00:00 2001 From: Griswald Brooks Date: Fri, 21 Aug 2026 19:58:47 -0400 Subject: [PATCH] feat(lerobot_v3): add image-mode output to the v3 writer Add `LeRobotV3WriterConfig.video` (default True). With `video=False` the writer PNG-encodes each camera frame into the data parquet as a `struct` cell and declares the feature as `dtype: "image"` in info.json, matching the LeRobot v3 image-dataset layout and upstream LeRobot's own `video=False` convention. The video encoder is bypassed entirely in that mode and no videos/ directory is created; `video_path` and `video_files_size_in_mb` are null while the keys are retained. Video mode remains the default and its output is unchanged. Also in this change: - carry `writer_config` through `ConversionConfig.to_dict()` and apply it in the parallel conversion worker, so `--workers > 1` honours image mode; - read inline image cells back in the v3 reader, tolerating null and path-only cells; - keep late-appearing parquet columns by filling the union of row keys before building a chunk table; - skip null rows when computing episode stats and when inferring features; - bound finalize's feature sampling to one batch instead of reading whole PNG payloads; - share one `_configure_writer` helper between the sequential and parallel converter paths. Tests cover image-mode dataset structure, a read-back through the LeRobot v3 reader, and a video-mode regression class. --- CHANGELOG.md | 18 + docs/configuration.md | 12 + forge/config/models.py | 3 + forge/convert/converter.py | 98 ++--- forge/formats/lerobot_v3/reader.py | 53 +++ forge/formats/lerobot_v3/writer.py | 173 ++++++-- pyproject.toml | 1 + tests/test_converter.py | 60 +++ tests/test_lerobot_v3_roundtrip.py | 80 ++++ tests/test_lerobot_v3_upstream_validity.py | 46 ++ tests/test_lerobot_v3_writer.py | 471 +++++++++++++++++++++ 11 files changed, 930 insertions(+), 85 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index cedf16b..80ac0f2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,6 +9,17 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +- **Image-mode output for the LeRobot v3 writer.** + `LeRobotV3WriterConfig(video=False)` PNG-encodes each camera frame inline into + the data parquet (`dtype: "image"` features) instead of encoding MP4s — no + `videos/` directory and no per-camera `videos/…` pointer columns are written, + and `info.json`'s `video_path` / `video_files_size_in_mb` are `null` (the keys + stay present, as v3 loaders require). Set it from a conversion config with + `writer_config: {video: false}`. Needs Pillow (now part of the `[lerobot]` + extra). Video mode stays the default and is unchanged. + The LeRobot v3 **reader** decodes these inline PNG columns, so image-mode + datasets round-trip through Forge. + - **Streaming reads for cloud datasets (LeRobot-v3, Zarr).** `forge inspect` **and `forge ingest`** on `s3://…` / `gs://…` now read over the network with **range requests** instead of downloading the whole dataset. @@ -148,6 +159,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 vars / profiles / IAM roles; GCP Application Default Credentials). Forge never handles credentials itself. +### Fixed + +- **`writer_config` now reaches parallel conversion workers.** `ConversionConfig` + dropped `writer_config` when serialized, so with `--workers > 1` every + writer-specific option (including the new `video: false`) was silently ignored + by the worker processes. + ### Changed - `fsspec` is now a core dependency. diff --git a/docs/configuration.md b/docs/configuration.md index d0f124e..147f362 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -100,6 +100,18 @@ video: compress: true # Enable compression ``` +### Writer Settings + +Options passed straight to the target format's writer config (unknown keys are +ignored): + +```yaml +writer_config: + video: false # lerobot-v3: store camera frames as PNG bytes inline in + # the data parquet (dtype "image") instead of MP4 videos. + # Default true. Needs the [lerobot] extra (Pillow). +``` + ### Behavior Settings ```yaml diff --git a/forge/config/models.py b/forge/config/models.py index 81f8cff..fff872c 100644 --- a/forge/config/models.py +++ b/forge/config/models.py @@ -309,6 +309,9 @@ def to_dict(self) -> dict[str, Any]: if self.num_workers > 1: result["num_workers"] = self.num_workers + if self.writer_config: + result["writer_config"] = self.writer_config + return result def to_yaml(self, path: Path | str) -> None: diff --git a/forge/convert/converter.py b/forge/convert/converter.py index 9cdfa72..5ac6f7f 100644 --- a/forge/convert/converter.py +++ b/forge/convert/converter.py @@ -62,6 +62,51 @@ class _EpisodeResult: error: str | None = None +def _configure_writer( + writer: Any, config: ConversionConfig, dataset_info: DatasetInfo +) -> None: + """Apply a conversion config to a writer's format-specific config. + + Shared by the sequential path and the parallel worker so writer options + cannot drift between them. + + Args: + writer: Writer instance to configure. + config: Conversion config supplying the overrides. + dataset_info: Dataset info for default values. + """ + # Check if writer has a config attribute (like LeRobotV3Writer) + if not hasattr(writer, "config"): + return + + wconfig = writer.config + + # Apply defaults from dataset_info + if hasattr(wconfig, "fps") and dataset_info.inferred_fps: + wconfig.fps = dataset_info.inferred_fps + + if hasattr(wconfig, "robot_type") and dataset_info.inferred_robot_type: + wconfig.robot_type = dataset_info.inferred_robot_type + + if hasattr(wconfig, "camera_name_mapping") and config.camera_mapping: + wconfig.camera_name_mapping = config.camera_mapping + + # Apply field mappings if writer supports them + if hasattr(wconfig, "field_mapping") and config.field_mapping: + wconfig.field_mapping = config.field_mapping + + if hasattr(wconfig, "action_field") and config.action_field: + wconfig.action_field = config.action_field + + if hasattr(wconfig, "state_field") and config.state_field: + wconfig.state_field = config.state_field + + # Apply format-specific config from conversion config + for key, value in config.writer_config.items(): + if hasattr(wconfig, key): + setattr(wconfig, key, value) + + def _process_episode_worker( source_path: str, output_path: str, @@ -131,20 +176,7 @@ def _process_episode_worker( dataset_info.inferred_robot_type = config.robot_type # Configure writer - if hasattr(writer, "config"): - wconfig = writer.config - if hasattr(wconfig, "fps") and dataset_info.inferred_fps: - wconfig.fps = dataset_info.inferred_fps - if hasattr(wconfig, "robot_type") and dataset_info.inferred_robot_type: - wconfig.robot_type = dataset_info.inferred_robot_type - if hasattr(wconfig, "camera_name_mapping") and config.camera_mapping: - wconfig.camera_name_mapping = config.camera_mapping - if hasattr(wconfig, "field_mapping") and config.field_mapping: - wconfig.field_mapping = config.field_mapping - if hasattr(wconfig, "action_field") and config.action_field: - wconfig.action_field = config.action_field - if hasattr(wconfig, "state_field") and config.state_field: - wconfig.state_field = config.state_field + _configure_writer(writer, config, dataset_info) # Read and write the specific episode episodes_iter = reader.read_episodes(source) @@ -291,7 +323,7 @@ def convert( self._apply_config_overrides(dataset_info) # 5. Configure writer with format-specific options - self._configure_writer(writer, dataset_info) + _configure_writer(writer, self.config, dataset_info) # 6. Convert episodes (parallel or sequential) total_episodes = dataset_info.num_episodes or 0 @@ -368,42 +400,6 @@ def _apply_config_overrides(self, dataset_info: DatasetInfo) -> None: if self.config.robot_type is not None: dataset_info.inferred_robot_type = self.config.robot_type - def _configure_writer(self, writer: Any, dataset_info: DatasetInfo) -> None: - """Configure writer with format-specific options. - - Args: - writer: Writer instance to configure. - dataset_info: Dataset info for default values. - """ - # Check if writer has a config attribute (like LeRobotV3Writer) - if hasattr(writer, "config"): - config = writer.config - - # Apply defaults from dataset_info - if hasattr(config, "fps") and dataset_info.inferred_fps: - config.fps = dataset_info.inferred_fps - - if hasattr(config, "robot_type") and dataset_info.inferred_robot_type: - config.robot_type = dataset_info.inferred_robot_type - - if hasattr(config, "camera_name_mapping") and self.config.camera_mapping: - config.camera_name_mapping = self.config.camera_mapping - - # Apply field mappings if writer supports them - if hasattr(config, "field_mapping") and self.config.field_mapping: - config.field_mapping = self.config.field_mapping - - if hasattr(config, "action_field") and self.config.action_field: - config.action_field = self.config.action_field - - if hasattr(config, "state_field") and self.config.state_field: - config.state_field = self.config.state_field - - # Apply format-specific config from conversion config - for key, value in self.config.writer_config.items(): - if hasattr(config, key): - setattr(config, key, value) - def _convert_parallel( self, source: Path, diff --git a/forge/formats/lerobot_v3/reader.py b/forge/formats/lerobot_v3/reader.py index 2c14bba..2310349 100644 --- a/forge/formats/lerobot_v3/reader.py +++ b/forge/formats/lerobot_v3/reader.py @@ -21,10 +21,15 @@ │ ├── episode_000000.mp4 │ └── ... └── ... + +Image-mode datasets (features with ``dtype: "image"``) have no videos/ +directory: each camera column holds PNG bytes inline in the data parquet, +decoded lazily per frame. """ from __future__ import annotations +import io import json from collections.abc import Iterator from pathlib import Path @@ -1144,12 +1149,17 @@ def _build_episode( # Get camera info from info.json for dimensions info_path = dataset_path / "meta" / "info.json" camera_features: dict[str, dict] = {} + # Image-mode cameras (dtype == "image"): inline PNG bytes live directly + # in the data parquet under their full feature key. + image_camera_features: dict[str, dict] = {} if info_path.exists(): with open(info_path) as f: info = json.load(f) for key, spec in info.get("features", {}).items(): if spec.get("dtype") == "video": camera_features[key] = spec + elif spec.get("dtype") == "image": + image_camera_features[key] = spec # Get video path template from info.json video_path_template = None @@ -1343,6 +1353,49 @@ def make_loader(vp: Path = video_path, fi: int = local_frame_idx, d: tuple = dim channels=dims[2], ) + # Inline image-mode cameras: PNG bytes live in this row's own + # struct column, decoded lazily on load(). + for image_key, spec in image_camera_features.items(): + if image_key not in ep_df.columns: + continue + cam_name = image_key.split(".")[-1] + cell = row[image_key] + png_bytes: bytes | None + png_path: str | None + if isinstance(cell, bytes): + png_bytes, png_path = cell, None + elif isinstance(cell, dict): + png_bytes, png_path = cell.get("bytes"), cell.get("path") + else: + continue + if png_bytes is None: + if not png_path: + continue + resolved = Path(png_path) + if not resolved.is_absolute(): + resolved = dataset_path / resolved + if not resolved.exists(): + continue + png_path = str(resolved) + shape = spec.get("shape") or [0, 0, 3] + + def make_image_loader( + data: bytes | None = png_bytes, path: str | None = png_path + ) -> NDArray[Any]: + import numpy as np + from PIL import Image + + source: Any = io.BytesIO(data) if data is not None else path + with Image.open(source) as img: + return np.array(img) + + images[cam_name] = LazyImage( + loader=make_image_loader, + height=shape[0], + width=shape[1], + channels=shape[2] if len(shape) > 2 else 3, + ) + # Extract state and action state = None action = None diff --git a/forge/formats/lerobot_v3/writer.py b/forge/formats/lerobot_v3/writer.py index 5f0e053..638fe63 100644 --- a/forge/formats/lerobot_v3/writer.py +++ b/forge/formats/lerobot_v3/writer.py @@ -15,11 +15,16 @@ └── videos/ └── observation.images.{camera}/ └── chunk-000/ - └── file-000.mp4 # Video data + └── file-000.mp4 # Video data (video mode, default) + +With ``LeRobotV3WriterConfig(video=False)``, camera frames are PNG-encoded +inline into the data parquet (as ``dtype: "image"`` features) instead of +being written as MP4s under videos/ — no videos/ directory is created. """ from __future__ import annotations +import io import json from collections.abc import Callable, Iterator from dataclasses import dataclass, field @@ -44,6 +49,18 @@ def _check_pyarrow() -> None: ) +def _check_pillow() -> None: + """Check if Pillow is available (required for image-mode writing).""" + try: + import PIL # noqa: F401 + except ImportError: + raise MissingDependencyError( + dependency="pillow", + feature="LeRobot v3 image-mode writing (video=False)", + install_hint="pip install forge-robotics[lerobot]", + ) + + @dataclass class LeRobotV3WriterConfig: """Configuration for LeRobot v3 writer. @@ -51,6 +68,10 @@ class LeRobotV3WriterConfig: Attributes: fps: Frames per second (required). robot_type: Robot type identifier (e.g., "franka", "so100"). + video: If True (default), camera frames are encoded to MP4 (v3 + "video" dtype features). If False, frames are written as PNG + bytes inline in the data parquet (v3 "image" dtype features), + matching upstream LeRobot's own ``video=False`` dataset mode. video_codec: Video codec for encoding (default: "libx264"). video_crf: Constant rate factor for video quality (default: 23). video_preset: FFmpeg preset for encoding speed (default: "medium"). @@ -65,6 +86,7 @@ class LeRobotV3WriterConfig: fps: float = 30.0 robot_type: str = "unknown" + video: bool = True video_codec: str = "libx264" video_crf: int = 23 video_preset: str = "medium" @@ -83,7 +105,8 @@ class LeRobotV3Writer: Converts Episode/Frame data to LeRobot v3 format with: - Chunked parquet files for state/action data - - MP4 videos for each camera (chunked) + - MP4 videos for each camera (chunked), or inline PNG bytes in the data + parquet when ``LeRobotV3WriterConfig(video=False)`` - Parquet metadata files Example: @@ -151,6 +174,22 @@ def _map_camera_name(self, source_name: str) -> str: return f"observation.images.{clean_name}" + def _encode_image_png(self, image_array: Any) -> bytes: + """Encode a frame array to PNG bytes for inline image-mode storage.""" + import numpy as np + from PIL import Image + + if image_array.dtype != np.uint8: + if image_array.dtype in (np.float32, np.float64) and image_array.max() <= 1.0: + image_array = (image_array * 255).astype(np.uint8) + else: + image_array = np.clip(image_array, 0, 255).astype(np.uint8) + + img = Image.fromarray(image_array) + buffer = io.BytesIO() + img.save(buffer, format="PNG") + return buffer.getvalue() + # Features included in stats columns. Image/video features are skipped # for now — upstream loaders only require *some* feature to have stats, # not all, and pixel-stat computation is expensive. @@ -171,7 +210,10 @@ def _compute_episode_stats(self, table: Any | None) -> dict[str, dict[str, list] for feat in self._STAT_FEATURES: if feat not in table.column_names: continue - arr = np.asarray(table[feat].to_pylist(), dtype=np.float64) + values = [v for v in table[feat].to_pylist() if v is not None] + if not values: + continue + arr = np.asarray(values, dtype=np.float64) if arr.ndim == 1: arr = arr.reshape(-1, 1) out[feat] = { @@ -265,6 +307,12 @@ def _flush_chunk(self, output_path: Path) -> None: data_dir.mkdir(parents=True, exist_ok=True) data_path = data_dir / f"file-{file_index:03d}.parquet" + all_keys = list(dict.fromkeys(k for r in self._current_chunk_frames for k in r)) + for r in self._current_chunk_frames: + if len(r) != len(all_keys): + for k in all_keys: + r.setdefault(k, None) + try: table = pa.Table.from_pylist(self._current_chunk_frames) pq.write_table(table, data_path) @@ -317,6 +365,8 @@ def write_episode( ConversionError: If writing fails. """ _check_pyarrow() + if not self.config.video: + _check_pillow() if episode_index is None: episode_index = len(self._episode_metadata) @@ -401,18 +451,12 @@ def write_episode( "fps": float(fps), } - # Collect camera frames for video encoding - # Note: Video features are NOT stored in parquet - only in the video files - # The mapping is done via the global 'index' column + # Collect camera frames. In video mode they are NOT stored in parquet — + # only in the video files, mapped via the global 'index' column. In image + # mode they are PNG-encoded straight into this row instead. for cam_name, lazy_img in frame.images.items(): mapped_name = self._map_camera_name(cam_name) - if mapped_name not in self._current_chunk_videos: - self._current_chunk_videos[mapped_name] = [] - self._current_chunk_videos[mapped_name].append(lazy_img) - if mapped_name not in self._video_keys: - self._video_keys.append(mapped_name) - # Track camera info if cam_name not in self._cameras: self._cameras[cam_name] = CameraInfo( @@ -420,22 +464,49 @@ def write_episode( height=lazy_img.height, width=lazy_img.width, channels=lazy_img.channels, + storage="mp4" if self.config.video else "image", ) - # Add feature definition (video features go in info.json, not parquet) - if mapped_name not in self._features: - self._features[mapped_name] = { - "dtype": "video", - "shape": [lazy_img.height, lazy_img.width, lazy_img.channels], - "names": ["height", "width", "channel"], - "video_info": { - "video.fps": float(fps), - "video.codec": "h264", - "video.pix_fmt": "yuv420p", - "video.is_depth_map": False, - "has_audio": False, - }, - } + if self.config.video: + if mapped_name not in self._current_chunk_videos: + self._current_chunk_videos[mapped_name] = [] + self._current_chunk_videos[mapped_name].append(lazy_img) + if mapped_name not in self._video_keys: + self._video_keys.append(mapped_name) + + # Add feature definition (video features go in info.json, not parquet) + if mapped_name not in self._features: + self._features[mapped_name] = { + "dtype": "video", + "shape": [lazy_img.height, lazy_img.width, lazy_img.channels], + "names": ["height", "width", "channel"], + "video_info": { + "video.fps": float(fps), + "video.codec": "h264", + "video.pix_fmt": "yuv420p", + "video.is_depth_map": False, + "has_audio": False, + }, + } + else: + # Image mode: PNG-encode inline into the data row, matching + # the HF `datasets.Image()` struct storage that + # upstream LeRobot uses for non-video (image) datasets. + image_array = lazy_img.load() + png_bytes = self._encode_image_png(image_array) + lazy_img.clear_cache() + row[mapped_name] = {"bytes": png_bytes, "path": None} + if mapped_name not in self._features: + self._features[mapped_name] = { + "dtype": "image", + "shape": ( + list(image_array.shape) + if image_array.ndim == 3 + else [*image_array.shape, 1] + ), + "names": ["height", "width", "channel"], + "fps": float(fps), + } self._current_chunk_frames.append(row) @@ -581,13 +652,31 @@ def finalize(self, output_path: Path, dataset_info: DatasetInfo) -> None: # In parallel mode, infer features from the first data file if not self._features or len(self._features) <= 5: # Only standard features first_data_file = output_path / "data" / "chunk-000" / "file-000.parquet" + sample_batch = None if first_data_file.exists(): - sample_table = pq.read_table(first_data_file) - for col_name in sample_table.column_names: + sample_batch = next( + pq.ParquetFile(first_data_file).iter_batches(batch_size=256), None + ) + if sample_batch is not None: + for col_name in sample_batch.schema.names: if col_name not in self._features: - col = sample_table.column(col_name) + col = sample_batch.column(col_name) # Infer dtype and shape from data - first_val = col[0].as_py() if len(col) > 0 else None + first_val = next((s.as_py() for s in col if s.is_valid), None) + if isinstance(first_val, dict) and "bytes" in first_val: + # Inline image-mode column (struct). + from PIL import Image + + with Image.open(io.BytesIO(first_val["bytes"])) as img: + width, height = img.size + channels = len(img.getbands()) + self._features[col_name] = { + "dtype": "image", + "shape": [height, width, channels], + "names": ["height", "width", "channel"], + "fps": float(fps), + } + continue if isinstance(first_val, list): shape = [len(first_val)] dtype = "float32" @@ -652,9 +741,15 @@ def finalize(self, output_path: Path, dataset_info: DatasetInfo) -> None: "train": f"0:{total_episodes}", }, "data_path": "data/chunk-{chunk_index:03d}/file-{file_index:03d}.parquet", - "video_path": "videos/{video_key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4", + "video_path": ( + "videos/{video_key}/chunk-{chunk_index:03d}/file-{file_index:03d}.mp4" + if self.config.video + else None + ), "data_files_size_in_mb": int(self.config.data_files_size_in_mb), - "video_files_size_in_mb": int(self.config.video_files_size_in_mb), + "video_files_size_in_mb": ( + int(self.config.video_files_size_in_mb) if self.config.video else None + ), "features": self._features, } @@ -689,8 +784,18 @@ def finalize(self, output_path: Path, dataset_info: DatasetInfo) -> None: data_file = ( output_path / "data" / f"chunk-{chunk_idx:03d}" / f"file-{file_idx:03d}.parquet" ) - data_table = pq.read_table(data_file) if data_file.exists() else None - length = data_table.num_rows if data_table is not None else 0 + data_table = None + length = 0 + if data_file.exists(): + parquet_file = pq.ParquetFile(data_file) + length = parquet_file.metadata.num_rows + stat_columns = [ + name + for name in (*self._STAT_FEATURES, "task_index") + if name in parquet_file.schema_arrow.names + ] + if stat_columns: + data_table = parquet_file.read(columns=stat_columns) # Resolve task for this episode. if episode_idx < len(self._episode_metadata): diff --git a/pyproject.toml b/pyproject.toml index e1c45cb..d2e7d85 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -40,6 +40,7 @@ rlds = [ lerobot = [ "pyarrow>=14.0.0", "datasets>=2.16.0", + "pillow>=10.0.0", ] hub = [ "huggingface_hub>=0.20.0", diff --git a/tests/test_converter.py b/tests/test_converter.py index a14d53b..cba01b5 100644 --- a/tests/test_converter.py +++ b/tests/test_converter.py @@ -147,6 +147,16 @@ def test_custom_config(self): assert config.robot_type == "franka" assert config.fail_on_error is True + def test_writer_config_round_trips_through_dict(self): + """writer_config must survive to_dict/from_dict (used to serialize + the config across process boundaries for parallel conversion).""" + config = ConversionConfig( + target_format="lerobot-v3", + writer_config={"video": False}, + ) + restored = ConversionConfig.from_dict(config.to_dict()) + assert restored.writer_config == {"video": False} + class TestConversionResult: """Tests for ConversionResult.""" @@ -350,3 +360,53 @@ def test_convert_function( assert result.success is True assert result.episodes_converted == 3 + + +class TestParallelWorkerWriterConfig: + """Regression: the parallel conversion path must honor writer_config, + same as the sequential path (converter.py's `_process_episode_worker` + used to only forward fps/robot_type/camera_mapping/field_mapping, so + e.g. `writer_config: {"video": False}` was silently dropped under + `--workers > 1`).""" + + @pytest.mark.skipif( + not _check_dependencies_available(), + reason="PyAV or PyArrow not installed", + ) + def test_worker_applies_writer_config_video_flag( + self, tmp_path: Path, mock_dataset: Path, register_mock_reader + ): + from forge.convert.converter import _process_episode_worker + + output_path = tmp_path / "output" + config = ConversionConfig( + target_format="lerobot-v3", + fps=30.0, + writer_config={"video": False}, + ) + + result = _process_episode_worker( + str(mock_dataset), + str(output_path), + 0, + "mock", + "lerobot-v3", + config.to_dict(), + {"inferred_fps": 30.0, "inferred_robot_type": "mock_robot", "cameras": {}}, + ) + + assert result.success is True + + # Image mode: inline PNG bytes in the parquet, no videos/ directory. + assert not (output_path / "videos").exists() + parquet_path = output_path / "data" / "chunk-000" / "file-000.parquet" + assert parquet_path.exists() + + import pyarrow.parquet as pq + + table = pq.read_table(parquet_path) + cam_col = "observation.images.camera0" + assert cam_col in table.column_names + row = table.to_pylist()[0] + assert row[cam_col]["path"] is None + assert isinstance(row[cam_col]["bytes"], bytes) diff --git a/tests/test_lerobot_v3_roundtrip.py b/tests/test_lerobot_v3_roundtrip.py index 27bec8b..553ef24 100644 --- a/tests/test_lerobot_v3_roundtrip.py +++ b/tests/test_lerobot_v3_roundtrip.py @@ -125,3 +125,83 @@ def test_video_frames_survive_round_trip(self, tmp_path: Path): f"episode {episode_index} decoded black frames (mean={mean})" ) + @pytest.mark.skipif( + not _check_dependencies_available(), + reason="PyAV or PyArrow not installed", + ) + def test_image_mode_frames_survive_round_trip(self, tmp_path: Path): + """Image mode (video=False) round-trips through the reader too. + + The reader used to build every frame's `images` dict exclusively + from `videos/*.mp4`, so an image-mode dataset (inline PNG bytes, + dtype: "image") read back with zero images per frame, even though + the reader already detected `storage="image"` at inspect time. + """ + output_dir = tmp_path / "dataset" + episodes = [ + _make_episode(0, 5, "first"), + _make_episode(1, 5, "second"), + ] + writer = LeRobotV3Writer( + LeRobotV3WriterConfig(fps=30.0, robot_type="test_robot", video=False) + ) + writer.write_dataset(iter(episodes), output_dir) + + read_back = list(LeRobotV3Reader().read_episodes(output_dir)) + + assert len(read_back) == 2 + for episode_index, episode in enumerate(read_back): + frame = list(episode.frames())[3] + assert "camera0" in frame.images + image = np.asarray(frame.images["camera0"].load()) + assert image.shape == (32, 32, 3) + # _make_episode paints frame N with value (N * 10) % 256; frame 3 -> 30. + assert float(image.mean()) == pytest.approx(30.0, abs=1.0) + + @pytest.mark.skipif( + not _check_dependencies_available(), + reason="PyAV or PyArrow not installed", + ) + def test_image_mode_tolerates_null_and_path_only_cells(self, tmp_path: Path): + """Null and path-only HF Image structs must not crash the reader. + + HuggingFace's `Image()` feature stores either embedded bytes or + `{"bytes": None, "path": ...}` for on-disk images, and a missing + image surfaces as a null struct. The reader used to assume every + cell was a dict carrying bytes. + """ + import pyarrow as pa + import pyarrow.parquet as pq + from PIL import Image + + output_dir = tmp_path / "dataset" + writer = LeRobotV3Writer( + LeRobotV3WriterConfig(fps=30.0, robot_type="test_robot", video=False) + ) + writer.write_dataset(iter([_make_episode(0, 5, "first")]), output_dir) + + external = output_dir / "images" / "frame1.png" + external.parent.mkdir(parents=True, exist_ok=True) + Image.fromarray(np.full((32, 32, 3), 77, dtype=np.uint8)).save(external) + + parquet_path = next(output_dir.glob("data/**/*.parquet")) + table = pq.read_table(parquet_path) + key = "observation.images.camera0" + col_idx = table.schema.get_field_index(key) + cells = table.column(key).to_pylist() + cells[0] = None + cells[1] = {"bytes": None, "path": "images/frame1.png"} + old_field = table.schema.field(col_idx) + # HF's own Image storage type: bytes may be null when path is set. + new_type = pa.struct([pa.field("bytes", pa.binary()), pa.field("path", pa.string())]) + new_field = pa.field(key, new_type, nullable=True, metadata=old_field.metadata) + table = table.set_column(col_idx, new_field, pa.array(cells, type=new_type)) + pq.write_table(table, parquet_path) + + frames = list(next(iter(LeRobotV3Reader().read_episodes(output_dir))).frames()) + + assert "camera0" not in frames[0].images + path_mean = float(np.asarray(frames[1].images["camera0"].load()).mean()) + inline_mean = float(np.asarray(frames[3].images["camera0"].load()).mean()) + assert path_mean == pytest.approx(77.0, abs=1.0) + assert inline_mean == pytest.approx(30.0, abs=1.0) diff --git a/tests/test_lerobot_v3_upstream_validity.py b/tests/test_lerobot_v3_upstream_validity.py index c18d8bc..1979387 100644 --- a/tests/test_lerobot_v3_upstream_validity.py +++ b/tests/test_lerobot_v3_upstream_validity.py @@ -86,6 +86,18 @@ def written_v3(tmp_path: Path) -> Path: return out +@pytest.fixture +def written_v3_image(tmp_path: Path) -> Path: + """Write a 3-episode v3 dataset in image mode (`video=False`).""" + out = tmp_path / "synth_v3_image" + writer = LeRobotV3Writer( + LeRobotV3WriterConfig(fps=30.0, robot_type="synth", video=False) + ) + episodes = (_synth_episode(f"ep_{i:03d}") for i in range(3)) + writer.write_dataset(episodes, out) + return out + + # --------------------------------------------------------------------------- # Constants — sourced from upstream lerobot==0.5.1 # --------------------------------------------------------------------------- @@ -302,3 +314,37 @@ def test_dataset_loads_and_iterates(self, written_v3: Path): sample = ds[0] assert "observation.state" in sample assert "action" in sample + + +# --------------------------------------------------------------------------- +# 6. Image mode (`LeRobotV3WriterConfig(video=False)`) — same gold-standard +# check, confirming stock LeRobotDataset reads inline PNG features too. +# --------------------------------------------------------------------------- + + +class TestUpstreamLoaderImageMode: + def test_metadata_loads(self, written_v3_image: Path): + from lerobot.datasets.lerobot_dataset import LeRobotDatasetMetadata + + meta = LeRobotDatasetMetadata( + repo_id="forge/synth_v3_image", + root=written_v3_image, + ) + assert meta.total_episodes == 3 + assert meta.fps == 30 + assert "observation.state" in meta.features + assert "action" in meta.features + assert meta.features["observation.images.top"]["dtype"] == "image" + + def test_dataset_loads_and_iterates(self, written_v3_image: Path): + from lerobot.datasets.lerobot_dataset import LeRobotDataset + + ds = LeRobotDataset( + repo_id="forge/synth_v3_image", + root=written_v3_image, + ) + assert len(ds) == 3 * 30 + sample = ds[0] + assert "observation.state" in sample + assert "action" in sample + assert "observation.images.top" in sample diff --git a/tests/test_lerobot_v3_writer.py b/tests/test_lerobot_v3_writer.py index 8e27a45..c76e0ff 100644 --- a/tests/test_lerobot_v3_writer.py +++ b/tests/test_lerobot_v3_writer.py @@ -21,6 +21,17 @@ def _check_dependencies_available() -> bool: return False +def _check_image_mode_dependencies_available() -> bool: + """Check dependencies for image mode (no `av` — video encoder isn't used).""" + try: + import PIL # noqa: F401 + import pyarrow # noqa: F401 + + return True + except ImportError: + return False + + @pytest.fixture def sample_episode() -> Episode: """Create a sample episode for testing.""" @@ -136,6 +147,11 @@ def test_custom_config(self): assert config.chunks_size == 500 assert config.camera_name_mapping == {"cam0": "observation.images.front"} + def test_video_defaults_true(self): + """Test that video mode is the default (backwards-compatible).""" + config = LeRobotV3WriterConfig() + assert config.video is True + class TestLeRobotV3Writer: """Tests for LeRobotV3Writer.""" @@ -346,3 +362,458 @@ def test_info_json_structure(self, tmp_path: Path, sample_episodes: list[Episode if feat.get("dtype") == "video": assert "video_info" in feat assert "video.fps" in feat["video_info"] + + +class TestLeRobotV3WriterImageMode: + """Tests for image mode (`LeRobotV3WriterConfig(video=False)`).""" + + @pytest.mark.skipif( + not _check_image_mode_dependencies_available(), + reason="PIL or PyArrow not installed", + ) + def test_no_video_directory_created(self, tmp_path: Path, sample_episode: Episode): + """Image mode never touches the video encoder or videos/ directory.""" + output_dir = tmp_path / "output" + config = LeRobotV3WriterConfig(fps=30.0, video=False) + writer = LeRobotV3Writer(config) + + writer.write_episode(sample_episode, output_dir, episode_index=0) + writer._flush_chunk(output_dir) + + assert not (output_dir / "videos").exists() + parquet_path = output_dir / "data" / "chunk-000" / "file-000.parquet" + assert parquet_path.exists() + + @pytest.mark.skipif( + not _check_image_mode_dependencies_available(), + reason="PIL or PyArrow not installed", + ) + def test_frames_stored_inline_as_png(self, tmp_path: Path, sample_episode: Episode): + """Camera columns hold PNG bytes decodable back to the source frame size.""" + import io + + import pyarrow.parquet as pq + from PIL import Image + + output_dir = tmp_path / "output" + config = LeRobotV3WriterConfig(fps=30.0, video=False) + writer = LeRobotV3Writer(config) + + writer.write_episode(sample_episode, output_dir, episode_index=0) + writer._flush_chunk(output_dir) + + table = pq.read_table(output_dir / "data" / "chunk-000" / "file-000.parquet") + rows = table.to_pylist() + assert len(rows) == 30 + + cam_col = "observation.images.camera0" + expected = [f.images["camera0"].load() for f in sample_episode.frames()] + for row, source in zip(rows, expected): + assert cam_col in row + assert row[cam_col]["path"] is None + img = Image.open(io.BytesIO(row[cam_col]["bytes"])) + assert img.format == "PNG" + assert img.size == (64, 64) + np.testing.assert_array_equal(np.asarray(img), source) + + @pytest.mark.skipif( + not _check_image_mode_dependencies_available(), + reason="PIL or PyArrow not installed", + ) + def test_float_frames_in_0_255_range_are_not_rescaled(self, tmp_path: Path): + """float32 frames already in 0-255 are clipped, not multiplied by 255.""" + import io + + import pyarrow.parquet as pq + from PIL import Image + + source = np.full((8, 8, 3), 200.0, dtype=np.float32) + + def frame_loader(): + yield Frame( + index=0, + timestamp=0.0, + images={ + "camera0": LazyImage( + loader=lambda: source, height=8, width=8, channels=3 + ), + }, + state=np.zeros(3, dtype=np.float32), + action=np.zeros(3, dtype=np.float32), + ) + + episode = Episode( + episode_id="ep_float", + language_instruction="float frames", + cameras={"camera0": CameraInfo(name="camera0", height=8, width=8)}, + fps=30.0, + _frame_loader=frame_loader, + ) + + output_dir = tmp_path / "output" + writer = LeRobotV3Writer(LeRobotV3WriterConfig(fps=30.0, video=False)) + writer.write_episode(episode, output_dir, episode_index=0) + writer._flush_chunk(output_dir) + + table = pq.read_table(output_dir / "data" / "chunk-000" / "file-000.parquet") + row = table.to_pylist()[0] + decoded = np.asarray(Image.open(io.BytesIO(row["observation.images.camera0"]["bytes"]))) + np.testing.assert_array_equal(decoded, np.full((8, 8, 3), 200, dtype=np.uint8)) + + @pytest.mark.skipif( + not _check_image_mode_dependencies_available(), + reason="PIL or PyArrow not installed", + ) + def test_info_json_dtype_is_image(self, tmp_path: Path, sample_episodes: list[Episode]): + """info.json declares image features with dtype 'image', not 'video'.""" + output_dir = tmp_path / "output" + config = LeRobotV3WriterConfig(fps=30.0, robot_type="test_robot", video=False) + writer = LeRobotV3Writer(config) + + writer.write_dataset(iter(sample_episodes), output_dir) + + with open(output_dir / "meta" / "info.json") as f: + info = json.load(f) + + cam_feature = info["features"]["observation.images.camera0"] + assert cam_feature["dtype"] == "image" + assert "video_info" not in cam_feature + assert cam_feature["fps"] == 30.0 + assert cam_feature["shape"] == [64, 64, 3] + + @pytest.mark.skipif( + not _check_image_mode_dependencies_available(), + reason="PIL or PyArrow not installed", + ) + def test_info_json_video_path_is_null(self, tmp_path: Path, sample_episodes: list[Episode]): + """info.json's video_path/video_files_size_in_mb are null in image + mode, matching upstream LeRobot's use_videos=False convention + (the keys stay present since v3 loaders require them to exist).""" + output_dir = tmp_path / "output" + config = LeRobotV3WriterConfig(fps=30.0, video=False) + writer = LeRobotV3Writer(config) + + writer.write_dataset(iter(sample_episodes), output_dir) + + with open(output_dir / "meta" / "info.json") as f: + info = json.load(f) + + assert "video_path" in info + assert info["video_path"] is None + assert "video_files_size_in_mb" in info + assert info["video_files_size_in_mb"] is None + + @pytest.mark.skipif( + not _check_image_mode_dependencies_available(), + reason="PIL or PyArrow not installed", + ) + def test_episodes_parquet_has_no_video_pointers( + self, tmp_path: Path, sample_episodes: list[Episode] + ): + """Image datasets carry no per-camera videos/ pointer columns.""" + import pyarrow.parquet as pq + + output_dir = tmp_path / "output" + config = LeRobotV3WriterConfig(fps=30.0, video=False) + writer = LeRobotV3Writer(config) + + writer.write_dataset(iter(sample_episodes), output_dir) + + episodes_path = output_dir / "meta" / "episodes" / "chunk-000" / "file-000.parquet" + table = pq.read_table(episodes_path) + video_cols = [c for c in table.column_names if c.startswith("videos/")] + assert video_cols == [] + + @pytest.mark.skipif( + not _check_image_mode_dependencies_available(), + reason="PIL or PyArrow not installed", + ) + def test_frame_count_matches_episode_length( + self, tmp_path: Path, sample_episodes: list[Episode] + ): + """Total frame count and per-episode lengths are unaffected by image mode.""" + import pyarrow.parquet as pq + + output_dir = tmp_path / "output" + config = LeRobotV3WriterConfig(fps=30.0, video=False) + writer = LeRobotV3Writer(config) + + writer.write_dataset(iter(sample_episodes), output_dir) + + with open(output_dir / "meta" / "info.json") as f: + info = json.load(f) + assert info["total_frames"] == sum(20 + i * 5 for i in range(3)) + + episodes_path = output_dir / "meta" / "episodes" / "chunk-000" / "file-000.parquet" + episodes_data = pq.read_table(episodes_path).to_pylist() + assert episodes_data[0]["length"] == 20 + + @pytest.mark.skipif( + not _check_image_mode_dependencies_available(), + reason="PIL or PyArrow not installed", + ) + def test_stats_cover_state_and_action( + self, tmp_path: Path, sample_episodes: list[Episode] + ): + """Normalization stats survive finalize's projected read of the data parquet.""" + import pyarrow.parquet as pq + + output_dir = tmp_path / "output" + writer = LeRobotV3Writer(LeRobotV3WriterConfig(fps=30.0, video=False)) + + writer.write_dataset(iter(sample_episodes), output_dir) + + with open(output_dir / "meta" / "stats.json") as f: + stats = json.load(f) + assert stats["observation.state"]["min"] == pytest.approx([0.1] * 7) + assert stats["observation.state"]["max"] == pytest.approx([0.1] * 7) + assert stats["action"]["mean"] == pytest.approx([0.01] * 7) + + episodes_path = output_dir / "meta" / "episodes" / "chunk-000" / "file-000.parquet" + episode_row = pq.read_table(episodes_path).to_pylist()[0] + assert episode_row["stats/observation.state/min"] == pytest.approx([0.1] * 7) + assert episode_row["stats/observation.state/count"] == [20] + + @pytest.mark.skipif( + not _check_image_mode_dependencies_available(), + reason="PIL or PyArrow not installed", + ) + def test_feature_shape_comes_from_decoded_frame(self, tmp_path: Path): + """Shape is taken from the decoded array, not the LazyImage's declared dims.""" + source = np.zeros((12, 20, 3), dtype=np.uint8) + + def frame_loader(): + yield Frame( + index=0, + timestamp=0.0, + images={ + "camera0": LazyImage( + loader=lambda: source, height=48, width=64, channels=3 + ), + }, + state=np.zeros(3, dtype=np.float32), + action=np.zeros(3, dtype=np.float32), + ) + + episode = Episode( + episode_id="ep_stale_dims", + language_instruction="stale declared dims", + cameras={"camera0": CameraInfo(name="camera0", height=48, width=64)}, + fps=30.0, + _frame_loader=frame_loader, + ) + + output_dir = tmp_path / "output" + writer = LeRobotV3Writer(LeRobotV3WriterConfig(fps=30.0, video=False)) + + writer.write_dataset(iter([episode]), output_dir) + + with open(output_dir / "meta" / "info.json") as f: + info = json.load(f) + assert info["features"]["observation.images.camera0"]["shape"] == [12, 20, 3] + + @pytest.mark.skipif( + not _check_image_mode_dependencies_available(), + reason="PIL or PyArrow not installed", + ) + def test_camera_missing_on_first_frame_still_written(self, tmp_path: Path): + """A camera absent from frame 0 is not dropped from the data parquet. + + pyarrow derives the parquet schema from the first row alone, so a + column that only shows up in later rows used to vanish silently while + info.json still declared the feature. + """ + import io + + import pyarrow.parquet as pq + from PIL import Image + + def make_frame(idx: int) -> Frame: + images = {} + if idx > 0: + images["camera0"] = LazyImage( + loader=lambda i=idx: np.full((8, 8, 3), i * 10, dtype=np.uint8), + height=8, + width=8, + channels=3, + ) + return Frame( + index=idx, + timestamp=idx / 30.0, + images=images, + state=np.zeros(3, dtype=np.float32), + action=np.zeros(3, dtype=np.float32), + ) + + def frame_loader(): + for i in range(4): + yield make_frame(i) + + episode = Episode( + episode_id="ep_late_camera", + language_instruction="camera appears late", + cameras={"camera0": CameraInfo(name="camera0", height=8, width=8)}, + fps=30.0, + _frame_loader=frame_loader, + ) + + output_dir = tmp_path / "output" + writer = LeRobotV3Writer(LeRobotV3WriterConfig(fps=30.0, video=False)) + writer.write_dataset(iter([episode]), output_dir) + + with open(output_dir / "meta" / "info.json") as f: + declared = set(json.load(f)["features"]) + rows = pq.read_table(output_dir / "data" / "chunk-000" / "file-000.parquet").to_pylist() + + cam_col = "observation.images.camera0" + assert cam_col in declared + assert declared <= set(rows[0]) + assert rows[0][cam_col] is None + decoded = np.asarray(Image.open(io.BytesIO(rows[2][cam_col]["bytes"]))) + np.testing.assert_array_equal(decoded, np.full((8, 8, 3), 20, dtype=np.uint8)) + + +class TestLeRobotV3WriterVideoModeUnchanged: + """Regression: video mode (the default) is untouched by image mode support.""" + + @pytest.mark.skipif( + not _check_dependencies_available(), + reason="PyAV or PyArrow not installed", + ) + def test_video_mode_still_writes_mp4_and_no_inline_bytes( + self, tmp_path: Path, sample_episode: Episode + ): + import pyarrow.parquet as pq + + output_dir = tmp_path / "output" + config = LeRobotV3WriterConfig(fps=30.0) # video=True by default + writer = LeRobotV3Writer(config) + + writer.write_episode(sample_episode, output_dir, episode_index=0) + writer._flush_chunk(output_dir) + + video_path = ( + output_dir / "videos" / "observation.images.camera0" / "chunk-000" / "file-000.mp4" + ) + assert video_path.exists() + + table = pq.read_table(output_dir / "data" / "chunk-000" / "file-000.parquet") + assert "observation.images.camera0" not in table.column_names + + @pytest.mark.skipif( + not _check_dependencies_available(), + reason="PyAV or PyArrow not installed", + ) + def test_video_mode_info_json_dtype_is_video( + self, tmp_path: Path, sample_episodes: list[Episode] + ): + output_dir = tmp_path / "output" + config = LeRobotV3WriterConfig(fps=30.0, robot_type="test_robot") + writer = LeRobotV3Writer(config) + + writer.write_dataset(iter(sample_episodes), output_dir) + + with open(output_dir / "meta" / "info.json") as f: + info = json.load(f) + + cam_feature = info["features"]["observation.images.camera0"] + assert cam_feature["dtype"] == "video" + assert "video_info" in cam_feature + + +class TestLeRobotV3WriterSparseColumns: + """Columns whose first row is null must not break stats or feature inference.""" + + @pytest.mark.skipif( + not _check_dependencies_available(), + reason="PyAV or PyArrow not installed", + ) + def test_action_missing_on_first_frame_still_finalizes(self, tmp_path: Path): + """An action absent from frame 0 must not fail stats computation. + + MCAP/RLDS readers emit ``action=None`` for frames with no action + sample, so a column can legitimately be null at row 0 and populated + later. + """ + + def frame_loader(): + for i in range(4): + yield Frame( + index=i, + timestamp=i / 30.0, + images={}, + state=np.full(3, float(i), dtype=np.float32), + action=None if i == 0 else np.full(3, float(i), dtype=np.float32), + ) + + episode = Episode( + episode_id="ep_sparse_action", + language_instruction="action appears late", + cameras={}, + fps=30.0, + _frame_loader=frame_loader, + ) + + output_dir = tmp_path / "output" + writer = LeRobotV3Writer(LeRobotV3WriterConfig(fps=30.0)) + writer.write_dataset(iter([episode]), output_dir) + + with open(output_dir / "meta" / "stats.json") as f: + stats = json.load(f) + assert stats["action"]["count"] == [3] + assert stats["action"]["min"] == [1.0, 1.0, 1.0] + assert stats["action"]["max"] == [3.0, 3.0, 3.0] + + @pytest.mark.skipif( + not _check_image_mode_dependencies_available(), + reason="PIL or PyArrow not installed", + ) + def test_parallel_inference_uses_first_non_null_row(self, tmp_path: Path): + """Parallel-mode feature inference reads past a null first row. + + Workers write the data parquet themselves, so ``finalize`` infers + features from the file. A camera missing on frame 0 leaves a null + there; keying off row 0 alone declared it as float32. + """ + import io + + import pyarrow as pa + import pyarrow.parquet as pq + from PIL import Image + + from forge.core.models import DatasetInfo + + def png(value: int) -> bytes: + buf = io.BytesIO() + Image.fromarray(np.full((8, 12, 3), value, dtype=np.uint8)).save(buf, format="PNG") + return buf.getvalue() + + cam_col = "observation.images.camera0" + rows = [ + { + "episode_index": 0, + "frame_index": i, + "index": i, + "task_index": 0, + "timestamp": i / 30.0, + "observation.state": None if i == 0 else [0.0, 0.0, 0.0], + cam_col: None if i == 0 else {"bytes": png(i * 10), "path": None}, + } + for i in range(3) + ] + data_dir = tmp_path / "output" / "data" / "chunk-000" + data_dir.mkdir(parents=True) + pq.write_table(pa.Table.from_pylist(rows), data_dir / "file-000.parquet") + + writer = LeRobotV3Writer(LeRobotV3WriterConfig(fps=30.0, video=False)) + writer._total_frames = 3 + writer.finalize( + tmp_path / "output", + DatasetInfo(path=tmp_path, format="test", num_episodes=1, total_frames=3), + ) + + with open(tmp_path / "output" / "meta" / "info.json") as f: + features = json.load(f)["features"] + assert features[cam_col]["dtype"] == "image" + assert features[cam_col]["shape"] == [8, 12, 3] + assert features["observation.state"]["shape"] == [3]