| |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
| import shutil |
|
|
| import numpy as np |
| import pandas as pd |
| from tqdm import tqdm |
|
|
|
|
| def _load_json(path: Path) -> dict: |
| with path.open() as f: |
| return json.load(f) |
|
|
|
|
| def _write_json(path: Path, obj: dict) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| with path.open("w") as f: |
| json.dump(obj, f, indent=4) |
| f.write("\n") |
|
|
|
|
| def _read_jsonl(path: Path) -> list[dict]: |
| rows = [] |
| with path.open() as f: |
| for raw_line in f: |
| stripped = raw_line.strip() |
| if stripped: |
| rows.append(json.loads(stripped)) |
| return rows |
|
|
|
|
| def _write_jsonl(path: Path, rows: list[dict]) -> None: |
| path.parent.mkdir(parents=True, exist_ok=True) |
| with path.open("w") as f: |
| for row in rows: |
| f.write(json.dumps(row) + "\n") |
|
|
|
|
| def _safe_name(path: Path) -> str: |
| return path.name.replace("/", "_") |
|
|
|
|
| def _numeric_stats(values: np.ndarray) -> dict: |
| arr = np.asarray(values) |
| flat = arr.reshape(-1, 1) if arr.ndim == 1 else arr.reshape(arr.shape[0], -1) |
| return { |
| "min": flat.min(axis=0).tolist(), |
| "max": flat.max(axis=0).tolist(), |
| "mean": flat.mean(axis=0).tolist(), |
| "std": flat.std(axis=0).tolist(), |
| "count": [int(arr.shape[0])], |
| } |
|
|
|
|
| def _episode_stats(df: pd.DataFrame, episode_index: int) -> dict: |
| stats = {} |
| for col in df.columns: |
| first = df[col].iloc[0] |
| if isinstance(first, np.ndarray): |
| stats[col] = _numeric_stats(np.stack(df[col].to_numpy())) |
| elif np.issubdtype(df[col].dtype, np.number): |
| stats[col] = _numeric_stats(df[col].to_numpy()) |
| return {"episode_index": int(episode_index), "stats": stats} |
|
|
|
|
| def _video_keys(info: dict) -> list[str]: |
| return [name for name, feature in info.get("features", {}).items() if feature.get("dtype") == "video"] |
|
|
|
|
| def _task_mapping(dataset: Path) -> dict[int, str]: |
| return {int(row["task_index"]): row["task"] for row in _read_jsonl(dataset / "meta" / "tasks.jsonl")} |
|
|
|
|
| def merge_datasets(inputs: list[Path], output: Path, *, overwrite: bool) -> None: |
| inputs = [path.resolve() for path in inputs] |
| output = output.resolve() |
| if output.exists(): |
| if not overwrite: |
| raise FileExistsError(f"{output} exists; pass --overwrite to replace it") |
| shutil.rmtree(output) |
|
|
| if len(inputs) < 2: |
| raise ValueError("At least two input datasets are required") |
|
|
| infos = [_load_json(path / "meta" / "info.json") for path in inputs] |
| chunks_size = int(infos[0].get("chunks_size", 1000)) |
| video_keys = _video_keys(infos[0]) |
| for dataset, info in zip(inputs, infos, strict=True): |
| if int(info.get("chunks_size", chunks_size)) != chunks_size: |
| raise ValueError(f"{dataset} uses a different chunks_size") |
| if _video_keys(info) != video_keys: |
| raise ValueError(f"{dataset} has different video keys") |
|
|
| output.mkdir(parents=True) |
| all_tasks: dict[str, int] = {} |
| task_rows: list[dict] = [] |
| episode_rows: list[dict] = [] |
| stats_rows: list[dict] = [] |
| speed_metrics_rows: list[dict] = [] |
|
|
| global_frame_index = 0 |
| global_episode_index = 0 |
| total_videos = 0 |
|
|
| for dataset_index, dataset in enumerate(inputs): |
| source_name = _safe_name(dataset) |
| old_tasks = _task_mapping(dataset) |
| old_to_new_task = {} |
| for old_task_index, task in sorted(old_tasks.items()): |
| if task not in all_tasks: |
| all_tasks[task] = len(all_tasks) |
| task_rows.append({"task_index": all_tasks[task], "task": task}) |
| old_to_new_task[old_task_index] = all_tasks[task] |
|
|
| info = infos[dataset_index] |
| input_chunks_size = int(info.get("chunks_size", chunks_size)) |
| for episode in tqdm(_read_jsonl(dataset / "meta" / "episodes.jsonl"), desc=source_name): |
| old_episode_index = int(episode["episode_index"]) |
| old_chunk = old_episode_index // input_chunks_size |
| new_chunk = global_episode_index // chunks_size |
| src_data = dataset / "data" / f"chunk-{old_chunk:03d}" / f"episode_{old_episode_index:06d}.parquet" |
| dst_data = output / "data" / f"chunk-{new_chunk:03d}" / f"episode_{global_episode_index:06d}.parquet" |
|
|
| frame = pd.read_parquet(src_data) |
| length = len(frame) |
| frame["episode_index"] = np.full(length, global_episode_index, dtype=np.int64) |
| frame["index"] = np.arange(global_frame_index, global_frame_index + length, dtype=np.int64) |
| frame["task_index"] = frame["task_index"].map(old_to_new_task).astype(np.int64) |
| dst_data.parent.mkdir(parents=True, exist_ok=True) |
| frame.to_parquet(dst_data, index=False) |
|
|
| for key in video_keys: |
| src_video = dataset / "videos" / f"chunk-{old_chunk:03d}" / key / f"episode_{old_episode_index:06d}.mp4" |
| dst_video = ( |
| output / "videos" / f"chunk-{new_chunk:03d}" / key / f"episode_{global_episode_index:06d}.mp4" |
| ) |
| dst_video.parent.mkdir(parents=True, exist_ok=True) |
| shutil.copy2(src_video, dst_video) |
| total_videos += 1 |
|
|
| episode_row = dict(episode) |
| episode_row.update( |
| { |
| "episode_index": global_episode_index, |
| "length": int(length), |
| "source_dataset": source_name, |
| "source_dataset_index": int(dataset_index), |
| "source_dataset_episode_index": old_episode_index, |
| } |
| ) |
| episode_rows.append(episode_row) |
| stats_rows.append(_episode_stats(frame, global_episode_index)) |
|
|
| global_frame_index += length |
| global_episode_index += 1 |
|
|
| metrics_path = dataset / "meta" / "speed_metrics.jsonl" |
| if metrics_path.exists(): |
| for row in _read_jsonl(metrics_path): |
| metric = dict(row) |
| metric["source_dataset"] = source_name |
| metric["source_dataset_index"] = int(dataset_index) |
| speed_metrics_rows.append(metric) |
|
|
| info = json.loads(json.dumps(infos[0])) |
| info["total_episodes"] = int(global_episode_index) |
| info["total_frames"] = int(global_frame_index) |
| info["total_tasks"] = len(task_rows) |
| info["total_videos"] = int(total_videos) |
| info["total_chunks"] = int(np.ceil(global_episode_index / chunks_size)) |
| info["chunks_size"] = int(chunks_size) |
| info["splits"] = {"train": f"0:{global_episode_index}"} |
| info["data_path"] = "data/chunk-{episode_chunk:03d}/episode_{episode_index:06d}.parquet" |
| info["video_path"] = "videos/chunk-{episode_chunk:03d}/{video_key}/episode_{episode_index:06d}.mp4" |
|
|
| _write_json(output / "meta" / "info.json", info) |
| _write_jsonl(output / "meta" / "tasks.jsonl", task_rows) |
| _write_jsonl(output / "meta" / "episodes.jsonl", episode_rows) |
| _write_jsonl(output / "meta" / "episodes_stats.jsonl", stats_rows) |
| if speed_metrics_rows: |
| _write_jsonl(output / "meta" / "speed_metrics.jsonl", speed_metrics_rows) |
|
|
| modality_path = inputs[0] / "meta" / "modality.json" |
| if modality_path.exists(): |
| shutil.copy2(modality_path, output / "meta" / "modality.json") |
|
|
| print(f"Wrote {global_episode_index} episodes / {global_frame_index} frames to {output}") |
|
|
|
|
| def parse_args() -> argparse.Namespace: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--inputs", nargs="+", required=True, help="Input LeRobot dataset roots") |
| parser.add_argument("--output", required=True, help="Merged output LeRobot dataset root") |
| parser.add_argument("--overwrite", action="store_true") |
| return parser.parse_args() |
|
|
|
|
| if __name__ == "__main__": |
| args = parse_args() |
| merge_datasets([Path(path) for path in args.inputs], Path(args.output), overwrite=args.overwrite) |
|
|
