diff --git a/.gitignore b/.gitignore index 9e1cc96..c646350 100644 --- a/.gitignore +++ b/.gitignore @@ -17,3 +17,4 @@ output/ .env !tests/files/* .idea +uv.lock \ No newline at end of file diff --git a/README.md b/README.md index f35e961..048fa8a 100644 --- a/README.md +++ b/README.md @@ -2261,6 +2261,61 @@ results = client.import_lerobot( > **Note:** Only LeRobot dataset v3 is supported. v2 datasets need to be converted to v3 before importing. +##### Customize the import with a converter + +By default, each frame's full `observation.state` and `action` are written to the telemetry JSON and every camera is uploaded. To control this, subclass `LeRobotConverter` and pass the class (the SDK constructs it with the parsed `meta/info.json`, so you can select by feature name): + +```python +from fastlabel.lerobot import LeRobotConverter + + +class JointsOnlyConverter(LeRobotConverter): + # keep only the 4 main cameras (drop tactile gel cameras, etc.) + CAMERA_KEYS = ("cam_high", "cam_low", "cam_left_wrist", "cam_right_wrist") + + # keep only joint values (names ending in ".pos") out of a large state vector + def select_observation_state_names(self, names): + return [n for n in names if n.endswith(".pos")] + + def select_action_names(self, names): + return [n for n in names if n.endswith(".pos")] + + # add an extra telemetry item + def build_telemetry_frame(self, frame): + telemetry = super().build_telemetry_frame(frame) + telemetry["gripper"] = self.build_action(frame)[-1] + return telemetry + + # set tags / metadata on each created task + def build_task_kwargs(self, *, episode_index, episode_name, frames): + return { + "tags": ["lerobot", self.meta.get("robot_type", "unknown")], + "metadatas": [{"key": "num_frames", "value": str(len(frames))}], + } + + +results = client.import_lerobot( + project="YOUR_PROJECT_SLUG", + lerobot_data_path="/path/to/lerobot/dataset", + converter=JointsOnlyConverter, +) +``` + +Overridable hooks: + +| Hook | Purpose | +| --- | --- | +| `OBSERVATION_STATE_NAMES` / `ACTION_NAMES` (class vars, tuples) | Keep only the named values (declared order; `None` = all). | +| `select_observation_state_names` / `select_action_names` | Pick kept feature names dynamically from the given name list (e.g. by suffix). | +| `build_observation_state` / `build_action` | Transform the kept values. | +| `build_telemetry_frame` | Shape each telemetry frame (call `super()` to extend). | +| `select_cameras` | Choose which cameras to upload. | +| `build_task_kwargs` | Keyword args forwarded to `create_robotics_task` (tags, metadatas, ...). | +| `select_episodes` | Which episodes to import when `episode_indices` is omitted. | +| `build_episode_name` | Task / artifact naming. | + +For constructor arguments, pass a factory instead of the class, e.g. `converter=functools.partial(MyConverter, threshold=3)`. + #### Find Task Find a single task. diff --git a/examples/import_lerobot.py b/examples/import_lerobot.py index 387a186..0105fdf 100644 --- a/examples/import_lerobot.py +++ b/examples/import_lerobot.py @@ -8,10 +8,11 @@ """ from fastlabel import Client +from fastlabel.lerobot import LeRobotConverter client = Client() -# Import all episodes +# Import all episodes (default: full telemetry, all cameras) results = client.import_lerobot( project="your-project-slug", lerobot_data_path="/path/to/lerobot/dataset", @@ -23,3 +24,62 @@ lerobot_data_path="/path/to/lerobot/dataset", episode_indices=[0, 1, 2], ) + + +# --- Customize with a converter -------------------------------------------- +# A converter controls which values go into the telemetry JSON, which cameras +# are uploaded, task naming and task keyword args. The SDK constructs it with +# the parsed meta/info.json, so selection can be done by feature name. +# +# The example dataset below has a 12014-dim observation.state; this keeps only +# the joint values (names ending in ".pos") and the 4 main cameras. +class JointsOnlyConverter(LeRobotConverter): + CAMERA_KEYS = ("cam_high", "cam_low", "cam_left_wrist", "cam_right_wrist") + + def select_observation_state_names(self, names): + return [n for n in names if n.endswith(".pos")] + + def select_action_names(self, names): + return [n for n in names if n.endswith(".pos")] + + def build_telemetry_frame(self, frame): + telemetry = super().build_telemetry_frame(frame) + telemetry["gripper"] = self.build_action(frame)[-1] # add an extra item + return telemetry + + def build_task_kwargs(self, *, episode_index, episode_name, frames): + return { + "tags": ["lerobot", self.meta.get("robot_type", "unknown")], + "metadatas": [{"key": "num_frames", "value": str(len(frames))}], + } + + +# Pass the class itself (the SDK constructs it with meta): +results = client.import_lerobot( + project="your-project-slug", + lerobot_data_path="/path/to/lerobot/dataset", + converter=JointsOnlyConverter, +) + +# For constructor arguments, pass a factory: +# from functools import partial +# converter=partial(MyConverter, threshold=3) + + +# Alternatively, declare exact names statically (no method override needed): +class StaticJointsConverter(LeRobotConverter): + OBSERVATION_STATE_NAMES = ( + "left_joint_0.pos", + "left_joint_1.pos", + # ... list the joint names you want ... + "right_joint_5.pos", + "right_left_carriage_joint.pos", + ) + CAMERA_KEYS = ("cam_high", "cam_low", "cam_left_wrist", "cam_right_wrist") + + +# Skip short episodes (select_episodes receives {episode_index: frame_num} and +# is only used when episode_indices is not passed): +class LongEpisodesOnly(LeRobotConverter): + def select_episodes(self, episode_lengths): + return [i for i, frame_num in episode_lengths.items() if frame_num >= 100] diff --git a/fastlabel/__init__.py b/fastlabel/__init__.py index ff80eed..923c5cb 100644 --- a/fastlabel/__init__.py +++ b/fastlabel/__init__.py @@ -3,11 +3,11 @@ import logging import os import re -import shutil +import tempfile import urllib.parse from concurrent.futures import ThreadPoolExecutor, wait from pathlib import Path -from typing import Dict, List, Literal, Optional, Union +from typing import Any, Callable, Dict, List, Literal, Optional, Union import cv2 import numpy as np @@ -2256,8 +2256,11 @@ def import_lerobot( self, project: str, lerobot_data_path: str, - episode_indices: list = None, - ) -> list: + episode_indices: list[int] | None = None, + converter: Callable[ + [dict[str, Any]], lerobot.LeRobotConverter + ] = lerobot.LeRobotConverter, + ) -> list[dict[str, Any]]: """ Import a LeRobot dataset into a FastLabel robotics project. @@ -2271,42 +2274,75 @@ def import_lerobot( lerobot_data_path is the path to the LeRobot dataset directory (Required). episode_indices is a list of episode indices to import. If None, all episodes are imported (Optional). + converter is a LeRobotConverter subclass (or any callable taking the + parsed meta/info.json dict and returning a LeRobotConverter). The SDK + constructs it with the dataset metadata. Override its hooks to control + telemetry values (build_observation_state / build_action / + build_telemetry_frame), video selection (select_cameras), task + keyword args (build_task_kwargs) and episode selection + (select_episodes). Defaults to LeRobotConverter (unchanged behaviour). """ data_path = Path(lerobot_data_path) + conv = converter(lerobot.load_info(data_path)) episode_map = lerobot.build_episode_map(data_path) if episode_indices is None: - episode_indices = sorted(episode_map.keys()) + episode_indices = conv.select_episodes( + {index: info["length"] for index, info in episode_map.items()} + ) results = [] for episode_index in episode_indices: - episode_name = lerobot.format_episode_name(episode_index) - self.create_robotics_task(project=project, name=episode_name) - - zip_path = lerobot.create_episode_zip( - lerobot_data_path=data_path, - episode_index=episode_index, - episode_map=episode_map, - ) + episode_name = None + # Build the ZIP before creating the task so a telemetry/video failure + # does not leave an orphan task. A per-episode FastLabelException is + # recorded and the remaining episodes still run; unexpected errors + # (e.g. connection errors) propagate so the caller can react. try: - result = self.import_robotics_contents_file( - project=project, file_path=zip_path + episode_name = conv.build_episode_name(episode_index) + raw_frames = lerobot.get_episode_raw_frames( + data_path, episode_index, episode_map + ) + task_kwargs = ( + conv.build_task_kwargs( + episode_index=episode_index, + episode_name=episode_name, + frames=raw_frames, + ) + or {} ) + with tempfile.TemporaryDirectory() as episode_dir: + zip_path = lerobot.create_episode_zip( + lerobot_data_path=data_path, + episode_index=episode_index, + episode_name=episode_name, + converter=conv, + episode_map=episode_map, + raw_frames=raw_frames, + output_dir=Path(episode_dir), + ) + self.create_robotics_task( + project=project, name=episode_name, **task_kwargs + ) + result = self.import_robotics_contents_file( + project=project, file_path=zip_path + ) results.append( - {"episode": episode_name, "success": True, "result": result} + { + "episode_index": episode_index, + "episode": episode_name, + "success": True, + "result": result, + } ) except FastLabelException as e: results.append( { + "episode_index": episode_index, "episode": episode_name, "success": False, "result": {"error": str(e)}, } ) - finally: - zip_file = Path(zip_path) - tmp_dir = zip_file.parent - if tmp_dir.exists(): - shutil.rmtree(tmp_dir) return results diff --git a/fastlabel/lerobot/__init__.py b/fastlabel/lerobot/__init__.py index 9654c44..287c011 100644 --- a/fastlabel/lerobot/__init__.py +++ b/fastlabel/lerobot/__init__.py @@ -1,23 +1,29 @@ +from pathlib import Path +from typing import Any + from fastlabel.exceptions import FastLabelInvalidException from fastlabel.lerobot import v3 from fastlabel.lerobot.common import ( + Camera, check_dependencies, detect_version, - format_episode_name, - get_camera_dirs, + load_info, ) +from fastlabel.lerobot.converter import LeRobotConverter +from fastlabel.lerobot.v3 import EpisodeMap, Frame __all__ = [ + "Camera", + "LeRobotConverter", "build_episode_map", "get_episode_indices", + "get_episode_raw_frames", "create_episode_zip", - "format_episode_name", - "get_camera_dirs", + "load_info", ] -def get_episode_indices(lerobot_data_path): - """Get all episode indices from a LeRobot v3 dataset.""" +def _ensure_v3(lerobot_data_path: Path) -> None: check_dependencies() version = detect_version(lerobot_data_path) if version == "v2": @@ -25,38 +31,75 @@ def get_episode_indices(lerobot_data_path): "LeRobot dataset v2 is not supported. Please convert to v3.", 422, ) + + +def get_episode_indices(lerobot_data_path: Path) -> list[int]: + """Get all episode indices from a LeRobot v3 dataset.""" + _ensure_v3(lerobot_data_path) return v3.get_episode_indices(lerobot_data_path) -def build_episode_map(lerobot_data_path): +def build_episode_map(lerobot_data_path: Path) -> EpisodeMap: """Build episode map from dataset. Returns a dict keyed by episode index.""" - check_dependencies() - version = detect_version(lerobot_data_path) - if version == "v2": - raise FastLabelInvalidException( - "LeRobot dataset v2 is not supported. Please convert to v3.", - 422, - ) + _ensure_v3(lerobot_data_path) return v3._build_episode_map(lerobot_data_path) -def create_episode_zip(lerobot_data_path, episode_index, episode_map=None): +def get_episode_raw_frames( + lerobot_data_path: Path, + episode_index: int, + episode_map: EpisodeMap, +) -> list[Frame]: + """Return every frame of an episode as native-Python dicts (all columns). + + Supports LeRobot dataset v3 only. The v3 check is not repeated here because + ``episode_map`` can only be obtained from ``build_episode_map`` (which runs + it once). + """ + 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"] + ) + + +def create_episode_zip( + lerobot_data_path: Path, + episode_index: int, + episode_name: str, + converter: LeRobotConverter, + output_dir: Path, + episode_map: EpisodeMap, + raw_frames: list[Frame], +) -> str: """Create a ZIP file for a single episode in the format expected by FastLabel. Supports LeRobot dataset v3 only. ZIP structure (files at root, ZIP name = episode name): - {content_name}.mp4 (one per camera) - {episode_name}.json (frame data) + {content_name}.mp4 (one per selected camera) + {episode_name}.json (telemetry frame data) - Returns the path to the created ZIP file. - The caller is responsible for cleaning up the returned ZIP file. + episode_name is the identifier for the task / JSON / ZIP; the caller passes + it (from ``converter.build_episode_name``) so the task and artifact names + stay consistent. + converter is a LeRobotConverter instance whose hooks control telemetry + (build_telemetry_frame) and video selection (select_cameras). Pass + ``LeRobotConverter()`` for default behaviour (all telemetry keys, all + cameras). raw_frames is the episode's raw frames (from + get_episode_raw_frames); it is reused here to avoid re-reading the parquet. + + output_dir is the directory the ZIP is written to; pass a directory you + manage (e.g. a ``tempfile.TemporaryDirectory``) so it is cleaned up + automatically. + + Returns the path to the created ZIP file (written under output_dir). + The v3 check is not repeated here (``episode_map`` implies it already ran). """ - check_dependencies() - version = detect_version(lerobot_data_path) - if version == "v2": - raise FastLabelInvalidException( - "LeRobot dataset v2 is not supported. Please convert to v3.", - 422, - ) - return v3.create_episode_zip(lerobot_data_path, episode_index, episode_map) + telemetry_frames: list[dict[str, Any]] = [ + converter.build_telemetry_frame(frame) for frame in raw_frames + ] + cameras = converter.select_cameras(v3.get_camera_dirs(lerobot_data_path)) + ep_info = v3.resolve_episode(episode_index, episode_map) + return v3._assemble_episode_zip( + ep_info, episode_name, cameras, telemetry_frames, output_dir + ) diff --git a/fastlabel/lerobot/common.py b/fastlabel/lerobot/common.py index d034a60..c967d19 100644 --- a/fastlabel/lerobot/common.py +++ b/fastlabel/lerobot/common.py @@ -1,9 +1,38 @@ +import json from pathlib import Path +from typing import Any, NamedTuple from fastlabel.exceptions import FastLabelInvalidException -def check_dependencies(): +class Camera(NamedTuple): + """A dataset camera. Shared contract between data access and converter. + + path is the video directory (v3: videos/observation.images.X). + key is the observation feature key (e.g. observation.images.cam_high), + matching ``meta["features"][key]`` so converters can look up resolution + etc. + content_name is the mp4 filename stem inside the ZIP (e.g. images_cam_high). + """ + + path: Path + key: str + content_name: str + + +def load_info(lerobot_data_path: Path) -> dict[str, Any]: + """Load meta/info.json for a LeRobot dataset. + + Returns the parsed dict, or {} when the file is absent (older datasets), + so name-based selection is opt-in and default behavior is unaffected. + """ + info_path = lerobot_data_path / "meta" / "info.json" + if not info_path.exists(): + return {} + return json.loads(info_path.read_text()) + + +def check_dependencies() -> None: try: import pandas # noqa: F401 import pyarrow # noqa: F401 @@ -43,31 +72,3 @@ def detect_version(lerobot_data_path: Path) -> str: "or data/chunk-XXX/file-*.parquet (v3).", 422, ) - - -def format_episode_name(episode_index: int) -> str: - return f"episode_{episode_index:06d}" - - -def get_camera_dirs(lerobot_data_path: Path) -> list: - """Get camera directories and their content names. - Returns [(camera_dir, content_name), ...]. - e.g. observation.images.top -> content_name = "images_top" - """ - videos_dir = lerobot_data_path / "videos" - if not videos_dir.exists(): - return [] - - results = [] - for obs_dir in sorted(videos_dir.iterdir()): - if not obs_dir.is_dir(): - continue - parts = obs_dir.name.split(".") - if parts[0] != "observation": - raise FastLabelInvalidException( - f"Unexpected camera dir name: {obs_dir.name}" - ) - - content_name = "_".join(parts[1:]) - results.append((obs_dir, content_name)) - return results diff --git a/fastlabel/lerobot/converter.py b/fastlabel/lerobot/converter.py new file mode 100644 index 0000000..367bed7 --- /dev/null +++ b/fastlabel/lerobot/converter.py @@ -0,0 +1,171 @@ +from collections.abc import Callable +from typing import Annotated, Any, ClassVar + +from fastlabel.exceptions import FastLabelInvalidException +from fastlabel.lerobot.common import Camera + +# Format-agnostic domain types. Deliberately generic (not v3-specific) so this +# module stays decoupled from any dataset version. +Meta = dict[str, Any] +Frame = dict[str, Any] +EpisodeIndex = Annotated[int, "episode_index"] +FrameNum = Annotated[int, "frame_num"] +# A name selector: given the full ordered feature names, return the ones to keep. +NameSelector = Callable[[list[str]], list[str]] + + +class LeRobotConverter: + """Customization hooks for ``Client.import_lerobot``. + + Override only the pieces you need; the overall flow (episode loop, task + creation, upload, cleanup) is owned by ``Client.import_lerobot`` and is not + overridable. A converter is a pure strategy object: everything is resolved + once in ``__init__`` from the dataset metadata (``meta/info.json``) and never + mutated afterwards. + + Instances are constructed by the SDK, not the caller. Pass the class itself + (or a ``functools.partial`` / factory) to ``import_lerobot``; the SDK calls + it with the parsed ``info.json`` as ``meta``. + + Selection of ``observation.state`` / ``action`` values is name-based and + follows a 3-layer structure (requires ``meta/info.json``): + + 1. Declare exact names via the ``OBSERVATION_STATE_NAMES`` / ``ACTION_NAMES`` + class variables (``None`` means keep everything). + 2. Or override ``select_observation_state_names`` / ``select_action_names`` to + pick names dynamically from the given full name list (e.g. by suffix). + 3. ``build_observation_state`` / ``build_action`` return the kept values; the + SDK resolves the selected names to indices once in ``__init__``. + """ + + # ---- declarative selection (static configuration) ---- + OBSERVATION_STATE_NAMES: ClassVar[tuple[str, ...] | None] = None + ACTION_NAMES: ClassVar[tuple[str, ...] | None] = None + CAMERA_KEYS: ClassVar[tuple[str, ...] | None] = None + + def __init__(self, meta: Meta | None = None) -> None: + self.meta: Meta = meta or {} + # Selected names -> positions, resolved once and immutable afterwards. + self._state_index: list[int] = self._names_to_index( + "observation.state", self.select_observation_state_names + ) + self._action_index: list[int] = self._names_to_index( + "action", self.select_action_names + ) + + # ---- value selection (by name; override for dynamic selection) ---- + def select_observation_state_names(self, names: list[str]) -> list[str]: + """Return the ``observation.state`` feature names to keep (default: all). + + ``names`` is the full ordered name list from ``meta/info.json``. Override + for dynamic selection, e.g. ``[n for n in names if n.endswith(".pos")]``. + """ + return ( + list(self.OBSERVATION_STATE_NAMES) + if self.OBSERVATION_STATE_NAMES + else names + ) + + def select_action_names(self, names: list[str]) -> list[str]: + """Return the ``action`` feature names to keep (default: all).""" + return list(self.ACTION_NAMES) if self.ACTION_NAMES else names + + def _names_to_index(self, key: str, select_names: NameSelector) -> list[int]: + """Resolve the selected names to their positions in ``key``'s full name + list (from ``self.meta``). Raises when the metadata lacks the names or a + selected name is absent. + """ + try: + all_names: list[str] = self.meta["features"][key]["names"] + except (KeyError, TypeError): + raise FastLabelInvalidException( + f"'{key}.names' not found in meta/info.json.", 422 + ) + position = {name: i for i, name in enumerate(all_names)} + index: list[int] = [] + for name in select_names(all_names): + if name not in position: + raise FastLabelInvalidException( + f"'{name}' not found in '{key}.names'.", 422 + ) + index.append(position[name]) + return index + + # ---- telemetry hooks ---- + def build_observation_state(self, frame: Frame) -> list[Any]: + values = frame["observation.state"] + return [values[i] for i in self._state_index] + + def build_action(self, frame: Frame) -> list[Any]: + values = frame["action"] + return [values[i] for i in self._action_index] + + def build_telemetry_frame(self, frame: Frame) -> dict[str, Any]: + """Build one telemetry frame written to the episode JSON. + + Override and call ``super()`` to add extra items (e.g. gripper). + """ + return { + "observation.state": self.build_observation_state(frame), + "action": self.build_action(frame), + "frame_index": int(frame["frame_index"]), + "timestamp": float(frame["timestamp"]), + } + + # ---- video hook ---- + def select_cameras(self, cameras: list[Camera]) -> list[Camera]: + """Return the cameras to include (default: all — the given list). + + Each ``Camera`` has ``path`` / ``key`` / ``content_name``; ``key`` matches + ``self.meta["features"][key]`` for looking up resolution etc. When + ``CAMERA_KEYS`` is set, keep cameras whose key suffix (e.g. ``cam_high`` + of ``observation.images.cam_high``) is in the set. + """ + if self.CAMERA_KEYS is None: + return cameras + return [ + camera + for camera in cameras + if camera.key.split(".")[-1] in self.CAMERA_KEYS + ] + + # ---- task hook ---- + def build_task_kwargs( + self, + *, + episode_index: int, + episode_name: str, + frames: list[Frame], + ) -> dict[str, Any]: + """Return keyword args forwarded to ``create_robotics_task``. + + Keyword-only: + episode_index is the episode index. + episode_name is the task name (e.g. ``episode_000001``). + frames is the list of raw native frames for the episode. + Dataset-level metadata is available via ``self.meta``. + e.g. return ``{"tags": [...], "metadatas": [...]}``. + """ + return {} + + # ---- episode hooks ---- + def select_episodes( + self, episode_lengths: dict[EpisodeIndex, FrameNum] + ) -> list[EpisodeIndex]: + """Return the episode indices to import (default: all). + + Only called when ``import_lerobot`` is invoked without + ``episode_indices``. ``episode_lengths`` maps each episode index to its + frame count, e.g. to skip short episodes. Kept intentionally minimal + (index -> frame count, no v3 layout details) so this class stays + decoupled from any dataset version. + """ + return sorted(episode_lengths) + + def build_episode_name(self, episode_index: int) -> str: + """Identifier used for the task name, episode JSON name and ZIP name. + + Defaults to the ``episode_000001`` form. Override for custom naming + (keep it filesystem-safe, as it is also used for artifact filenames). + """ + return f"episode_{episode_index:06d}" diff --git a/fastlabel/lerobot/v3.py b/fastlabel/lerobot/v3.py index 9a51801..c564de7 100644 --- a/fastlabel/lerobot/v3.py +++ b/fastlabel/lerobot/v3.py @@ -2,15 +2,31 @@ import shutil import tempfile from pathlib import Path +from typing import Any, TypedDict import cv2 from fastlabel.exceptions import FastLabelInvalidException -from fastlabel.lerobot.common import format_episode_name, get_camera_dirs +from fastlabel.lerobot.common import Camera -def _build_episode_map(lerobot_data_path: Path) -> dict: - """Build a mapping of episode_index -> {chunk, file_stem, frame_offset, length}. +class EpisodeInfo(TypedDict): + """Per-episode location within the v3 layout.""" + + chunk: str + file_stem: str + frame_offset: int + length: int + + +# episode_index -> EpisodeInfo +EpisodeMap = dict[int, EpisodeInfo] +# One frame's parquet columns converted to native Python (JSON-serialisable). +Frame = dict[str, Any] + + +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). @@ -20,7 +36,7 @@ def _build_episode_map(lerobot_data_path: Path) -> dict: import pandas as pd data_dir = lerobot_data_path / "data" - episode_map = {} + episode_map: EpisodeMap = {} for chunk_dir in sorted(data_dir.iterdir()): if not chunk_dir.is_dir() or not chunk_dir.name.startswith("chunk-"): @@ -46,35 +62,69 @@ def _build_episode_map(lerobot_data_path: Path) -> dict: return episode_map -def get_episode_indices(lerobot_data_path: Path) -> list: +def get_episode_indices(lerobot_data_path: Path) -> list[int]: """Get all episode indices from a v3 dataset.""" episode_map = _build_episode_map(lerobot_data_path) return sorted(episode_map.keys()) -def _convert_episode_frames( +def get_camera_dirs(lerobot_data_path: Path) -> list[Camera]: + """Get camera directories and their content names (v3 video layout). + + v3: videos/{observation.images.X}/chunk-XXX/file-YYY.mp4 + Returns [(camera_dir, content_name), ...]. + e.g. observation.images.top -> content_name = "images_top" + """ + videos_dir = lerobot_data_path / "videos" + if not videos_dir.exists(): + return [] + + results: list[Camera] = [] + for obs_dir in sorted(videos_dir.iterdir()): + if not obs_dir.is_dir(): + continue + parts = obs_dir.name.split(".") + if parts[0] != "observation": + raise FastLabelInvalidException( + f"Unexpected camera dir name: {obs_dir.name}" + ) + + content_name = "_".join(parts[1:]) + results.append( + Camera(path=obs_dir, key=obs_dir.name, content_name=content_name) + ) + return results + + +def _row_to_native(row: Any) -> Frame: + """Convert a pandas row to a dict of JSON-serialisable native Python values. + + numpy arrays become lists and numpy scalars become Python scalars so the + result is safe to hand to converter hooks and to json.dumps. + """ + return { + key: (value.tolist() if hasattr(value, "tolist") else value) + for key, value in row.items() + } + + +def get_episode_raw_frames( lerobot_data_path: Path, episode_index: int, chunk: str, file_stem: str -) -> list: - """Extract frame dicts for a single episode from a v3 consolidated parquet.""" +) -> list[Frame]: + """Return every frame of a single episode as native-Python dicts. + + Each dict contains all parquet columns for the row (observation.state, + action, frame_index, timestamp, episode_index, ...). This is the raw data + passed to the converter hooks. + """ import pandas as pd parquet_path = lerobot_data_path / "data" / chunk / f"{file_stem}.parquet" - df = pd.read_parquet(parquet_path) + # Predicate pushdown so a consolidated file holding many episodes only reads + # the relevant row groups instead of the whole file per episode. + df = pd.read_parquet(parquet_path, filters=[("episode_index", "==", episode_index)]) ep_df = df[df["episode_index"] == episode_index] - - required_keys = ["observation.state", "action", "frame_index", "timestamp"] - if not all(key in ep_df.columns for key in required_keys): - return [] - - return [ - { - "observation.state": row["observation.state"].tolist(), - "action": row["action"].tolist(), - "frame_index": int(row["frame_index"]), - "timestamp": float(row["timestamp"]), - } - for _, row in ep_df.iterrows() - ] + return [_row_to_native(row) for _, row in ep_df.iterrows()] def _extract_video_segment( @@ -104,52 +154,55 @@ def _extract_video_segment( cap.release() -def create_episode_zip( - lerobot_data_path: Path, episode_index: int, episode_map: dict = None +def _assemble_episode_zip( + ep_info: EpisodeInfo, + episode_name: str, + cameras: list[Camera], + telemetry_frames: list[dict[str, Any]], + output_dir: Path, ) -> str: - """Create a ZIP for a single v3 episode. + """Stage selected video segments + telemetry JSON, then archive as a ZIP. - v3 video layout: videos/{key}/chunk-XXX/file-YYY.mp4 + ``telemetry_frames`` is the list of frame dicts written to the episode JSON. + The staging directory is removed automatically; the ZIP is written under + ``output_dir`` (owned by the caller) and its path returned. """ - if episode_map is None: - episode_map = _build_episode_map(lerobot_data_path) - - if episode_index not in episode_map: - raise FastLabelInvalidException( - f"Episode index {episode_index} not found in dataset.", - 422, - ) - - ep_info = episode_map[episode_index] chunk = ep_info["chunk"] file_stem = ep_info["file_stem"] frame_offset = ep_info["frame_offset"] length = ep_info["length"] - episode_name = format_episode_name(episode_index) - tmp_dir = tempfile.mkdtemp() - content_dir = Path(tmp_dir) / "content" - content_dir.mkdir() + with tempfile.TemporaryDirectory() as staging: + content_dir = Path(staging) + + # Extract video segments + # v3: videos/{key}/chunk-XXX/file-YYY.mp4 + for camera in cameras: + video_path = camera.path / chunk / f"{file_stem}.mp4" + if not video_path.exists(): + continue + output_path = content_dir / f"{camera.content_name}.mp4" + _extract_video_segment(video_path, frame_offset, length, output_path) + + json_path = content_dir / f"{episode_name}.json" + json_path.write_text(json.dumps(telemetry_frames, ensure_ascii=False)) + + # Create ZIP (files at root, ZIP name = episode name) + return shutil.make_archive( + base_name=str(output_dir / episode_name), + format="zip", + root_dir=str(content_dir), + ) - # Extract video segments - # v3: videos/{key}/chunk-XXX/file-YYY.mp4 - for camera_dir, content_name in get_camera_dirs(lerobot_data_path): - video_path = camera_dir / chunk / f"{file_stem}.mp4" - if not video_path.exists(): - continue - output_path = content_dir / f"{content_name}.mp4" - _extract_video_segment(video_path, frame_offset, length, output_path) - - # Convert parquet to JSON - frames = _convert_episode_frames(lerobot_data_path, episode_index, chunk, file_stem) - json_path = content_dir / f"{episode_name}.json" - json_path.write_text(json.dumps(frames, ensure_ascii=False)) - - # Create ZIP (files at root, ZIP name = episode name) - zip_path = shutil.make_archive( - base_name=str(Path(tmp_dir) / episode_name), - format="zip", - root_dir=str(content_dir), - ) - shutil.rmtree(content_dir) - return zip_path + +def resolve_episode( + episode_index: int, + episode_map: EpisodeMap, +) -> EpisodeInfo: + """Look up an episode's info in the episode map (raises if absent).""" + if episode_index not in episode_map: + raise FastLabelInvalidException( + f"Episode index {episode_index} not found in dataset.", + 422, + ) + return episode_map[episode_index] diff --git a/tests/test_lerobot_converter.py b/tests/test_lerobot_converter.py new file mode 100644 index 0000000..f07e18e --- /dev/null +++ b/tests/test_lerobot_converter.py @@ -0,0 +1,187 @@ +"""Tests for LeRobotConverter hooks and load_info. + +These are pure-logic tests (no pandas/opencv needed): the converter operates on +plain frame dicts and camera tuples, independent of the v3 data-access layer. +""" + +import json +from pathlib import Path + +import pytest + +from fastlabel.exceptions import FastLabelInvalidException +from fastlabel.lerobot import LeRobotConverter, load_info +from fastlabel.lerobot.common import Camera + +META = { + "robot_type": "bi_widowxai", + "features": { + "observation.state": {"names": ["j0.pos", "sensor_fx", "j1.pos"]}, + "action": {"names": ["a.pos", "b"]}, + }, +} + +FRAME = { + "observation.state": [10.0, 20.0, 30.0], + "action": [7.0, 8.0], + "frame_index": 3, + "timestamp": 0.3, + "episode_index": 0, +} + + +class TestDefaults: + def test_passthrough_keeps_all_values(self): + c = LeRobotConverter(META) + assert c._state_index == [0, 1, 2] + assert c._action_index == [0, 1] + assert c.build_telemetry_frame(FRAME) == { + "observation.state": [10.0, 20.0, 30.0], + "action": [7.0, 8.0], + "frame_index": 3, + "timestamp": 0.3, + } + + def test_no_meta_raises(self): + with pytest.raises(FastLabelInvalidException): + LeRobotConverter() + + def test_select_cameras_keeps_all_by_default(self): + cams = [ + Camera( + Path("videos/observation.images.cam_high"), + "observation.images.cam_high", + "images_cam_high", + ) + ] + assert LeRobotConverter(META).select_cameras(cams) == cams + + def test_select_episodes_returns_all_sorted_by_default(self): + assert LeRobotConverter(META).select_episodes({2: 5, 0: 3, 1: 4}) == [0, 1, 2] + + def test_build_episode_name_default_format(self): + assert LeRobotConverter(META).build_episode_name(7) == "episode_000007" + + def test_build_task_kwargs_default_empty(self): + assert ( + LeRobotConverter(META).build_task_kwargs( + episode_index=0, episode_name="episode_000000", frames=[] + ) + == {} + ) + + +class TestStaticNameSelection: + def test_preserves_declared_order(self): + class Sel(LeRobotConverter): + OBSERVATION_STATE_NAMES = ("j1.pos", "j0.pos") + ACTION_NAMES = ("a.pos",) + + c = Sel(META) + assert c._state_index == [2, 0] + assert c._action_index == [0] + assert c.build_telemetry_frame(FRAME) == { + "observation.state": [30.0, 10.0], + "action": [7.0], + "frame_index": 3, + "timestamp": 0.3, + } + + def test_missing_name_raises(self): + class Sel(LeRobotConverter): + OBSERVATION_STATE_NAMES = ("does_not_exist",) + + with pytest.raises(FastLabelInvalidException): + Sel(META) + + def test_missing_names_metadata_raises(self): + class Sel(LeRobotConverter): + ACTION_NAMES = ("a.pos",) + + with pytest.raises(FastLabelInvalidException): + Sel({}) # no features/names in meta + + +class TestDynamicSelection: + def test_suffix_name_override(self): + class Pos(LeRobotConverter): + def select_observation_state_names(self, names): + return [n for n in names if n.endswith(".pos")] + + c = Pos(META) + assert c._state_index == [0, 2] + assert c.build_observation_state(FRAME) == [10.0, 30.0] + + +class TestExtraTelemetry: + def test_super_extend(self): + class WithGripper(LeRobotConverter): + def build_telemetry_frame(self, frame): + telemetry = super().build_telemetry_frame(frame) + telemetry["gripper"] = self.build_action(frame)[-1] + return telemetry + + assert WithGripper(META).build_telemetry_frame(FRAME)["gripper"] == 8.0 + + +class TestCameraSelection: + def test_keeps_only_listed_keys(self): + class Cam(LeRobotConverter): + CAMERA_KEYS = ("cam_high", "cam_left_wrist") + + def cam(name): + return Camera( + Path(f"videos/observation.images.{name}"), + f"observation.images.{name}", + f"images_{name}", + ) + + cams = [cam("cam_high"), cam("gel_left_near"), cam("cam_left_wrist")] + kept = Cam(META).select_cameras(cams) + assert [c.content_name for c in kept] == [ + "images_cam_high", + "images_cam_left_wrist", + ] + + +class TestEpisodeHooks: + def test_select_episodes_by_length(self): + class LongOnly(LeRobotConverter): + def select_episodes(self, episode_lengths): + return [i for i, n in episode_lengths.items() if n >= 3] + + episode_lengths = {0: 1, 1: 3, 2: 5} + assert LongOnly(META).select_episodes(episode_lengths) == [1, 2] + + def test_build_episode_name_override_uses_meta(self): + class Named(LeRobotConverter): + def build_episode_name(self, episode_index): + return f"{self.meta['robot_type']}_{episode_index:03d}" + + assert Named(META).build_episode_name(2) == "bi_widowxai_002" + + +class TestTaskKwargs: + def test_keyword_only(self): + class K(LeRobotConverter): + def build_task_kwargs(self, *, episode_index, episode_name, frames): + return {"tags": [self.meta["robot_type"], episode_name]} + + assert K(META).build_task_kwargs( + episode_index=1, episode_name="episode_000001", frames=[] + ) == {"tags": ["bi_widowxai", "episode_000001"]} + + def test_positional_call_rejected(self): + with pytest.raises(TypeError): + LeRobotConverter(META).build_task_kwargs(0, "episode_000000", []) + + +class TestLoadInfo: + def test_missing_returns_empty(self, tmp_path): + assert load_info(tmp_path) == {} + + def test_parses_info_json(self, tmp_path): + meta_dir = tmp_path / "meta" + meta_dir.mkdir() + (meta_dir / "info.json").write_text(json.dumps(META)) + assert load_info(tmp_path) == META diff --git a/tests/test_lerobot_v3_parquet.py b/tests/test_lerobot_v3_parquet.py index d8e9a7d..e52c7ed 100644 --- a/tests/test_lerobot_v3_parquet.py +++ b/tests/test_lerobot_v3_parquet.py @@ -1,10 +1,12 @@ """Tests for v3 pandas/pyarrow code paths. -Covers _build_episode_map, get_episode_indices, _convert_episode_frames, and +Covers _build_episode_map, get_episode_indices, get_episode_raw_frames, and check_dependencies so that pandas/pyarrow major-version bumps surface breakage in CI. """ +import json + import pytest pd = pytest.importorskip("pandas") @@ -52,6 +54,19 @@ def v3_dataset(tmp_path): ] _write_parquet(chunk1 / "file-000.parquet", rows) + meta_dir = tmp_path / "meta" + meta_dir.mkdir() + (meta_dir / "info.json").write_text( + json.dumps( + { + "features": { + "observation.state": {"names": ["s0", "s1"]}, + "action": {"names": ["a0", "a1"]}, + } + } + ) + ) + return tmp_path @@ -83,35 +98,55 @@ def test_get_episode_indices_sorted(self, v3_dataset): assert v3.get_episode_indices(v3_dataset) == [0, 1, 2] -class TestConvertEpisodeFrames: - def test_extracts_frame_dicts(self, v3_dataset): - frames = v3._convert_episode_frames( +class TestGetEpisodeRawFrames: + def test_extracts_all_columns_as_native(self, v3_dataset): + frames = v3.get_episode_raw_frames( v3_dataset, episode_index=1, chunk="chunk-000", file_stem="file-000" ) assert len(frames) == 3 for i, frame in enumerate(frames): + assert frame["episode_index"] == 1 assert frame["frame_index"] == i assert frame["timestamp"] == pytest.approx(i * 0.1) assert frame["action"] == [1.0, 2.0] assert isinstance(frame["observation.state"], list) - def test_missing_required_columns_returns_empty(self, tmp_path): - chunk = tmp_path / "data" / "chunk-000" - chunk.mkdir(parents=True) - _write_parquet( - chunk / "file-000.parquet", - [{"episode_index": 0, "frame_index": 0}], - ) - - assert ( - v3._convert_episode_frames( - tmp_path, episode_index=0, chunk="chunk-000", file_stem="file-000" - ) - == [] - ) - class TestCheckDependencies: def test_returns_when_available(self): common.check_dependencies() + + +class TestCreateEpisodeZip: + def test_writes_zip_into_output_dir_without_leaking_staging( + self, v3_dataset, tmp_path + ): + import zipfile + from pathlib import Path + + from fastlabel.lerobot import LeRobotConverter, create_episode_zip + + out = tmp_path / "out" + 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"] + ) + zip_path = create_episode_zip( + v3_dataset, + episode_index=1, + episode_name="episode_000001", + converter=LeRobotConverter(common.load_info(v3_dataset)), + output_dir=out, + episode_map=episode_map, + raw_frames=raw_frames, + ) + + # ZIP is written under output_dir; the staging dir is not leaked there. + assert Path(zip_path).parent == out + assert [p.name for p in out.iterdir()] == ["episode_000001.zip"] + + # No videos in this fixture, so the ZIP holds only the telemetry JSON. + with zipfile.ZipFile(zip_path) as zf: + assert zf.namelist() == ["episode_000001.json"]