scripts/concat_episodes.py

"""Stitch grape->box episode videos into one labelled clip with a SUCCESS/FAIL title card
per episode and a final success-rate card.

Usage:  OUT=/path/to/out python scripts/concat_episodes.py
        (reads OUT/ep_results.txt and OUT/ep*.mp4, writes OUT/episodes.mp4)
"""
import imageio.v2 as imageio, numpy as np, re, os
from PIL import Image, ImageDraw

OUT = os.environ.get("OUT", "/tmp/yam_grape_eps")
eps = []
for ln in open(os.path.join(OUT, "ep_results.txt")):
    m = re.search(r"EP(\d+) grape=(\S+) box=(\S+) ::.*EPISODE_RESULT:\s*(\w+)", ln)
    if m:
        eps.append((int(m.group(1)), m.group(2), m.group(3), m.group(4)))
    else:
        m2 = re.search(r"EP(\d+) grape=(\S+) box=(\S+)", ln)
        if m2:
            eps.append((int(m2.group(1)), m2.group(2), m2.group(3), "FAIL"))
eps.sort()
succ = sum(1 for e in eps if e[3] == "SUCCESS")

W, H = 720, 544
for (i, _, _, _) in eps:
    fp = os.path.join(OUT, f"ep{i}.mp4")
    if os.path.exists(fp):
        f0 = next(iter(imageio.get_reader(fp))); H, W = f0.shape[0], f0.shape[1]; break

def fit(a):
    if a.shape[0] != H or a.shape[1] != W:
        a = np.array(Image.fromarray(a).resize((W, H)))
    return a[..., :3]

def card(lines, color=(255, 235, 60)):
    im = Image.new("RGB", (W, H), (10, 12, 20)); d = ImageDraw.Draw(im)
    y = H // 2 - len(lines) * 18
    for ln in lines:
        d.text((max(8, W // 2 - len(ln) * 5), y), ln, fill=color); y += 36
    return np.array(im)

allf = [card(["YAM  grape -> box", "PRM + real friction grasp (no attach)", f"{len(eps)} episodes, varied positions"])] * 10
for (i, gx, bx, res) in eps:
    col = (90, 230, 120) if res == "SUCCESS" else (240, 90, 90)
    allf += [card([f"EPISODE {i}", f"grape=({gx})  box=({bx})", f"result: {res}"], col)] * 8
    fp = os.path.join(OUT, f"ep{i}.mp4")
    if os.path.exists(fp):
        for f in imageio.get_reader(fp): allf.append(fit(f))
allf += [card([f"SUCCESS RATE:  {succ} / {len(eps)}"], (90, 230, 120) if succ > len(eps) // 2 else (255, 200, 60))] * 26
out = os.path.join(OUT, "episodes.mp4")
imageio.mimsave(out, allf, fps=7, macro_block_size=1)
print(f"[concat] {succ}/{len(eps)} success -> {out}  ({len(allf)} frames, {W}x{H})")