Blame for code/scripts/make_goldens.py

#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Goldens of streampetr-p150 from the bundle's fp32 CPU reference (``tt_streampetr.reference``), cross-checked
against the research ORT goldens (research/streampetr/public_data/golden).

    OMP_NUM_THREADS=3 PYTHONPATH=$PWD/code python code/scripts/make_goldens.py [--threads 3] [--out DIR]
    # research venv or tt-env (torch + onnx); the weights from $STREAMPETR_WEIGHTS_DIR or the workspace download

Writes (``--out``, default ``research/streampetr/goldens``; large, never shipped):
- ``pandaset_019/f40.npz``: the shipped sample's cold-start frame with the network input, the PE inputs / outputs /
  taps, every stage-level backbone tap, every head tap (tokens, temporal, the 6 decoder layers, branches,
  selection) and the 14 head outputs; ``f41..f43.npz``: the gate taps of the next frames of the same stream
  (input, ``img_feats``, memory inputs, ``outs_dec``, head outputs, selection) + ``objects.json``;
- ``public/<run>_tf_<frame>.npz``: the head taps of research teacher-forcing records (PandaSet 019 frames 1 / 40,
  nuScenes scene-0103 frame 111: internal validation data, CC BY-NC-SA, never shipped), teacher-forced from the
  record's inputs, plus the port's relative memory state of the record (``reference.port_form.global_to_relative``);
- ``report.json``: generator, versions, sha256 of every file, and the cross-checks against the research ORT goldens.

And the small goldens of the PandaSet sample (PandaSet-derived, CC BY 4.0; written beside the sample, i.e. into the
git-ignored ``staging_samples_pandaset/`` until the user approves shipping it: ``tests/paths.py`` PANDASET_*; the
``code/tt_streampetr/...`` paths below are relative to that root):
- ``code/tt_streampetr/tests/goldens/pandaset_019_f40_head.npz``: ``all_cls_scores`` / ``all_bbox_preds`` of frame 40;
- ``code/tt_streampetr/tests/goldens/pandaset_019_f40-43.json``: per frame the objects, counts, the sha256 of the
  network input / PE / head outputs and the decoder-layer histogram;
- ``code/tt_streampetr/samples/pandaset_019_f40.reference.json``: the ``to_dict()`` of the reference on the sample
  request of frame 40 (a cold start): the container smoke test's stored reference.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import platform
import sys
import time
from pathlib import Path

import numpy as np

BUNDLE = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(BUNDLE / "code"))

from tt_streampetr.host import inputs as hin  # noqa: E402
from tt_streampetr.host.agreement import detection_f1, recover_topk_indices, topk_overlap  # noqa: E402
from tt_streampetr.reference import port_form as pf  # noqa: E402
from tt_streampetr.reference.config import CAMERA_ORDER, default_config  # noqa: E402
from tt_streampetr.tests import paths  # noqa: E402
from tt_streampetr.ttaw.golden import TapRegistry, save_goldens  # noqa: E402
from tt_streampetr.ttaw.metrics import pcc  # noqa: E402

CFG = default_config()
FRAMES = (40, 41, 42, 43)
STAGE_TAPS = ("backbone.stem", "backbone.stage?", "neck.*", "img_feats", "pe.*", "pos_embed", "cone", "head.*",
              "all_cls_scores", "all_bbox_preds")
GATE_TAPS = ("img_feats", "head.outs_dec", "head.cls", "head.reg", "head.rec_score", "head.topk_idx",
             "head.reference_points", "head.mem_ego", "head.mem_ref", "all_cls_scores", "all_bbox_preds")
TF_TAPS = ("head.*", "all_cls_scores", "all_bbox_preds")
TF_EXCLUDE = ("head.tokens", "head.memory_embed", "head.sa_gamma", "head.sa_beta")     # [6000, 256] intermediates
OUT_KEYS = ("all_cls_scores", "all_bbox_preds", "post_memory_embedding", "post_memory_reference_point",
            "post_memory_timestamp", "post_memory_egopose", "post_memory_velo", "reference_points", "outs_dec",
            "topk_idx", "rec_score")


def sha(a) -> str:
    return hashlib.sha256(np.ascontiguousarray(a).tobytes()).hexdigest()


