File size: 2,197 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 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | """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})")
|