File size: 5,235 Bytes
68eed61 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | #!/usr/bin/env python3
"""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 only after the source check: this script remains usable for dry inspection.
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)
# Keep all rows but rewrite episode/task ids to a compact standalone set.
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()
|