feat: import_lerobot に LeRobotConverter を追加（取り込みのカスタマイズ）

.gitignore

@@ -17,3 +17,4 @@ output/
 .env
 !tests/files/*
 .idea
+uv.lock

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.

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]

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
 

fastlabel/lerobot/__init__.py

@@ -1,62 +1,105 @@
+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":
         raise FastLabelInvalidException(
             "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
+    )