| """Render a task video into one contact sheet for visual review. |
| |
| The numeric check only reads final object poses, and it will pass an episode where the cup landed |
| upside down or a gripper hung empty in the air. Looking at the frames is the only way to catch |
| those, so every task gets a sheet before it is called done. |
| """ |
| import argparse, os |
| import imageio.v2 as iio |
| from PIL import Image, ImageDraw |
|
|
| ap = argparse.ArgumentParser() |
| ap.add_argument("videos", nargs="+") |
| ap.add_argument("--out", default="outputs/sheets") |
| ap.add_argument("--cols", type=int, default=6) |
| ap.add_argument("--w", type=int, default=340) |
| a = ap.parse_args() |
| os.makedirs(a.out, exist_ok=True) |
|
|
| for v in a.videos: |
| if not os.path.exists(v): |
| print(f"[sheet] MISSING {v}"); continue |
| r = iio.get_reader(v) |
| n = r.count_frames() |
| idx = [min(n-1, int(n*f)) for f in [i/(a.cols-1) for i in range(a.cols)]] |
| idx[0] = min(2, n-1) |
| frames = [Image.fromarray(r.get_data(i)) for i in idx] |
| fw, fh = frames[0].size |
| h = int(a.w*fh/fw) |
| sheet = Image.new("RGB", (a.w*len(frames), h), (20, 20, 22)) |
| for k, f in enumerate(frames): |
| sheet.paste(f.resize((a.w, h), Image.LANCZOS), (a.w*k, 0)) |
| ImageDraw.Draw(sheet).text((a.w*k+6, h-16), f"{idx[k]}/{n}", fill=(255, 220, 60)) |
| p = os.path.join(a.out, os.path.basename(v).replace(".mp4", ".png")) |
| sheet.save(p) |
| print(f"[sheet] {p} ({n} frames)") |
|
|