def file_sha(p: Path) -> str:
    h = hashlib.sha256()
    with open(p, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def bundle_commit() -> str:
    import subprocess

    try:
        return subprocess.run(["git", "-C", str(BUNDLE), "rev-parse", "HEAD"], capture_output=True, text=True,
                              check=True).stdout.strip()
    except (OSError, subprocess.CalledProcessError):
        return "unknown"


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--threads", type=int, default=3)
    ap.add_argument("--out", type=Path, default=paths.golden_dir())
    ap.add_argument("--skip-public", action="store_true")
    a = ap.parse_args()
    import torch

    from tt_streampetr.reference.pipeline import StreamPETRReference
    from tt_streampetr.ttaw import __version__ as ttaw_version

    torch.set_num_threads(a.threads)
    wdir = paths.weights_dir()
    if wdir is None:
        sys.exit("weights not found (STREAMPETR_WEIGHTS_DIR)")
    ref = StreamPETRReference(wdir, CFG, threads=a.threads, preprocess_workers=min(a.threads, 4))
    out_dir = a.out / "pandaset_019"
    out_dir.mkdir(parents=True, exist_ok=True)
    small = paths.PANDASET_GOLDENS
    small.mkdir(parents=True, exist_ok=True)
    meta_common = {"generator": "code/scripts/make_goldens.py", "bundle_commit": bundle_commit(),
                   "ttaw": ttaw_version, "torch": torch.__version__, "numpy": np.__version__,
                   "python": platform.python_version(), "weights": ref.net.weights.sha256(),
                   "normalization": CFG.normalization, "threads": a.threads}
    report = {"meta": meta_common, "files": {}, "checks": {}}

    # ---- A. the shipped sample: one stream from a cold start --------------------------------------------------
    ref.streams.clear()
    frames_json = []
    research_cold = paths.public_golden("ps019_f40_43_cold")
    gold_head = np.load(research_cold / "head_cls_bbox.npz") if research_cold else None
    gold_dets = ({r["frame"]: r for r in json.loads((research_cold / "detections.json").read_text())["results"]}
                 if research_cold else {})
    gold_hash = ({r["frame"]: r["img_f32_sha256"] for r in json.loads((research_cold / "inputs_sha256.json")
                                                                         .read_text())} if research_cold else {})
    checks = []
    for k, f in enumerate(FRAMES):
        t0 = time.time()
        req = hin.load_sample(paths.pandaset_sample(f))
        taps = TapRegistry(include=STAGE_TAPS if f == 40 else GATE_TAPS)
        det, outs, plan, img = ref(**req, taps=taps, return_outputs=True)
        frame = hin.prepare_frame(**req, cfg=CFG)
        tensors = {"input.img": img, "input.data_ego_pose": plan.E, "input.data_ego_pose_inv": plan.Einv,
                   "input.ts": np.float32(plan.ts), "input.M": plan.M.astype(np.float64)}
        tensors.update({f"input.{kk}": v for kk, v in plan.pre.items()})
        tensors.update({f"input.pe.{kk}": v for kk, v in frame.rig.pe_inputs().items()})
        tensors.update({f"output.{kk}": outs[kk] for kk in OUT_KEYS})
        tensors.update(taps.to_dict())
        objs = ref.postprocess(outs)
        meta = dict(meta_common, frame=f, stream_frame=plan.frame_index, fresh=bool(plan.fresh),
                    sample=f"samples/pandaset_019_f{f}.json", num_objects=len(objs["objects"]))
        path = out_dir / f"f{f}.npz"
        save_goldens(path, tensors, meta)
        recs = objs["objects"].to_records(CFG.class_names)
        entry = {"frame": f, "fresh": bool(plan.fresh), "ts": float(plan.ts), "num_pre_nms": objs["num_pre_nms"],
                 "num_objects": len(recs), "counts": {n: sum(r["label_name"] == n for r in recs)
                                                      for n in CFG.class_names},
                 "decoder_layer_histogram": np.bincount(objs["objects"].layer, minlength=6).tolist(),
                 "from_propagated_queries": int(np.sum(objs["objects"].query >= CFG.num_query)),
                 "sha256": {"img": {c: sha(img[0, i]) for i, c in enumerate(CAMERA_ORDER)},
                            "all_cls_scores": sha(outs["all_cls_scores"]), "all_bbox_preds": sha(outs["all_bbox_preds"])},
                 "objects": [{kk: (round(v, 5) if isinstance(v, float) else v) for kk, v in r.items()} for r in recs]}
        frames_json.append(entry)
        if f == 40:
            np.savez_compressed(small / "pandaset_019_f40_head.npz", all_cls_scores=outs["all_cls_scores"],
                                all_bbox_preds=outs["all_bbox_preds"])
            pos, cone = ref.position_embedding(frame.rig)
            entry["sha256"].update({"pos_embed": sha(pos), "cone": sha(cone)})
            body = det.to_dict()
            body["timing_ms"] = {}
            (paths.PANDASET_SAMPLES / "pandaset_019_f40.reference.json").write_text(json.dumps(body, indent=1) + "\n")
        c = {"frame": f, "seconds": round(time.time() - t0, 1), "golden_bytes": path.stat().st_size}
        if gold_head is not None:
            i = list(gold_head["frame"]).index(f)
            c["inputs_bit_identical"] = entry["sha256"]["img"] == gold_hash[f]
            c["cls_pcc_vs_research_ort"] = pcc(outs["all_cls_scores"], gold_head["all_cls_scores"][i])
            c["bbox_pcc_vs_research_ort"] = pcc(outs["all_bbox_preds"], gold_head["all_bbox_preds"][i])
            c["objects_vs_research_ort"] = detection_f1(objs["objects"], gold_dets[f]["dets"], dist=1.0)
        checks.append(c)
        print(json.dumps(c), flush=True)
    (out_dir / "objects.json").write_text(json.dumps({"meta": meta_common, "frames": frames_json}, indent=1))
    (small / "pandaset_019_f40-43.json").write_text(json.dumps(
        {"meta": dict(meta_common, source="PandaSet sequence 019 frames 40-43 (CC BY 4.0 + PandaSet Dataset Terms; "
                                          "samples/pandaset_019/LICENSE.txt), one stream from a cold start"),
         "frames": frames_json}, indent=1) + "\n")
    report["checks"]["pandaset_019_sample"] = checks

    # ---- B. teacher-forced head taps of research records (public datasets) ------------------------------------
    if not a.skip_public:
        pub = a.out / "public"
        pub.mkdir(exist_ok=True)
        tf_checks = []
        for run, f in (("ps019", 1), ("ps019", 40), ("ns0103", 111)):
            g = paths.public_golden(run)
            if g is None or not (g / f"tf_{f:04d}.npz").is_file():
                continue
            rec, pe = np.load(g / f"tf_{f:04d}.npz"), np.load(g / "pe_inputs_outputs.npz")
            pre = {kk: rec[kk] for kk in rec.files if kk.startswith("pre_")}
            taps = TapRegistry(include=TF_TAPS, exclude=TF_EXCLUDE)
            out = ref.head(rec["img_feats"], pe["pos_embed"], pe["cone"], rec["data_ego_pose"],
                           rec["data_ego_pose_inv"], pre, taps)
            gold_idx = recover_topk_indices(rec["post_memory_embedding"], rec["outs_dec"])
            st, M = pf.global_to_relative(pre, rec["data_ego_pose_inv"])
            tensors = {"input.img_feats": rec["img_feats"], "input.pos_embed": pe["pos_embed"], "input.cone": pe["cone"],
                       "input.data_ego_pose": rec["data_ego_pose"], "input.data_ego_pose_inv": rec["data_ego_pose_inv"],
                       "input.ts": rec["ts"], "input.golden_topk_idx": gold_idx,
                       "input.relative.emb": st.emb, "input.relative.ref": st.ref, "input.relative.ego": st.ego,
                       "input.relative.velo": st.velo, "input.relative.M": M}
            tensors.update({f"input.{kk}": v for kk, v in pre.items()})
            tensors.update({f"output.{kk}": out[kk] for kk in OUT_KEYS})
            tensors.update(taps.to_dict())
            path = pub / f"{run}_tf_{f:04d}.npz"
            save_goldens(path, tensors, dict(meta_common, run=run, frame=f, source=str(g / f"tf_{f:04d}.npz"),
                                              license="nuScenes CC BY-NC-SA 4.0: internal only" if run.startswith("ns")
                                              else "PandaSet CC BY 4.0 + Dataset Terms"))
            c = {"run": run, "frame": f, "golden_bytes": path.stat().st_size,
                 "cls_pcc_vs_research_ort": pcc(out["all_cls_scores"], rec["all_cls_scores"]),
                 "bbox_pcc_vs_research_ort": pcc(out["all_bbox_preds"], rec["all_bbox_preds"]),
                 "outs_dec_pcc_vs_research_ort": pcc(out["outs_dec"], rec["outs_dec"]),
                 "top256_overlap_vs_research_ort": topk_overlap(out["topk_idx"].reshape(-1), gold_idx)}
            tf_checks.append(c)
            print(json.dumps(c), flush=True)
        report["checks"]["teacher_forced_public"] = tf_checks

    for p in sorted(a.out.rglob("*")):
        if p.is_file() and p.name != "report.json":
            report["files"][str(p.relative_to(a.out))] = {"bytes": p.stat().st_size, "sha256": file_sha(p)}
    (a.out / "report.json").write_text(json.dumps(report, indent=1, default=float))
    print(f"wrote {len(report['files'])} files to {a.out}")


if __name__ == "__main__":
    os.environ.setdefault("OMP_NUM_THREADS", "3")
    main()