Skip to content

Commit fe6fc04

Browse files
committed
feat(convert): lerobot data convert video codec
1 parent 41f09fa commit fe6fc04

2 files changed

Lines changed: 6 additions & 0 deletions

File tree

python/rcs/__main__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -147,6 +147,7 @@ def lerobot_convert(
147147
success: Annotated[bool, typer.Option(help="Only include successful episodes. Example: --success")] = True,
148148
n: Annotated[int, typer.Option(help="Maximum number of episodes to convert. -1 means all. Example: --n 50")] = -1,
149149
video_encoding: Annotated[bool, typer.Option(help="Should the image data be video encoded")] = False,
150+
video_backend: Annotated[str | None, typer.Option(help="Video backend to use if image data is video encoded e.g. torchcodec")] = None,
150151
):
151152
cameras = camera_specs_to_configs(camera_specs) if camera_specs is not None else list(DEFAULT_CAMERAS)
152153
run_conversion(
@@ -164,6 +165,7 @@ def lerobot_convert(
164165
success=success,
165166
n=n,
166167
video_encoding=video_encoding,
168+
video_backend=video_backend,
167169
)
168170

169171

python/rcs/lerobot_joint_converter.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -99,6 +99,7 @@ def __init__(
9999
image_batch_size: int = DEFAULT_IMAGE_BATCH_SIZE,
100100
per_robot_arm_dim: int = DEFAULT_PER_ROBOT_ARM_DIM,
101101
video_encoding: bool = False,
102+
video_backend: str | None = None,
102103
):
103104
self.root = Path(root)
104105
self.conn = duckdb.connect()
@@ -133,6 +134,7 @@ def __init__(
133134
features=self._build_features(),
134135
image_writer_threads=10,
135136
image_writer_processes=5,
137+
video_backend=video_backend,
136138
)
137139

138140
def _build_features(self) -> dict[str, dict[str, Any]]:
@@ -416,6 +418,7 @@ def run_conversion(
416418
success: bool = True,
417419
n: int = -1,
418420
video_encoding: bool = False,
421+
video_backend: str | None = None,
419422
) -> None:
420423
robot_type_converted = RobotType(robot_type)
421424
gripper_type_converted = GripperType(gripper_type)
@@ -432,6 +435,7 @@ def run_conversion(
432435
image_batch_size=image_batch_size,
433436
per_robot_arm_dim=per_robot_arm_dim,
434437
video_encoding=video_encoding,
438+
video_backend=video_backend
435439
)
436440
converter.generate_examples(success=success, n=n)
437441

0 commit comments

Comments
 (0)