diff --git a/ds_version_convert/v21_to_v30/convert_dataset_v21_to_v30.py b/ds_version_convert/v21_to_v30/convert_dataset_v21_to_v30.py index aebd596..dea8f5f 100644 --- a/ds_version_convert/v21_to_v30/convert_dataset_v21_to_v30.py +++ b/ds_version_convert/v21_to_v30/convert_dataset_v21_to_v30.py @@ -26,11 +26,23 @@ This script will help you convert any LeRobot dataset already pushed to the hub Usage: +Convert a dataset from the hub: ```bash -python src/lerobot/datasets/v30/convert_dataset_v21_to_v30.py \ +python src/lerobot/scripts/convert_dataset_v21_to_v30.py \ --repo-id=lerobot/pusht ``` +Convert a local dataset (works in place): +```bash +python src/lerobot/scripts/convert_dataset_v21_to_v30.py \ + --repo-id=lerobot/pusht \ + --root=/path/to/local/dataset/directory \ + --push-to-hub=false + +N.B. Path semantics (v2): --root is the exact dataset folder containing +meta/, data/, videos/. When omitted, defaults to $HF_LEROBOT_HOME/{repo_id}. +``` + """ import argparse @@ -39,39 +51,47 @@ import shutil from pathlib import Path from typing import Any +from lerobot.utils.import_utils import require_package + +require_package("jsonlines", extra="dataset") + import jsonlines import pandas as pd import pyarrow as pa import tqdm 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 requests import HTTPError + +from lerobot.datasets import CODEBASE_VERSION, LeRobotDataset, aggregate_stats +from lerobot.datasets.io_utils import ( + cast_stats_to_numpy, + get_file_size_in_mb, + get_parquet_file_size_in_mb, + get_parquet_num_frames, + load_info, + load_json, + write_episodes, + write_info, + write_stats, + write_tasks, +) 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, + INFO_PATH, LEGACY_EPISODES_PATH, LEGACY_EPISODES_STATS_PATH, LEGACY_TASKS_PATH, - cast_stats_to_numpy, - flatten_dict, - get_file_size_in_mb, - get_parquet_file_size_in_mb, - get_parquet_num_frames, - load_info, + DatasetInfo, 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 +from lerobot.utils.utils import flatten_dict, init_logging V21 = "v2.1" V30 = "v3.0" @@ -95,7 +115,7 @@ episodes.jsonl {"episode_index": 1, "tasks": ["Put the blue block in the green bowl"], "length": 266} NEW -meta/episodes/chunk-000/episodes_000.parquet +meta/episodes/chunk-000/file_000.parquet episode_index | video_chunk_index | video_file_index | data_chunk_index | data_file_index | tasks | length ------------------------- OLD @@ -103,15 +123,16 @@ tasks.jsonl {"task_index": 1, "task": "Put the blue block in the green bowl"} NEW -meta/tasks/chunk-000/file_000.parquet +meta/tasks.parquet task_index | task ------------------------- OLD episodes_stats.jsonl +{"episode_index": 1, "stats": {"feature_name": {"min": ..., "max": ..., "mean": ..., "std": ..., "count": ...}}} NEW -meta/episodes_stats/chunk-000/file_000.parquet -episode_index | mean | std | min | max +meta/episodes/chunk-000/file_000.parquet +episode_index | feature_name/min | feature_name/max | feature_name/mean | feature_name/std | feature_name/count ------------------------- UPDATE meta/info.json @@ -147,7 +168,7 @@ def legacy_load_tasks(local_dir: Path) -> tuple[dict, dict]: 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") + dataset_version = info.codebase_version or "unknown" if dataset_version != V21: raise ValueError( f"Local dataset has codebase version '{dataset_version}', expected '{V21}'. " @@ -160,7 +181,7 @@ def convert_tasks(root, 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) + df_tasks = pd.DataFrame({"task_index": task_indices}, index=pd.Index(task_strings, name="task")) write_tasks(df_tasks, new_root) @@ -191,7 +212,6 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int): image_keys = get_image_keys(root) - ep_idx = 0 chunk_idx = 0 file_idx = 0 size_in_mb = 0 @@ -201,9 +221,23 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int): logging.info(f"Converting data files from {len(ep_paths)} episodes") - for ep_path in tqdm.tqdm(ep_paths, desc="convert data files"): + for ep_idx, ep_path in enumerate(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) + + # Check if we need to start a new file BEFORE creating metadata + if size_in_mb + ep_size_in_mb >= data_file_size_in_mb and len(paths_to_cat) > 0: + # Write the accumulated data files + concat_data_files(paths_to_cat, new_root, chunk_idx, file_idx, image_keys) + + # Move to next file + chunk_idx, file_idx = update_chunk_file_indices(chunk_idx, file_idx, DEFAULT_CHUNK_SIZE) + + # Reset for the next file + size_in_mb = 0 + paths_to_cat = [] + + # Now create metadata with correct chunk/file indices ep_metadata = { "episode_index": ep_idx, "data/chunk_index": chunk_idx, @@ -214,20 +248,7 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int): size_in_mb += ep_size_in_mb num_frames += ep_num_frames episodes_metadata.append(ep_metadata) - ep_idx += 1 - - if size_in_mb < data_file_size_in_mb: - paths_to_cat.append(ep_path) - continue - - 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) + paths_to_cat.append(ep_path) # Write remaining data if any if paths_to_cat: @@ -238,14 +259,14 @@ def convert_data(root: Path, new_root: Path, data_file_size_in_mb: int): def get_video_keys(root): info = load_info(root) - features = info["features"] + features = info.features video_keys = [key for key, ft in features.items() if ft["dtype"] == "video"] return video_keys def get_image_keys(root): info = load_info(root) - features = info["features"] + features = info.features image_keys = [key for key, ft in features.items() if ft["dtype"] == "image"] return image_keys @@ -268,7 +289,7 @@ def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int): 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 = [] + episodes_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"): @@ -281,9 +302,9 @@ def convert_videos(root: Path, new_root: Path, video_file_size_in_mb: int): 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) + episodes_metadata.append(ep_dict) - return episods_metadata + return episodes_metadata def convert_videos_of_camera(root: Path, new_root: Path, video_key: str, video_file_size_in_mb: int): @@ -416,7 +437,8 @@ def convert_episodes_metadata(root, new_root, episodes_metadata, episodes_video_ def convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb): - info = load_info(root) + # Load as raw dict to remove legacy v2.1 fields before constructing DatasetInfo. + info = load_json(root / INFO_PATH) info["codebase_version"] = V30 del info["total_chunks"] del info["total_videos"] @@ -431,7 +453,9 @@ def convert_info(root, new_root, data_file_size_in_mb, video_file_size_in_mb): # already has fps in video_info continue info["features"][key]["fps"] = info["fps"] - write_info(info, new_root) + # Convert raw dict to typed DatasetInfo before writing + dataset_info = DatasetInfo.from_dict(info) + write_info(dataset_info, new_root) def convert_dataset( @@ -459,7 +483,7 @@ def convert_dataset( # 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 + root = HF_LEROBOT_HOME / repo_id if root is None else Path(root) if root.exists(): validate_local_dataset_version(root) use_local_dataset = True @@ -519,7 +543,7 @@ if __name__ == "__main__": type=str, required=True, 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`).", + "(e.g. `lerobot/pusht`, `/aloha_sim_insertion_human`).", ) parser.add_argument( "--branch", @@ -543,7 +567,7 @@ if __name__ == "__main__": "--root", type=str, default=None, - help="Local directory to use for downloading/writing the dataset.", + help="Local directory to use for downloading/writing the dataset. Defaults to $HF_LEROBOT_HOME/repo_id.", ) parser.add_argument( "--push-to-hub", @@ -558,4 +582,4 @@ if __name__ == "__main__": ) args = parser.parse_args() - convert_dataset(**vars(args)) + convert_dataset(**vars(args)) \ No newline at end of file