feat(video re-encoding): Adding utility and dataset edition tool for video re-encoding (#3611)

* feat(utility): adding video re-encode utility

* feat(edit): adding a new lerobot-edit-dataset tool to re-encode all the videos of a dataset

* chore(format): formatting code

* chore(review): fix Claude reviews

* test(reencode dataset): adding missing test for reencode dataset
This commit is contained in:
Caroline Pascal
2026-05-19 14:46:14 +02:00
committed by GitHub
parent 7ab4936b1b
commit d38eb89f71
6 changed files with 342 additions and 10 deletions

View File

@@ -178,6 +178,31 @@ Recompute stats for relative actions and push to hub:
--operation.num_workers 4 \
--push_to_hub true
Re-encode all videos in a dataset (saves to lerobot/pusht_reencoded by default):
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--operation.type reencode_videos \
--operation.camera_encoder.vcodec h264 \
--operation.camera_encoder.pix_fmt yuv420p \
--operation.camera_encoder.crf 23
Re-encode videos into a new dataset using 4 parallel processes:
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--new_repo_id lerobot/pusht_h264 \
--operation.type reencode_videos \
--operation.camera_encoder.vcodec h264 \
--operation.camera_encoder.crf 23 \
--operation.num_workers 4
Re-encode videos in-place (overwrites original dataset):
lerobot-edit-dataset \
--repo_id lerobot/pusht \
--new_repo_id lerobot/pusht \
--operation.type reencode_videos \
--operation.camera_encoder.vcodec h264 \
--operation.overwrite true
Using JSON config file:
lerobot-edit-dataset \
--config_path path/to/edit_config.json
@@ -200,6 +225,7 @@ from lerobot.datasets import (
merge_datasets,
modify_tasks,
recompute_stats,
reencode_dataset,
remove_feature,
split_dataset,
)
@@ -268,6 +294,15 @@ class RecomputeStatsConfig(OperationConfig):
overwrite: bool = False
@OperationConfig.register_subclass("reencode_videos")
@dataclass
class ReencodeVideosConfig(OperationConfig):
camera_encoder: VideoEncoderConfig = field(default_factory=camera_encoder_defaults)
num_workers: int = 0
encoder_threads: int | None = None
overwrite: bool = False
@OperationConfig.register_subclass("info")
@dataclass
class InfoConfig(OperationConfig):
@@ -634,6 +669,58 @@ def handle_recompute_stats(cfg: EditDatasetConfig) -> None:
dataset.push_to_hub()
def handle_reencode_videos(cfg: EditDatasetConfig) -> None:
if not isinstance(cfg.operation, ReencodeVideosConfig):
raise ValueError("Operation config must be ReencodeVideosConfig")
output_repo_id, input_root, output_root = _resolve_io_paths(
cfg.repo_id,
cfg.new_repo_id,
cfg.root,
cfg.new_root,
default_new_repo_id=f"{cfg.repo_id}_reencoded",
)
in_place = output_root == input_root
if in_place and not cfg.operation.overwrite:
raise ValueError(
f"reencode_videos would overwrite the dataset in-place at {input_root}. "
"Pass --operation.overwrite true to allow in-place modification, "
"or use --new_repo_id / --new_root to write to a different location. "
f"Default output repo_id when neither is set: '{cfg.repo_id}_reencoded'."
)
if in_place:
logging.warning(
f"Overwriting dataset videos in-place at {input_root}. The original videos will be lost."
)
dataset = LeRobotDataset(cfg.repo_id, root=input_root)
else:
logging.info(f"Copying dataset from {input_root} to {output_root}")
if output_root.exists():
backup_path = output_root.with_name(output_root.name + "_old")
logging.warning(f"Output directory {output_root} already exists. Moving to {backup_path}")
if backup_path.exists():
shutil.rmtree(backup_path)
shutil.move(output_root, backup_path)
shutil.copytree(input_root, output_root)
dataset = LeRobotDataset(output_repo_id, root=output_root)
logging.info(f"Re-encoding videos in {output_repo_id} with {cfg.operation.camera_encoder}")
reencode_dataset(
dataset,
camera_encoder=cfg.operation.camera_encoder,
encoder_threads=cfg.operation.encoder_threads,
num_workers=cfg.operation.num_workers,
)
logging.info(f"All videos re-encoded at {dataset.root}")
if cfg.push_to_hub:
logging.info(f"Pushing to hub as {output_repo_id}...")
dataset.push_to_hub()
def _get_dataset_size(repo_path):
import os
@@ -707,6 +794,8 @@ def edit_dataset(cfg: EditDatasetConfig) -> None:
handle_convert_image_to_video(cfg)
elif operation_type == "recompute_stats":
handle_recompute_stats(cfg)
elif operation_type == "reencode_videos":
handle_reencode_videos(cfg)
elif operation_type == "info":
handle_info(cfg)
else: