Add generic converter adapter hooks

Co-authored-by: Codex <codex@openai.com>
This commit is contained in:
Tavish
2026-06-21 15:44:14 +08:00
parent 7a8642edfc
commit 65296e75cb
3 changed files with 71 additions and 29 deletions
+33 -3
View File
@@ -1,6 +1,7 @@
from abc import ABC, abstractmethod
from collections.abc import Iterable, Sequence
from pathlib import Path
from typing import Any
from .utils import ConversionTask, FeatureSpec
@@ -26,7 +27,36 @@ class BaseAdapter(ABC):
"""Build conversion tasks from dataset-specific inputs."""
@abstractmethod
def load_subset(
self, task: ConversionTask
) -> Iterable[Sequence[dict]]:
def load_subset(self, task: ConversionTask) -> Iterable[Any]:
"""Yield LeRobot episodes for one raw input path."""
def create_dataset(self, task: ConversionTask):
"""Create the temporary LeRobot dataset for one conversion task."""
from lerobot.datasets import LeRobotDataset
return LeRobotDataset.create(
repo_id=task.local_repo_id,
root=task.output_path,
fps=self.fps,
robot_type=self.robot_type,
features=self.features,
)
def save_episode(
self,
dataset: Any,
episode_data: Sequence[dict],
task: ConversionTask,
) -> bool:
"""Save one episode to the temporary dataset.
Adapters can override this when a dataset needs extra per-episode
arguments or a non-standard writer.
"""
for frame in episode_data:
dataset.add_frame(frame)
dataset.save_episode()
return True
def get_episode_length(self, episode_data: Any) -> int:
return len(episode_data)