visualization/tools/annotate_rollouts.py

import json
import sys
from collections import defaultdict
from pathlib import Path

import imageio
import numpy as np
from PIL import Image, ImageDraw, ImageFont

ROOT = Path("/pfss/mlde/workspaces/mlde_wsp_IAS_SAMMerge/triquang/graphvla")
VIDEO = ROOT / "runs/viz_rollouts/ctrl5wngml_k32_alr5e-05_mlr0.0003_s0/checkpoint-72000/video"
OUT = ROOT / "viz/out/07_rollouts"
KINDS = ["pick", "place", "table", "press", "move", "static", "insert", "none"]
ORDINALS = ["-", "first", "second", "third", "fourth", "fifth"]
MARGIN, SPAN, SCALE, BAR = 6.4, 243.2, 2, 150


def font():
    from matplotlib import font_manager

    regular = font_manager.findfont(font_manager.FontProperties(family="DejaVu Sans Mono"))
    bold = font_manager.findfont(font_manager.FontProperties(family="DejaVu Sans Mono", weight="bold"))
    return ImageFont.truetype(regular, 15), ImageFont.truetype(bold, 17)


def plog():
    out = defaultdict(list)
    for line in open(ROOT / "logs/planner/plog_viz.jsonl"):
        d = json.loads(line)
        out[d["session"]].append(d)
    return out


def describe(p):
    if p["kind"] is None:
        return "none"
    k = KINDS[p["kind"]]
    o = ORDINALS[p["ordinal"]] if p["ordinal"] else ""
    return f"{k} {o}".strip()


def mark(draw, point, color, shape):
    if point is None or min(point) < 0:
        return
    y, x = (point[0] * SPAN + MARGIN) * SCALE, (point[1] * SPAN + MARGIN) * SCALE
    r = 11
    if shape == "o":
        draw.ellipse([x - r, y - r, x + r, y + r], outline=color, width=3)
    else:
        draw.line([x - r, y - r, x + r, y + r], fill=color, width=4)
        draw.line([x - r, y + r, x + r, y - r], fill=color, width=4)


def annotate(mp4, log):
    task, _, episode, status = mp4.stem.split("_")
    trace = json.loads(mp4.with_suffix(".json").read_text())
    entries = log.get(f"{task}:test:{episode}", [])
    start = entries[0]["frame"] if entries else 0
    regular, bold = font()
    out = []
    j = -1
    for t, frame in enumerate(imageio.v3.imiter(mp4)):
        while j + 1 < len(entries) and entries[j + 1]["frame"] - start <= t:
            j += 1
        h, w = frame.shape[:2]
        img = Image.new("RGB", (w * SCALE, h * SCALE + BAR), "white")
        img.paste(Image.fromarray(frame).resize((w * SCALE, h * SCALE), Image.NEAREST), (0, 0))
        draw = ImageDraw.Draw(img)
        y0 = h * SCALE + 6
        step = trace[min(t, len(trace) - 1)]
        draw.text((8, y0), f"{task}  test episode {episode}  ({status})   step {t}   controller query {step['query']}", font=bold, fill="black")
        draw.text((8, y0 + 26), f"label subgoal  : {step['subgoal'][:80]}", font=regular, fill=(0, 120, 0))
        if j >= 0:
            e = entries[j]
            mark(draw, e["truth"]["point"], (0, 255, 0), "o")
            mark(draw, e["pred"]["point"], (0, 255, 255), "x")
            ok = (e["pred"]["kind"], e["pred"]["ordinal"]) == (e["truth"]["kind"], e["truth"]["ordinal"])
            draw.text((8, y0 + 50), f"planner step   : {describe(e['pred'])}   (label: {describe(e['truth'])})", font=regular, fill="black" if ok else (200, 0, 0))
            draw.text((8, y0 + 74), f"graph nodes    : {e['nodes']}   last planner call at step {e['frame'] - start}", font=regular, fill="black")
        draw.text((8, y0 + 100), "front camera (left) and wrist camera (right).  green circle = label target, cyan x = planner target", font=regular, fill=(90, 90, 90))
        out.append(np.asarray(img))
    OUT.mkdir(parents=True, exist_ok=True)
    imageio.mimwrite(OUT / mp4.name, out, fps=20, codec="libx264")


def main():
    log = plog()
    names = sys.argv[1:] or sorted(p.name for p in VIDEO.glob("*.mp4"))
    for name in names:
        annotate(VIDEO / name, log)
        print("done", name, flush=True)


if __name__ == "__main__":
    main()