From 7db3a4b453d5b51b87c29fae4937d719bc188f01 Mon Sep 17 00:00:00 2001 From: Qizhi Chen Date: Tue, 23 Jun 2026 15:28:39 -0700 Subject: [PATCH] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20Update=20dataset=20version?= =?UTF-8?q?=20conversion=20script=20(#110)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Codex --- .../v30_to_v21/convert_dataset_v30_to_v21.py | 923 +++++++++--------- 1 file changed, 472 insertions(+), 451 deletions(-) diff --git a/ds_version_convert/v30_to_v21/convert_dataset_v30_to_v21.py b/ds_version_convert/v30_to_v21/convert_dataset_v30_to_v21.py index c0909af..16d999b 100644 --- a/ds_version_convert/v30_to_v21/convert_dataset_v30_to_v21.py +++ b/ds_version_convert/v30_to_v21/convert_dataset_v30_to_v21.py @@ -1,540 +1,561 @@ -"""Utilities to convert a LeRobot dataset from codebase version v3.0 back to v2.1. +#!/usr/bin/env python -The script mirrors :mod:`lerobot.datasets.v21.convert_dataset_v21_to_v30` but applies the reverse -transformations so an existing dataset created with the new consolidated file -layout can be ported back to the legacy per-episode structure. +# Copyright 2025 The HuggingFace Inc. team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. -Usage examples --------------- +""" +This script will help you convert any LeRobot dataset already pushed to the hub from codebase version 2.1 to +3.0. It will: -Convert a dataset that already exists locally:: +- Generate per-episodes stats and writes them in `episodes_stats.jsonl` +- Check consistency between these new stats and the old ones. +- Remove the deprecated `stats.json`. +- Update codebase_version in `info.json`. +- Push this new version to the hub on the 'main' branch and tags it with "v3.0". - python convert_dataset_v30_to_v21.py \ - --repo-id=lerobot/pusht \ - --root=/path/to/dataset +Usage: + +```bash +python src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py \ + --repo-id=lerobot/pusht +``` """ -from __future__ import annotations - import argparse import logging -import math import shutil -import subprocess -from collections import defaultdict from pathlib import Path -from typing import Any, Iterable +from typing import Any import jsonlines -import numpy as np -import pyarrow.parquet as pq +import pandas as pd +import pyarrow as pa import tqdm -from datasets import Dataset -from huggingface_hub import snapshot_download -from lerobot.datasets.io_utils import ( - load_info, - load_tasks, - write_info, -) +from datasets import Dataset, Features, Image +from huggingface_hub import HfApi, snapshot_download +from lerobot.datasets.compute_stats import aggregate_stats +from lerobot.datasets.lerobot_dataset import CODEBASE_VERSION, LeRobotDataset from lerobot.datasets.utils import ( DEFAULT_CHUNK_SIZE, + DEFAULT_DATA_FILE_SIZE_IN_MB, DEFAULT_DATA_PATH, + DEFAULT_VIDEO_FILE_SIZE_IN_MB, DEFAULT_VIDEO_PATH, - EPISODES_DIR, LEGACY_EPISODES_PATH, LEGACY_EPISODES_STATS_PATH, LEGACY_TASKS_PATH, - serialize_dict, - unflatten_dict, + cast_stats_to_numpy, + flatten_dict, + get_file_size_in_mb, + get_parquet_file_size_in_mb, + get_parquet_num_frames, + load_info, + update_chunk_file_indices, + write_episodes, + write_info, + write_stats, + write_tasks, ) +from lerobot.datasets.video_utils import concatenate_video_files, get_video_duration_in_s from lerobot.utils.constants import HF_LEROBOT_HOME from lerobot.utils.utils import init_logging +from requests import HTTPError V21 = "v2.1" V30 = "v3.0" -LEGACY_DATA_PATH_TEMPLATE = ( - "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet" -) -LEGACY_VIDEO_PATH_TEMPLATE = ( - "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4" -) -MIN_VIDEO_DURATION = 1e-6 -LEGACY_STATS_KEYS = ("mean", "std", "min", "max", "count") +""" +------------------------- +OLD +data/chunk-000/episode_000000.parquet + +NEW +data/chunk-000/file_000.parquet +------------------------- +OLD +videos/chunk-000/CAMERA/episode_000000.mp4 + +NEW +videos/CAMERA/chunk-000/file_000.mp4 +------------------------- +OLD +episodes.jsonl +{"episode_index": 1, "tasks": ["Put the blue block in the green bowl"], "length": 266} + +NEW +meta/episodes/chunk-000/episodes_000.parquet +episode_index | video_chunk_index | video_file_index | data_chunk_index | data_file_index | tasks | length +------------------------- +OLD +tasks.jsonl +{"task_index": 1, "task": "Put the blue block in the green bowl"} + +NEW +meta/tasks/chunk-000/file_000.parquet +task_index | task +------------------------- +OLD +episodes_stats.jsonl + +NEW +meta/episodes_stats/chunk-000/file_000.parquet +episode_index | mean | std | min | max +------------------------- +UPDATE +meta/info.json +------------------------- +""" -def _to_serializable(value: Any) -> Any: - """Convert numpy/pyarrow values into standard Python types for JSON dumps.""" - - if isinstance(value, np.ndarray): - return value.tolist() - if isinstance(value, np.generic): - return value.item() - if isinstance(value, (list, tuple)): - return [_to_serializable(item) for item in value] - if isinstance(value, dict): - return {key: _to_serializable(val) for key, val in value.items()} - return value +def load_jsonlines(fpath: Path) -> list[Any]: + with jsonlines.open(fpath, "r") as reader: + return list(reader) -def load_episode_records(root: Path) -> list[dict[str, Any]]: - """Load the consolidated metadata rows stored in ``meta/episodes``.""" - - episodes_dir = root / EPISODES_DIR - pq_paths = sorted(episodes_dir.glob("chunk-*/file-*.parquet")) - if not pq_paths: - raise FileNotFoundError(f"No episode parquet files found in {episodes_dir}.") - - records: list[dict[str, Any]] = [] - for pq_path in pq_paths: - table = pq.read_table(pq_path) - records.extend(table.to_pylist()) - - records.sort(key=lambda rec: int(rec["episode_index"])) - return records +def legacy_load_episodes(local_dir: Path) -> dict: + episodes = load_jsonlines(local_dir / LEGACY_EPISODES_PATH) + return {item["episode_index"]: item for item in sorted(episodes, key=lambda x: x["episode_index"])} -def convert_tasks(root: Path, new_root: Path) -> None: - logging.info("Converting tasks parquet to legacy JSONL") - tasks = load_tasks(root) - tasks = tasks.sort_values("task_index") - - out_path = new_root / LEGACY_TASKS_PATH - out_path.parent.mkdir(parents=True, exist_ok=True) - - with jsonlines.open(out_path, mode="w") as writer: - for task, row in tasks.iterrows(): - writer.write( - { - "task_index": int(row["task_index"]), - "task": _to_serializable(task), - } - ) +def legacy_load_episodes_stats(local_dir: Path) -> dict: + episodes_stats = load_jsonlines(local_dir / LEGACY_EPISODES_STATS_PATH) + return { + item["episode_index"]: cast_stats_to_numpy(item["stats"]) + for item in sorted(episodes_stats, key=lambda x: x["episode_index"]) + } -def convert_info( - root: Path, - new_root: Path, - episode_records: list[dict[str, Any]], - video_keys: list[str], -) -> None: - info = load_info(root) - logging.info("Converting info.json metadata to v2.1 schema") +def legacy_load_tasks(local_dir: Path) -> tuple[dict, dict]: + tasks = load_jsonlines(local_dir / LEGACY_TASKS_PATH) + tasks = {item["task_index"]: item["task"] for item in sorted(tasks, key=lambda x: x["task_index"])} + task_to_task_index = {task: task_index for task_index, task in tasks.items()} + return tasks, task_to_task_index - total_episodes = info.get("total_episodes") or len(episode_records) - chunks_size = info.get("chunks_size", DEFAULT_CHUNK_SIZE) - info["codebase_version"] = V21 +def validate_local_dataset_version(local_path: Path) -> None: + """Validate that the local dataset has the expected v2.1 version.""" + info = load_info(local_path) + dataset_version = info.get("codebase_version", "unknown") + if dataset_version != V21: + raise ValueError( + f"Local dataset has codebase version '{dataset_version}', expected '{V21}'. " + f"This script is specifically for converting v2.1 datasets to v3.0." + ) - # Restore legacy layout templates. - info["data_path"] = LEGACY_DATA_PATH_TEMPLATE - if info.get("video_path") is not None and len(video_keys) > 0: - info["video_path"] = LEGACY_VIDEO_PATH_TEMPLATE + +def convert_tasks(root, new_root): + logging.info(f"Converting tasks from {root} to {new_root}") + tasks, _ = legacy_load_tasks(root) + task_indices = tasks.keys() + task_strings = tasks.values() + df_tasks = pd.DataFrame({"task_index": task_indices}, index=task_strings) + write_tasks(df_tasks, new_root) + + +def concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys): + # TODO(rcadene): to save RAM use Dataset.from_parquet(file) and concatenate_datasets + dataframes = [pd.read_parquet(file) for file in paths_to_cat] + # Concatenate all DataFrames along rows + concatenated_df = pd.concat(dataframes, ignore_index=True) + + path = new_root / DEFAULT_DATA_PATH.format(chunk_index=chunk_idx, file_index=file_idx) + path.parent.mkdir(parents=True, exist_ok=True) + + if len(image_keys) > 0: + schema = pa.Schema.from_pandas(concatenated_df) + features = Features.from_arrow_schema(schema) + for key in image_keys: + features[key] = Image() + schema = features.arrow_schema else: - info["video_path"] = None + schema = None - # Remove v3-specific sizing hints which do not exist in v2.1. - info.pop("data_files_size_in_mb", None) - info.pop("video_files_size_in_mb", None) - - # Restore per-feature metadata: camera entries already contain their own fps. - for key, ft in info["features"].items(): - if ft.get("dtype") != "video": - ft.pop("fps", None) - - info["total_chunks"] = ( - math.ceil(total_episodes / chunks_size) if total_episodes > 0 else 0 - ) - info["total_videos"] = total_episodes * len(video_keys) - - write_info(info, new_root) + concatenated_df.to_parquet(path, index=False, schema=schema) -def _group_episodes_by_data_file( - episode_records: Iterable[dict[str, Any]], -) -> dict[tuple[int, int], list[dict[str, Any]]]: - grouped: dict[tuple[int, int], list[dict[str, Any]]] = defaultdict(list) - for record in episode_records: - key = ( - int(record["data/chunk_index"]), - int(record["data/file_index"]), - ) - grouped[key].append(record) - return grouped +def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int): + data_dir = root / "data" + ep_paths = sorted(data_dir.glob("*/*.parquet")) + image_keys = get_image_keys(root) -def convert_data( - root: Path, new_root: Path, episode_records: list[dict[str, Any]] -) -> None: - logging.info("Converting consolidated parquet files back to per-episode files") - grouped = _group_episodes_by_data_file(episode_records) + ep_idx = 0 + chunk_idx = 0 + file_idx = 0 + size_in_mb = 0 + num_frames = 0 + paths_to_cat = [] + episodes_metadata = [] - for (chunk_idx, file_idx), records in tqdm.tqdm( - grouped.items(), desc="convert data files" - ): - source_path = root / DEFAULT_DATA_PATH.format( - chunk_index=chunk_idx, file_index=file_idx - ) - if not source_path.exists(): - raise FileNotFoundError( - f"Expected source parquet file not found: {source_path}" - ) + logging.info(f"Converting data files from {len(ep_paths)} episodes") - table = pq.read_table(source_path) - records = sorted(records, key=lambda rec: int(rec["dataset_from_index"])) - file_offset = int(records[0]["dataset_from_index"]) + for ep_path in tqdm.tqdm(ep_paths, desc="convert data files"): + ep_size_in_mb = get_parquet_file_size_in_mb(ep_path) + ep_num_frames = get_parquet_num_frames(ep_path) + ep_metadata = { + "episode_index": ep_idx, + "data/chunk_index": chunk_idx, + "data/file_index": file_idx, + "dataset_from_index": num_frames, + "dataset_to_index": num_frames + ep_num_frames, + } + size_in_mb += ep_size_in_mb + num_frames += ep_num_frames + episodes_metadata.append(ep_metadata) + ep_idx += 1 - for record in records: - episode_index = int(record["episode_index"]) - start = int(record["dataset_from_index"]) - file_offset - stop = int(record["dataset_to_index"]) - file_offset - length = stop - start - - if length <= 0: - raise ValueError( - "Invalid episode length computed during data conversion: " - f"episode_index={episode_index}, length={length}" - ) - - episode_table = table.slice(start, length) - - dest_chunk = episode_index // DEFAULT_CHUNK_SIZE - dest_path = new_root / LEGACY_DATA_PATH_TEMPLATE.format( - episode_chunk=dest_chunk, - episode_index=episode_index, - ) - dest_path.parent.mkdir(parents=True, exist_ok=True) - Dataset(episode_table).to_parquet(dest_path) - - -def _group_episodes_by_video_file( - episode_records: Iterable[dict[str, Any]], - video_key: str, -) -> dict[tuple[int, int], list[dict[str, Any]]]: - grouped: dict[tuple[int, int], list[dict[str, Any]]] = defaultdict(list) - chunk_column = f"videos/{video_key}/chunk_index" - file_column = f"videos/{video_key}/file_index" - - for record in episode_records: - if chunk_column not in record or file_column not in record: + if size_in_mb < data_file_size_in_mb: + paths_to_cat.append(ep_path) continue - chunk_idx = record.get(chunk_column) - file_idx = record.get(file_column) - if chunk_idx is None or file_idx is None: - continue - grouped[(int(chunk_idx), int(file_idx))].append(record) - return grouped + + if paths_to_cat: + concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys) + + # Reset for the next file + size_in_mb = ep_size_in_mb + paths_to_cat = [ep_path] + + chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE) + + # Write remaining data if any + if paths_to_cat: + concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys) + + return episodes_metadata -def _validate_video_paths(src: Path, dst: Path) -> None: - """Validate source and destination paths to prevent security issues.""" - - # Convert to Path objects if they aren't already - src = Path(src) - dst = Path(dst) - - # Resolve paths to handle symlinks and normalize them - try: - src_resolved = src.resolve() - dst_resolved = dst.resolve() - except OSError as exc: - raise ValueError(f"Invalid path provided: {exc}") from exc - - # Check that source file exists and is a regular file - if not src_resolved.exists(): - raise FileNotFoundError(f"Source video file does not exist: {src_resolved}") - - if not src_resolved.is_file(): - raise ValueError(f"Source path is not a regular file: {src_resolved}") - - # Validate file extensions for video files - valid_video_extensions = {".mp4", ".avi", ".mov", ".mkv", ".webm", ".m4v"} - if src_resolved.suffix.lower() not in valid_video_extensions: - raise ValueError( - f"Source file does not have a valid video extension: {src_resolved}" - ) - - if dst_resolved.suffix.lower() not in valid_video_extensions: - raise ValueError( - f"Destination file does not have a valid video extension: {dst_resolved}" - ) - - # Check for path traversal attempts in the original paths - src_str = str(src) - dst_str = str(dst) - - # Ensure paths don't contain null bytes or other control characters - for path_str, name in [(src_str, "source"), (dst_str, "destination")]: - if "\0" in path_str: - raise ValueError(f"Path contains null bytes: {name} path") - if any(ord(c) < 32 and c not in ["\t", "\n", "\r"] for c in path_str): - raise ValueError(f"Path contains invalid control characters: {name} path") - - # Additional check: ensure resolved paths don't point to system directories - system_dirs = {"/etc", "/sys", "/proc", "/dev", "/boot", "/root"} - for resolved_path, name in [ - (src_resolved, "source"), - (dst_resolved, "destination"), - ]: - path_str = str(resolved_path) - for sys_dir in system_dirs: - if path_str.startswith(sys_dir + "/") or path_str == sys_dir: - raise ValueError( - f"Path points to system directory: {name} path {resolved_path}" - ) - - # Ensure the destination directory can be created safely - try: - dst_parent = dst_resolved.parent - if not dst_parent.exists(): - # Check if we can create the parent directory structure - dst_parent.resolve() - except OSError as exc: - raise ValueError(f"Cannot create destination directory: {exc}") from exc +def get_video_keys(root): + info = load_info(root) + features = info["features"] + video_keys = [key for key, ft in features.items() if ft["dtype"] == "video"] + return video_keys -def _extract_video_segment( - src: Path, - dst: Path, - start: float, - end: float, -) -> None: - # Validate paths to prevent security issues - _validate_video_paths(src, dst) - - # Validate numeric parameters to prevent injection - if not (0 <= start <= 86400): # 24 hours max - raise ValueError(f"Invalid start time: {start}") - if not (0 <= end <= 86400): # 24 hours max - raise ValueError(f"Invalid end time: {end}") - if start >= end: - raise ValueError(f"Start time {start} must be less than end time {end}") - - duration = max(end - start, MIN_VIDEO_DURATION) - - # Validate duration is reasonable - if duration > 3600: # 1 hour max - raise ValueError(f"Video segment duration too long: {duration} seconds") - - dst.parent.mkdir(parents=True, exist_ok=True) - - # Build command with validated parameters - cmd = [ - "ffmpeg", - "-hide_banner", - "-loglevel", - "debug", - "-ss", - f"{start:.6f}", - "-i", - str(src), - "-t", - f"{duration:.6f}", - "-c", - "copy", - "-avoid_negative_ts", - "1", - "-y", - str(dst), - ] - - try: - # Use more secure subprocess call with explicit timeout - result = subprocess.run( - cmd, - check=True, - timeout=300, # 5 minute timeout - capture_output=True, - text=True, - ) - except subprocess.TimeoutExpired as exc: - raise RuntimeError( - f"ffmpeg timed out while processing video '{src}' -> '{dst}'" - ) from exc - except FileNotFoundError as exc: - raise RuntimeError( - "ffmpeg executable not found; it is required for video conversion" - ) from exc - except subprocess.CalledProcessError as exc: - error_msg = f"ffmpeg failed while splitting video '{src}' into '{dst}'" - if exc.stderr: - error_msg += f". Error: {exc.stderr.strip()}" - raise RuntimeError(error_msg) from exc +def get_image_keys(root): + info = load_info(root) + features = info["features"] + image_keys = [key for key, ft in features.items() if ft["dtype"] == "image"] + return image_keys -def convert_videos( - root: Path, - new_root: Path, - episode_records: list[dict[str, Any]], - video_keys: list[str], -) -> None: +def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int): + logging.info(f"Converting videos from {root} to {new_root}") + + video_keys = get_video_keys(root) if len(video_keys) == 0: - logging.info("No video features detected; skipping video conversion") - return + return None - logging.info("Converting concatenated MP4 files back to per-episode videos") + video_keys = sorted(video_keys) - for video_key in video_keys: - grouped = _group_episodes_by_video_file(episode_records, video_key) - if len(grouped) == 0: - logging.info("No video metadata found for key '%s'; skipping", video_key) + eps_metadata_per_cam = [] + for camera in video_keys: + eps_metadata = convert_videos_of_camera(root, new_root, camera, video_file_size_in_mb) + eps_metadata_per_cam.append(eps_metadata) + + num_eps_per_cam = [len(eps_cam_map) for eps_cam_map in eps_metadata_per_cam] + if len(set(num_eps_per_cam)) != 1: + raise ValueError(f"All cams dont have same number of episodes ({num_eps_per_cam}).") + + episods_metadata = [] + num_cameras = len(video_keys) + num_episodes = num_eps_per_cam[0] + for ep_idx in tqdm.tqdm(range(num_episodes), desc="convert videos"): + # Sanity check + ep_ids = [eps_metadata_per_cam[cam_idx][ep_idx]["episode_index"] for cam_idx in range(num_cameras)] + ep_ids += [ep_idx] + if len(set(ep_ids)) != 1: + raise ValueError(f"All episode indices need to match ({ep_ids}).") + + ep_dict = {} + for cam_idx in range(num_cameras): + ep_dict.update(eps_metadata_per_cam[cam_idx][ep_idx]) + episods_metadata.append(ep_dict) + + return episods_metadata + + +def convert_videos_of_camera(root: Path, new_root: Path, video_key: str, video_file_size_in_mb: int): + # Access old paths to mp4 + videos_dir = root / "videos" + ep_paths = sorted(videos_dir.glob(f"*/{video_key}/*.mp4")) + + ep_idx = 0 + chunk_idx = 0 + file_idx = 0 + size_in_mb = 0 + duration_in_s = 0.0 + paths_to_cat = [] + episodes_metadata = [] + + for ep_path in tqdm.tqdm(ep_paths, desc=f"convert videos of {video_key}"): + ep_size_in_mb = get_file_size_in_mb(ep_path) + ep_duration_in_s = get_video_duration_in_s(ep_path) + + # Check if adding this episode would exceed the limit + if size_in_mb + ep_size_in_mb >= video_file_size_in_mb and len(paths_to_cat) > 0: + # Size limit would be exceeded, save current accumulation WITHOUT this episode + concatenate_video_files( + paths_to_cat, + new_root + / DEFAULT_VIDEO_PATH.format(video_key=video_key, chunk_index=chunk_idx, file_index=file_idx), + ) + + # Update episodes metadata for the file we just saved + for i, _ in enumerate(paths_to_cat): + past_ep_idx = ep_idx - len(paths_to_cat) + i + episodes_metadata[past_ep_idx][f"videos/{video_key}/chunk_index"] = chunk_idx + episodes_metadata[past_ep_idx][f"videos/{video_key}/file_index"] = file_idx + + # Move to next file and start fresh with current episode + chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE) + size_in_mb = 0 + duration_in_s = 0.0 + paths_to_cat = [] + + # Add current episode metadata + ep_metadata = { + "episode_index": ep_idx, + f"videos/{video_key}/chunk_index": chunk_idx, # Will be updated when file is saved + f"videos/{video_key}/file_index": file_idx, # Will be updated when file is saved + f"videos/{video_key}/from_timestamp": duration_in_s, + f"videos/{video_key}/to_timestamp": duration_in_s + ep_duration_in_s, + } + episodes_metadata.append(ep_metadata) + + # Add current episode to accumulation + paths_to_cat.append(ep_path) + size_in_mb += ep_size_in_mb + duration_in_s += ep_duration_in_s + ep_idx += 1 + + # Write remaining videos if any + if paths_to_cat: + concatenate_video_files( + paths_to_cat, + new_root + / DEFAULT_VIDEO_PATH.format(video_key=video_key, chunk_index=chunk_idx, file_index=file_idx), + ) + + # Update episodes metadata for the final file + for i, _ in enumerate(paths_to_cat): + past_ep_idx = ep_idx - len(paths_to_cat) + i + episodes_metadata[past_ep_idx][f"videos/{video_key}/chunk_index"] = chunk_idx + episodes_metadata[past_ep_idx][f"videos/{video_key}/file_index"] = file_idx + + return episodes_metadata + + +def generate_episode_metadata_dict( + episodes_legacy_metadata, episodes_metadata, episodes_stats, episodes_videos=None +): + num_episodes = len(episodes_metadata) + episodes_legacy_metadata_vals = list(episodes_legacy_metadata.values()) + episodes_stats_vals = list(episodes_stats.values()) + episodes_stats_keys = list(episodes_stats.keys()) + + for i in range(num_episodes): + ep_legacy_metadata = episodes_legacy_metadata_vals[i] + ep_metadata = episodes_metadata[i] + ep_stats = episodes_stats_vals[i] + + ep_ids_set = { + ep_legacy_metadata["episode_index"], + ep_metadata["episode_index"], + episodes_stats_keys[i], + } + + if episodes_videos is None: + ep_video = {} + else: + ep_video = episodes_videos[i] + ep_ids_set.add(ep_video["episode_index"]) + + if len(ep_ids_set) != 1: + raise ValueError(f"Number of episodes is not the same ({ep_ids_set}).") + + ep_dict = {**ep_metadata, **ep_video, **ep_legacy_metadata, **flatten_dict({"stats": ep_stats})} + ep_dict["meta/episodes/chunk_index"] = 0 + ep_dict["meta/episodes/file_index"] = 0 + yield ep_dict + + +def convert_episodes_metadata(root, new_root, episodes_metadata, episodes_video_metadata=None): + logging.info(f"Converting episodes metadata from {root} to {new_root}") + + episodes_legacy_metadata = legacy_load_episodes(root) + episodes_stats = legacy_load_episodes_stats(root) + + num_eps_set = {len(episodes_legacy_metadata), len(episodes_metadata)} + if episodes_video_metadata is not None: + num_eps_set.add(len(episodes_video_metadata)) + + if len(num_eps_set) != 1: + raise ValueError(f"Number of episodes is not the same ({num_eps_set}).") + + ds_episodes = Dataset.from_generator( + lambda: generate_episode_metadata_dict( + episodes_legacy_metadata, episodes_metadata, episodes_stats, episodes_video_metadata + ) + ) + write_episodes(ds_episodes, new_root) + + stats = aggregate_stats(list(episodes_stats.values())) + write_stats(stats, new_root) + + +def convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb): + info = load_info(root) + info["codebase_version"] = V30 + del info["total_chunks"] + del info["total_videos"] + info["data_files_size_in_mb"] = data_file_size_in_mb + info["video_files_size_in_mb"] = video_file_size_in_mb + info["data_path"] = DEFAULT_DATA_PATH + info["video_path"] = DEFAULT_VIDEO_PATH if info["video_path"] is not None else None + info["fps"] = int(info["fps"]) + logging.info(f"Converting info from {root} to {new_root}") + for key in info["features"]: + if info["features"][key]["dtype"] == "video": + # already has fps in video_info continue - - for (chunk_idx, file_idx), records in tqdm.tqdm( - grouped.items(), desc=f"convert videos ({video_key})" - ): - src_path = root / DEFAULT_VIDEO_PATH.format( - video_key=video_key, - chunk_index=chunk_idx, - file_index=file_idx, - ) - if not src_path.exists(): - raise FileNotFoundError(f"Expected MP4 file not found: {src_path}") - - records = sorted( - records, - key=lambda rec: float(rec[f"videos/{video_key}/from_timestamp"]), - ) - - for record in records: - episode_index = int(record["episode_index"]) - start = float(record[f"videos/{video_key}/from_timestamp"]) - end = float(record[f"videos/{video_key}/to_timestamp"]) - - dest_chunk = episode_index // DEFAULT_CHUNK_SIZE - dest_path = new_root / LEGACY_VIDEO_PATH_TEMPLATE.format( - episode_chunk=dest_chunk, - video_key=video_key, - episode_index=episode_index, - ) - - _extract_video_segment(src_path, dest_path, start=start, end=end) - - -def convert_episodes_metadata( - new_root: Path, episode_records: list[dict[str, Any]] -) -> None: - logging.info("Reconstructing legacy episodes and episodes_stats JSONL files") - - episodes_path = new_root / LEGACY_EPISODES_PATH - stats_path = new_root / LEGACY_EPISODES_STATS_PATH - episodes_path.parent.mkdir(parents=True, exist_ok=True) - - def _filter_stats(stats: dict[str, Any]) -> dict[str, Any]: - """Remove v3-only statistics keys so output matches the v2.1 schema.""" - - filtered: dict[str, Any] = {} - for feature, values in stats.items(): - if not isinstance(values, dict): - continue - keep = {k: v for k, v in values.items() if k in LEGACY_STATS_KEYS} - if keep: - filtered[feature] = keep - return filtered - - with ( - jsonlines.open(episodes_path, mode="w") as episodes_writer, - jsonlines.open(stats_path, mode="w") as stats_writer, - ): - for record in sorted( - episode_records, key=lambda rec: int(rec["episode_index"]) - ): - legacy_episode = { - key: value - for key, value in record.items() - if not key.startswith("data/") - and not key.startswith("videos/") - and not key.startswith("stats/") - and not key.startswith("meta/") - and key not in {"dataset_from_index", "dataset_to_index"} - } - - serializable_episode = { - key: _to_serializable(value) for key, value in legacy_episode.items() - } - episodes_writer.write(serializable_episode) - - stats_flat = { - key: record[key] for key in record if key.startswith("stats/") - } - stats_nested = unflatten_dict(stats_flat).get("stats", {}) - stats_serialized = serialize_dict(_filter_stats(stats_nested)) - stats_writer.write( - { - "episode_index": int(record["episode_index"]), - "stats": stats_serialized, - } - ) - - -def copy_ancillary_directories(root: Path, new_root: Path) -> None: - for subdir in ["images"]: - source = root / subdir - if source.exists(): - shutil.copytree(source, new_root / subdir, dirs_exist_ok=True) + info["features"][key]["fps"] = info["fps"] + write_info(info, new_root) def convert_dataset( repo_id: str, + branch: str | None = None, + data_file_size_in_mb: int | None = None, + video_file_size_in_mb: int | None = None, root: str | Path | None = None, -) -> None: - root = HF_LEROBOT_HOME / repo_id if root is None else Path(root) + push_to_hub: bool = True, + force_conversion: bool = False, +): + if data_file_size_in_mb is None: + data_file_size_in_mb = DEFAULT_DATA_FILE_SIZE_IN_MB + if video_file_size_in_mb is None: + video_file_size_in_mb = DEFAULT_VIDEO_FILE_SIZE_IN_MB - if not root.exists(): - snapshot_download( - repo_id, - repo_type="dataset", - revision=V30, - local_dir=root, - ) + # First check if the dataset already has a v3.0 version + if root is None and not force_conversion: + try: + print("Trying to download v3.0 version of the dataset from the hub...") + snapshot_download(repo_id, repo_type="dataset", revision=V30, local_dir=HF_LEROBOT_HOME / repo_id) + return + except Exception: + print("Dataset does not have an uploaded v3.0 version. Continuing with conversion.") - old_root = root.parent / f"{root.name}_{V30}" - new_root = root.parent / f"{root.name}_{V21}" + # Set root based on whether local dataset path is provided + use_local_dataset = False + root = HF_LEROBOT_HOME / repo_id if root is None else Path(root) / repo_id + if root.exists(): + validate_local_dataset_version(root) + use_local_dataset = True + print(f"Using local dataset at {root}") + + old_root = root.parent / f"{root.name}_old" + new_root = root.parent / f"{root.name}_v30" + + # Handle old_root cleanup if both old_root and root exist + if old_root.is_dir() and root.is_dir(): + shutil.rmtree(str(root)) + shutil.move(str(old_root), str(root)) - if old_root.is_dir(): - shutil.rmtree(old_root) if new_root.is_dir(): shutil.rmtree(new_root) - new_root.mkdir(parents=True, exist_ok=True) + if not use_local_dataset: + snapshot_download( + repo_id, + repo_type="dataset", + revision=V21, + local_dir=root, + ) - episode_records = load_episode_records(root) - video_keys = [ - key - for key, ft in load_info(root)["features"].items() - if ft.get("dtype") == "video" - ] - - convert_info(root, new_root, episode_records, video_keys) + convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb) convert_tasks(root, new_root) - convert_data(root, new_root, episode_records) - convert_videos(root, new_root, episode_records, video_keys) - convert_episodes_metadata(new_root, episode_records) - copy_ancillary_directories(root, new_root) + episodes_metadata = convert_data(root, new_root, data_file_size_in_mb) + episodes_videos_metadata = convert_videos(root, new_root, video_file_size_in_mb) + convert_episodes_metadata(root, new_root, episodes_metadata, episodes_videos_metadata) shutil.move(str(root), str(old_root)) shutil.move(str(new_root), str(root)) + if push_to_hub: + hub_api = HfApi() + try: + hub_api.delete_tag(repo_id, tag=CODEBASE_VERSION, repo_type="dataset") + except HTTPError as e: + print(f"tag={CODEBASE_VERSION} probably doesn't exist. Skipping exception ({e})") + pass + hub_api.delete_files( + delete_patterns=["data/chunk*/episode_*", "meta/*.jsonl", "videos/chunk*"], + repo_id=repo_id, + revision=branch, + repo_type="dataset", + ) + hub_api.create_tag(repo_id, tag=CODEBASE_VERSION, revision=branch, repo_type="dataset") -def parse_args() -> argparse.Namespace: + LeRobotDataset(repo_id).push_to_hub() + + +if __name__ == "__main__": + init_logging() parser = argparse.ArgumentParser() parser.add_argument( "--repo-id", type=str, required=True, - help="Repository identifier on Hugging Face (e.g. `lerobot/pusht`).", + help="Repository identifier on Hugging Face: a community or a user name `/` the name of the dataset " + "(e.g. `lerobot/pusht`, `cadene/aloha_sim_insertion_human`).", + ) + parser.add_argument( + "--branch", + type=str, + default=None, + help="Repo branch to push your dataset. Defaults to the main branch.", + ) + parser.add_argument( + "--data-file-size-in-mb", + type=int, + default=None, + help="File size in MB. Defaults to 100 for data and 500 for videos.", + ) + parser.add_argument( + "--video-file-size-in-mb", + type=int, + default=None, + help="File size in MB. Defaults to 100 for data and 500 for videos.", ) parser.add_argument( "--root", type=str, default=None, - help="Path to the local dataset root directory. If not provided, the script will use the dataset from local.", + help="Local directory to use for downloading/writing the dataset.", + ) + parser.add_argument( + "--push-to-hub", + type=lambda input: input.lower() == "true", + default=True, + help="Push the converted dataset to the hub.", + ) + parser.add_argument( + "--force-conversion", + action="store_true", + help="Force conversion even if the dataset already has a v3.0 version.", ) - return parser.parse_args() - -if __name__ == "__main__": - init_logging() - args = parse_args() - convert_dataset(**vars(args)) + args = parser.parse_args() + convert_dataset(**vars(args)) \ No newline at end of file