File size: 1,418 Bytes
7399b6f | 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 | """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)")
|