| |
| """Extract a low-variation Bridge/LeRobot subset by episode task text. |
| |
| The default selection is the plate-placement family (``put ... on plate``), |
| which is close to the current carrot-on-plate evaluation while retaining |
| multiple objects and camera scenes. Episodes are copied whole, so no video |
| re-encoding or temporal duplication is introduced. |
| """ |
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import os |
| import re |
| import shutil |
| from pathlib import Path |
|
|
|
|
| def main() -> None: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--src", default="/opt/zhangchenyu/datasets/bridge_orig_lerobot_smoke") |
| ap.add_argument("--dst", default=".artifacts/bridge_put_on_plate_259") |
| ap.add_argument("--pattern", default=r"^put .+ on plate(?: sink)?$") |
| ap.add_argument("--max-episodes", type=int, default=None) |
| ap.add_argument("--min-length", type=int, default=None, |
| help="keep only episodes with at least this many frames") |
| ap.add_argument("--hardlink-videos", action="store_true", |
| help="hardlink videos instead of duplicating them when possible") |
| args = ap.parse_args() |
| src, dst = Path(args.src), Path(args.dst) |
| meta = src / "meta" |
| episodes = [json.loads(x) for x in (meta / "episodes.jsonl").read_text().splitlines()] |
| tasks = {int(x["task_index"]): x["task"] for x in |
| (json.loads(x) for x in (meta / "tasks.jsonl").read_text().splitlines())} |
| rx = re.compile(args.pattern, re.I) |
| selected = [] |
| for ep in episodes: |
| text = " ".join(ep.get("tasks", [])) |
| if rx.fullmatch(text.strip()): |
| selected.append((ep, text.strip())) |
| selected.sort(key=lambda x: int(x[0]["episode_index"])) |
| if args.min_length is not None: |
| selected = [x for x in selected if int(x[0]["length"]) >= args.min_length] |
| if args.max_episodes is not None: |
| selected = selected[:args.max_episodes] |
| if not selected: |
| raise SystemExit("no matching episodes") |
|
|
| missing = [] |
| for ep, _ in selected: |
| i = int(ep["episode_index"]) |
| chunk = i // 1000 |
| dp = src / f"data/chunk-{chunk:03d}/episode_{i:06d}.parquet" |
| vp = src / f"videos/chunk-{chunk:03d}/observation.images.image_0/episode_{i:06d}.mp4" |
| if not dp.exists() or not vp.exists(): |
| missing.extend(str(p) for p in (dp, vp) if not p.exists()) |
| if missing: |
| raise SystemExit("raw parquet/video files are missing; metadata-only source. First missing: " + missing[0]) |
|
|
| |
| import pyarrow.parquet as pq |
| import pyarrow as pa |
|
|
| (dst / "meta").mkdir(parents=True, exist_ok=True) |
| out_eps, out_tasks = [], [] |
| total_frames = 0 |
| for new_i, (ep, text) in enumerate(selected): |
| old_i = int(ep["episode_index"]) |
| chunk = old_i // 1000 |
| old_dp = src / f"data/chunk-{chunk:03d}/episode_{old_i:06d}.parquet" |
| old_vp = src / f"videos/chunk-{chunk:03d}/observation.images.image_0/episode_{old_i:06d}.mp4" |
| new_dp = dst / f"data/chunk-000/episode_{new_i:06d}.parquet" |
| new_vp = dst / f"videos/chunk-000/observation.images.image_0/episode_{new_i:06d}.mp4" |
| new_dp.parent.mkdir(parents=True, exist_ok=True) |
| new_vp.parent.mkdir(parents=True, exist_ok=True) |
| tab = pq.read_table(old_dp) |
| |
| cols = {name: tab[name] for name in tab.column_names} |
| if "episode_index" in cols: |
| cols["episode_index"] = pa.array([new_i] * tab.num_rows, type=cols["episode_index"].type) |
| task_col = pa.array([new_i] * tab.num_rows, type=cols["task_index"].type) |
| cols["task_index"] = task_col |
| pq.write_table(pa.table(cols), new_dp) |
| if args.hardlink_videos: |
| os.link(old_vp, new_vp) |
| else: |
| shutil.copy2(old_vp, new_vp) |
| out_eps.append({"episode_index": new_i, "tasks": [text], "length": int(ep["length"])}) |
| out_tasks.append({"task_index": new_i, "task": text}) |
| total_frames += int(ep["length"]) |
|
|
| info = json.loads((meta / "info.json").read_text()) |
| info["total_episodes"] = len(out_eps) |
| info["total_frames"] = total_frames |
| info["total_tasks"] = len(out_tasks) |
| info["total_videos"] = len(out_eps) |
| info["total_chunks"] = 1 |
| info["splits"] = {"train": f"0:{len(out_eps)}"} |
| 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" |
| (dst / "meta/info.json").write_text(json.dumps(info, indent=2) + "\n") |
| (dst / "meta/episodes.jsonl").write_text("".join(json.dumps(x) + "\n" for x in out_eps)) |
| (dst / "meta/tasks.jsonl").write_text("".join(json.dumps(x) + "\n" for x in out_tasks)) |
| for name in ("modality.json", "stats.json", "stats_gr00t.json"): |
| if (meta / name).exists(): shutil.copy2(meta / name, dst / "meta" / name) |
| print(json.dumps({"episodes": len(out_eps), "frames": total_frames, "dst": str(dst)}, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
