code/scripts/make_demo.py

#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Demo media of the model card and the wider TT-vs-CPU agreement check on public frames (workspace tool: it reads
the public-dataset inputs and CPU runs of ``research/frnet/public_data`` and the raw datasets for the camera images;
none of them ship).

Two steps, so the device is held only for inference:

    # 1. device (tt-env, under bin/devrun): the production API on the shipped samples, the demo frames and every
    #    converted public frame (PandaSet 019 / 029 every 5th frame, all 81 nuScenes mini_val key-frames), both
    #    profiles; per-frame p150 classes + agreement with the CPU reference (ONNX Runtime fp32 runs of the research
    #    pipeline: research/frnet/public_data/runs)
    bin/devrun -t 1800 -- python code/scripts/make_demo.py run --out logs/frnet/m5_docs/demo
    # 2. CPU (research venv: OpenCV for the camera JPEGs): the renders into media/
    tools/research-venv/bin/python code/scripts/make_demo.py render --out logs/frnet/m5_docs/demo --media media

``run`` writes ``<out>/run/<frame>_<variant>.npz`` (``pred``: the p150 class of every input row, 255 = not scored:
ego crop box or profile skip; ``cpu``: the reference's, same layout) and ``<out>/agreement.json`` (per frame and per
set: raw-point argmax agreement of p150 vs CPU, N / M, timing). Agreement here is on the class decision of every raw
point the node scores; the frozen device gates (``tests/test_e2e_device.py``) additionally compare the probabilities.
"""
from __future__ import annotations

import argparse
import json
import sys
import time
from pathlib import Path

import numpy as np

HERE = Path(__file__).resolve()
CODE = HERE.parents[1]
sys.path.insert(0, str(CODE))

WS = Path("/home/ubuntu/experiments/tt-models")
PUB = WS / "research" / "frnet" / "public_data"
IGNORE = 255
VARIANTS = ("ot128", "qt128")
# camera + BEV renders (PandaSet front camera, nuScenes CAM_FRONT); TT vs CPU compares; the PandaSet BEV sequence
CAMERA_FRAMES = ["pandaset_019_40", "pandaset_029_40", "nuscenes_scene-0103_20"]
COMPARE_FRAMES = ["pandaset_019_40", "nuscenes_scene-0916_40"]
GIF_FRAMES = [f"pandaset_019_{k:02d}" for k in range(0, 80, 5)]


def public_frames():
    return sorted(p.stem for p in (PUB / "inputs").glob("*.npz"))


# ------------------------------------------------------------------------------------------------ run (device)
def cmd_run(a) -> int:
    from tt_frnet import FRNet
    from tt_frnet.ttaw.server.smoke import decode_array

    out = Path(a.out)
    (out / "run").mkdir(parents=True, exist_ok=True)
    frames = public_frames() if not a.frames else a.frames
    report = {"frames": {}, "sets": {}, "shipped": {}}
    for variant in a.variants:
        t0 = time.perf_counter()
        with FRNet.from_pretrained(variant=variant, weights_dir=a.weights_dir) as model:
            load_s = time.perf_counter() - t0
            dev = model.info.get("device")
            print(f"[{variant}] loaded in {load_s:.1f} s, device {dev}", flush=True)
            report.setdefault("device", {})[variant] = dev
            # the shipped synthetic sample vs its stored reference body
            sample = CODE / "tt_frnet" / "samples" / f"synthetic_{variant}.npz"
            res = model(sample, calibration={"preset": f"synthetic_{variant}"})
            ref = json.loads(sample.with_name(f"synthetic_{variant}.{variant}.reference.json").read_text())
            ref_labels = np.asarray(decode_array(ref["labels"]).values, np.uint8)
            np.savez_compressed(out / "run" / f"synthetic_{variant}.npz", points=res.points, pred=res.label_ids,
                                cpu=ref_labels, num_points_raw=res.num_points_raw)
            report["shipped"][variant] = {"agreement": float(np.mean(res.label_ids == ref_labels)),
                                          "n": len(res), "n_raw": int(res.num_points_raw),
                                          "timing_ms": res.timing_ms}
            print(f"[{variant}] shipped: {report['shipped'][variant]}", flush=True)
            agree_rows = {"pandaset": [0, 0], "nuscenes": [0, 0]}
            per_set = {"pandaset": [], "nuscenes": []}
            for name in frames:
                z = np.load(PUB / "inputs" / f"{name}.npz", allow_pickle=False)
                xyzi, T = z["xyzi"], z["T_base_link_sensor"]
                run = np.load(PUB / "runs" / f"{name}_{variant}.npz", allow_pickle=False)
                cpu = np.full(len(xyzi), IGNORE, np.uint8)
                cpu[run["input_rows"]] = run["class_id"]
                t1 = time.perf_counter()
                res = model(xyzi, calibration={"T_base_link_from_sensor": T.tolist()})
                ms = (time.perf_counter() - t1) * 1e3
                pred = np.full(len(xyzi), IGNORE, np.uint8)
                skipped = res.meta.get("skipped")
                if not skipped:
                    pred[res.input_index] = res.label_ids[: res.num_points_raw]
                same_rows = bool(np.array_equal(pred != IGNORE, cpu != IGNORE))
                scored = cpu != IGNORE
                n_eq = int(np.sum(pred[scored] == cpu[scored]))
                n_sc = int(scored.sum())
                rec = {"variant": variant, "skipped": skipped, "same_rows": same_rows, "n_raw": n_sc,
                       "agree": n_eq / n_sc if n_sc else None, "n": len(res), "ms": round(ms, 1)}
                report["frames"][f"{name}.{variant}"] = rec
                kind = name.split("_")[0]
                if n_sc:
                    agree_rows[kind][0] += n_eq
                    agree_rows[kind][1] += n_sc
                    per_set[kind].append(rec["agree"])
                np.savez_compressed(out / "run" / f"{name}_{variant}.npz", pred=pred, cpu=cpu)
                print(f"[{variant}] {name}: agree {rec['agree']} rows_equal {same_rows} skipped {skipped} "
                      f"{ms:.0f} ms", flush=True)
            for kind, (eq, n) in agree_rows.items():
                if n:
                    report["sets"][f"{kind}.{variant}"] = {
                        "frames": len(per_set[kind]), "points": n, "pooled": eq / n,
                        "min_frame": min(per_set[kind]), "median_frame": float(np.median(per_set[kind]))}
            print(f"[{variant}] sets: { {k: v for k, v in report['sets'].items() if k.endswith(variant)} }",
                  flush=True)
    (out / "agreement.json").write_text(json.dumps(report, indent=1, default=float))
    bad = [k for k, r in report["frames"].items() if not r["same_rows"]]
    if bad:
        print(f"scored rows differ from the CPU reference on {bad}", flush=True)
        return 1
    return 0


# ------------------------------------------------------------------------------------------------ render (CPU)
ATTR_PANDASET = ("Contains data from PandaSet (Scale AI and Hesai), https://pandaset.org, licensed under CC BY 4.0 and "
                 "the PandaSet Dataset Terms. Changes: Pandar64 sweep re-expressed in an estimated sensor frame, "
                 "front-camera image downscaled and dimmed, points and labels overlaid. Scale AI and Hesai do not "
                 "endorse this work. Cite: P. Xiao et al., PandaSet: Advanced Sensor Suite Dataset for Autonomous "
                 "Driving, ITSC 2021.")
ATTR_NUSCENES = ("Rendered from the nuScenes dataset, (c) Motional AD Inc., CC BY-NC-SA 4.0 and the nuScenes Terms of "
                 "Use (https://www.nuscenes.org/terms-of-use). Non-commercial use only; adaptations under the same "
                 "license. Motional does not endorse this work. Changes: camera image downscaled and dimmed, LiDAR "
                 "points and labels overlaid. Cite: H. Caesar et al., nuScenes: A Multimodal Dataset for Autonomous "
                 "Driving, CVPR 2020; W. K. Fong et al., Panoptic nuScenes, RA-L 2022 (lidarseg labels).")
NC_BADGE = "CC BY-NC-SA 4.0 · non-commercial"


def _research_render():
    sys.path.insert(0, str(PUB / "scripts"))
    sys.path.insert(0, str(WS / "datasets" / "tools"))
    import frnet_public_common as C  # noqa: E402  (workspace code)
    import frnet_public_render as R  # noqa: E402

    return C, R


def _agree_text(rep, key):
    r = rep["frames"].get(key)
    if not r or r["agree"] is None:
        return ""
    return f"p150 vs fp32 CPU reference: {r['agree'] * 100:.2f} % of the {r['n_raw']:,} scored raw points agree"


def render_camera(name: str, run_dir: Path, rep: dict, out_path: Path) -> None:
    import cv2  # noqa: F401  (read_jpeg)

    C, R = _research_render()
    inp = np.load(PUB / "inputs" / f"{name}.npz", allow_pickle=False)
    meta = json.loads(str(inp["meta_json"]))
    pred = np.load(run_dir / f"{name}_ot128.npz", allow_pickle=False)["pred"]
    gjson = json.loads((PUB / "runs" / f"{name}_ot128.json").read_text())
    proc = pred != IGNORE
    xyz = inp["xyzi"][:, :3]
    agree = _agree_text(rep, f"{name}.ot128")
    if name.startswith("pandaset"):
        from pandaset_meta_loader import project
        from pandaset_seq_loader import load_frame, load_sequence, read_jpeg

        seq, k = meta["sequence"], meta["frame"]
        m = load_sequence(C.PANDASET_ROOT, seq)
        fr = load_frame(C.PANDASET_ROOT, seq, k, m)
        world = fr["lidar"][inp["src_index"], :3]
        cam = fr["cameras"]["front_camera"]
        img = read_jpeg(cam["path"])
        uv, idx = project(world[proc], cam["pose"], cam["K"], img.shape[1], img.shape[0])
        sel = np.flatnonzero(proc)[idx]
        depth = np.linalg.norm(inp["xyzi"][sel, :3], axis=1)
        cam_im = R.camera_overlay(img, uv, depth, C.PALETTE_RGB[pred[sel]], R.W / img.shape[1])
        bl, br = R.bev(xyz, R.colors_for(pred, ignore_rgb=R.UNPROC)), R.bev(xyz, R.colors_for(inp["gt27"]))
        subtitle = (f"PandaSet sequence {seq}, frame {k:02d} · Hesai Pandar64, one sweep in the sensor frame · Autoware "
                    f"lidar_frnet v2.0 weights, Autoware pre- and post-processing, network on one Tenstorrent "
                    f"Blackhole p150 (N={gjson['N']:,}, M={gjson['M']:,}). {agree}.")
        cam_caption = "front_camera (52° HFOV), points coloured by the p150 class"
        cap_r = "PandaSet semseg mapped to FRNet classes (dark grey: no T4 counterpart, e.g. Ground, Pickup)"
        footer, badge = ATTR_PANDASET, None
    else:
        import nuscenes_mini as nm
        from pandaset_seq_loader import read_jpeg

        ds = nm.NuScenesMini(C.NUSC_ROOT)
        s = ds.get("sample", meta["sample_token"])
        sd = ds.sensor_data(s, "LIDAR_TOP")
        pts = ds.lidar(s)
        T_lidar = ds.T_global_ego(sd) @ ds.T_ego_sensor(sd)
        c = ds.cameras(s, ["CAM_FRONT"])["CAM_FRONT"]
        T_cam_lidar = np.linalg.inv(c["T_global_ego"] @ c["T_ego_sensor"]) @ T_lidar
        img = read_jpeg(c["path"])
        xyz_c = nm.transform_points(T_cam_lidar, pts[:, :3].astype(np.float64))
        ok = proc & (xyz_c[:, 2] > 1.0)
        uvw = xyz_c[ok] @ c["K"].T
        uv = uvw[:, :2] / uvw[:, 2:3]
        ins = (uv[:, 0] >= 0) & (uv[:, 0] < img.shape[1]) & (uv[:, 1] >= 0) & (uv[:, 1] < img.shape[0])
        sel = np.flatnonzero(ok)[ins]
        cam_im = R.camera_overlay(img, uv[ins], xyz_c[sel, 2], C.PALETTE_RGB[pred[sel]], R.W / img.shape[1],
                                  radius=2)
        bl = R.bev(xyz, R.colors_for(pred, ignore_rgb=R.UNPROC), x_range=(-30.0, 50.0), y_half=40.0)
        br = R.bev(xyz, R.colors_for(inp["gt27"]), x_range=(-30.0, 50.0), y_half=40.0)
        subtitle = (f"nuScenes {meta['scene']}, key-frame {meta['keyframe_index']} (mini_val) · Velodyne HDL-32E "
                    f"LIDAR_TOP rotated to x forward · Autoware lidar_frnet v2.0, network on one Tenstorrent Blackhole "
                    f"p150 (N={gjson['N']:,}, M={gjson['M']:,}). Domain gap: a 32-beam sensor (+10/-30°) vs the "
                    f"128-beam OT128 of the training data. {agree}.")
        cam_caption = "CAM_FRONT, points coloured by the p150 class"
        cap_r = "nuScenes-lidarseg mapped to FRNet classes (dark grey: no T4 counterpart, e.g. terrain, barrier)"
        footer, badge = ATTR_NUSCENES, NC_BADGE
    R.compose("FRNet OT128 · Tenstorrent Blackhole p150 (tt-nn)", subtitle, badge, cam_im, cam_caption, bl, br,
              "p150 prediction (dark grey: not scored, ego crop box)", cap_r, pred[proc], footer, str(out_path))


def render_compare_frame(name: str, run_dir: Path, rep: dict, out_path: Path, variant: str = "ot128") -> None:
    from tt_frnet import viz
    from tt_frnet.api import LABELS

    inp = np.load(PUB / "inputs" / f"{name}.npz", allow_pickle=False)
    r = np.load(run_dir / f"{name}_{variant}.npz", allow_pickle=False)
    pred, cpu = r["pred"], r["cpu"]
    scored = cpu != IGNORE
    xyz = inp["xyzi"][scored, :3]
    meta = json.loads(str(inp["meta_json"]))
    nusc = name.startswith("nuscenes")
    where = (f"nuScenes {meta['scene']}, key-frame {meta['keyframe_index']} (mini_val)" if nusc else
             f"PandaSet sequence {meta['sequence']}, frame {meta['frame']:02d}")
    header = (f"FRNet {variant.upper()} on {where}: every raw point the node scores, coloured by the p150 class (left) "
              f"and by the fp32 CPU reference's (ONNX Runtime on the deployed ONNX, middle); right: the points whose "
              f"class differs. {_agree_text(rep, f'{name}.{variant}')}.")
    if nusc:
        header += " Non-commercial render (CC BY-NC-SA 4.0)."
    rng = ((-30.0, 50.0), 40.0) if nusc else ((-40.0, 60.0), 50.0)
    img = viz.render_compare(xyz, pred[scored], cpu[scored], class_names=LABELS,
                             titles=(f"p150 ({variant.upper()})", "fp32 CPU reference"), header=header,
                             footer=(ATTR_NUSCENES.replace("camera image downscaled and dimmed, LiDAR points and "
                                                           "labels overlaid", "LiDAR points rendered from above")
                                     if nusc else ATTR_PANDASET.replace("front-camera image downscaled and dimmed, "
                                                                        "points and labels overlaid",
                                                                        "points rendered from above")),
                             x_range=rng[0], y_half=rng[1], size=420)
    img.save(out_path, optimize=True)
    print("wrote", out_path, img.size)


def render_shipped(run_dir: Path, rep: dict, media: Path) -> None:
    from tt_frnet import viz
    from tt_frnet.api import LABELS

    for variant in VARIANTS:
        r = np.load(run_dir / f"synthetic_{variant}.npz", allow_pickle=False)
        s = rep["shipped"][variant]
        header = (f"Shipped sample synthetic_{variant}.npz (generated by code/scripts/make_sample.py, no dataset): "
                  f"{s['n']:,} output points ({s['n_raw']:,} raw + Autoware's interpolated points). Left: the p150 "
                  f"classes ({variant.upper()} profile); middle: the stored fp32 CPU reference; right: the points "
                  f"whose class differs ({s['agreement'] * 100:.2f} % agree).")
        img = viz.render_compare(r["points"], r["pred"], r["cpu"], class_names=LABELS,
                                 titles=(f"p150 ({variant.upper()})", "fp32 CPU reference (stored)"), header=header,
                                 x_range=(-60.0, 60.0), y_half=60.0, size=420)
        out = media / f"frnet_synthetic_{variant}_tt_vs_cpu.png"
        img.save(out, optimize=True)
        print("wrote", out, img.size)


def render_gif(run_dir: Path, rep: dict, out_path: Path) -> None:
    from PIL import Image, ImageDraw

    C, R = _research_render()
    frames = []
    for name in GIF_FRAMES:
        inp = np.load(PUB / "inputs" / f"{name}.npz", allow_pickle=False)
        pred = np.load(run_dir / f"{name}_ot128.npz", allow_pickle=False)["pred"]
        im = R.bev(inp["xyzi"][:, :3], R.colors_for(pred, ignore_rgb=R.UNPROC), size=480)
        canvas = Image.new("RGB", (480, 480 + 44), R.BG)
        canvas.paste(im, (0, 22))
        d = ImageDraw.Draw(canvas)
        k = int(name.split("_")[-1])
        a = rep["frames"][f"{name}.ot128"]["agree"]
        d.text((6, 3), f"FRNet OT128 on p150 · PandaSet 019, frame {k:02d} (every 5th: 2 Hz)", font=R.font(13),
               fill=R.INK)
        d.text((6, 504), f"agreement with the fp32 CPU reference {a * 100:.2f} % · CC BY 4.0, PandaSet (Scale AI, "
                         f"Hesai)", font=R.font(11), fill=R.INK2)
        frames.append(canvas.convert("P", palette=Image.ADAPTIVE, colors=128))
    frames[0].save(out_path, save_all=True, append_images=frames[1:], duration=500, loop=0, optimize=True)
    print("wrote", out_path, f"{out_path.stat().st_size / 1e6:.2f} MB")


def cmd_render(a) -> int:
    out, media = Path(a.out), Path(a.media)
    media.mkdir(parents=True, exist_ok=True)
    run_dir = out / "run"
    rep = json.loads((out / "agreement.json").read_text())
    render_shipped(run_dir, rep, media)
    for name in CAMERA_FRAMES:
        nc = "_NC" if name.startswith("nuscenes") else ""
        render_camera(name, run_dir, rep, media / f"frnet_ot128_{name}_tt{nc}.png")
    for name in COMPARE_FRAMES:
        nc = "_NC" if name.startswith("nuscenes") else ""
        render_compare_frame(name, run_dir, rep, media / f"frnet_ot128_{name}_tt_vs_cpu{nc}.png")
    render_gif(run_dir, rep, media / "frnet_ot128_pandaset_019_bev_tt.gif")
    return 0


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    sub = ap.add_subparsers(dest="cmd", required=True)
    r = sub.add_parser("run")
    r.add_argument("--out", required=True)
    r.add_argument("--variants", nargs="*", default=list(VARIANTS))
    r.add_argument("--frames", nargs="*", help="input stems (default: every converted public frame)")
    r.add_argument("--weights-dir", default=None)
    d = sub.add_parser("render")
    d.add_argument("--out", required=True)
    d.add_argument("--media", required=True)
    a = ap.parse_args()
    return cmd_run(a) if a.cmd == "run" else cmd_render(a)


if __name__ == "__main__":
    sys.exit(main())