code/scripts/make_goldens.py

#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""CPU-reference goldens of transfusion-p150 (research venv: torch + onnxruntime, CPU only, <= 4 threads).

    cd bundles/transfusion-p150 && OMP_NUM_THREADS=4 PYTHONPATH=code \\
        /home/ubuntu/experiments/tt-models/tools/research-venv/bin/python -I code/scripts/make_goldens.py \\
        --weights /home/ubuntu/experiments/tt-models/assets/transfusion/hf_lidar_transfusion \\
        --research /home/ubuntu/experiments/tt-models/research/transfusion \\
        --out /home/ubuntu/experiments/tt-models/research/transfusion/goldens

For every frame the fp32 torch reference (``reference/pipeline.py``: host pre-processing, the deployed network rebuilt
from the ONNX initializers, host post-processing) runs once; ONNX Runtime runs the deployed ONNX on the same voxel
tensors twice (plain: the bitwise output golden; with the module intermediates exposed: the agreement check), and the
reference runs once more teacher-forced with ORT's top-k indices. The agreement (PCC per module tap before the
selection, top-500 set overlap and strong-proposal coverage, teacher-forced PCC after it, detection match) is recorded
with the golden. Written:

- ``<out>/<frame>.io.npz``: the ONNX inputs (``voxels``, ``num_points``, ``coors``), the decorated features of the
  kept pillars (``features``), the reference outputs (``cls_score0``, ``bbox_pred0``, ``dir_cls_pred0``) and the
  plain ORT outputs (``ort_cls_score0``, ...; float32);
- ``<out>/<frame>.detections.json``: the reference output (``Detections3D.to_dict``), per-stage counts, pillar and
  sampling statistics, the ORT agreement;
- ``<out>/<frame>.taps.npz`` (``--taps`` frames only; ~100-200 MB each): every tap of ``reference.model.MODULE_TAPS``
  plus ``pfn.layer0`` and ``decoder.key_pos_embed`` (float32; ``head.topk_indices`` int64): the per-module goldens of
  the device PCC gates, including the indices that teacher-force the device decoder;
- ``<out>/manifest.json``: sha256 of every file, tool versions, inputs;
- in the bundle (small, shippable: Apache-2.0 / CC BY 4.0 inputs only): ``tests/goldens/voxels.json`` (voxel tensor
  hashes of the samples), ``tests/goldens/detections/<sample>.json`` and the smoke reference
  ``samples/<stem>.reference.json``. The PandaSet sample is not shipped until the user approves PandaSet samples
  (research/DATASETS.md N5): its input is read from, and its small golden and reference are written to, the
  git-ignored ``staging_samples_pandaset/`` of the bundle (``goldens/`` there) unless it already sits in ``samples/``.

``--set selection`` (fix round 1) adds ``<out>/<frame>.selection.npz`` to every development frame above: the
reference's top-500 (``head.topk_indices`` int64, ``head.topk_values``) from a fresh reference run, whose outputs must
equal the stored ``io.npz`` bit for bit (nothing else is rewritten): the proposal identities of the flip-free
proposal-level agreement (``tests/_agreement.py``).

``--set heldout`` (fix round 1) writes compact goldens of the HELD-OUT frames (``tests/_data.py HELDOUT``: 34 public
frames never used for development or precision choices, converted by the research converter into
``research/transfusion/heldout/inputs``) to ``<out>/heldout/``: ``<frame>.io.npz`` (reference outputs, top-500
indices / values, plain ORT outputs) and ``<frame>.detections.json`` (output, counts, proposals and the same ORT
agreement record as the development goldens: pre-selection PCC, top-k sets, teacher-forced PCC, detections), plus
``heldout/manifest.json``. No module taps are stored.

Frames: the shipped samples (``pandaset_019_f40``, ``test_pcd``), the 16 public frames of
research/transfusion/public_data (PandaSet CC BY 4.0; nuScenes / Argoverse 2 non-commercial: research only) and the
Autoware sample-rosbag frames 50 / 150 / 250 (2 sweeps; frame 150 has 62,527 pillars; license unstated: research
only), plus frame 150 through the ``autoware_compat`` shuffle emulation (offset 12345, the research's).
"""
from __future__ import annotations

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

import numpy as np

HERE = Path(__file__).resolve().parent
CODE = HERE.parent
sys.path.insert(0, str(CODE))

from tt_transfusion import __version__  # noqa: E402
from tt_transfusion import io as tio  # noqa: E402
from tt_transfusion.reference.model import MODULE_TAPS, SELECTION_TAPS  # noqa: E402
from tt_transfusion.reference.ort import OrtTransFusion  # noqa: E402
from tt_transfusion.reference.pipeline import TransFusionReference  # noqa: E402
from tt_transfusion.ttaw import __version__ as ttaw_version  # noqa: E402
from tt_transfusion.ttaw.golden import TapRegistry, save_goldens  # noqa: E402
from tt_transfusion.ttaw.metrics import match_detections, pcc  # noqa: E402

SAMPLES = CODE / "tt_transfusion" / "samples"
BUNDLE_GOLDENS = CODE / "tt_transfusion" / "tests" / "goldens"
PANDASET_STAGING = CODE.parent / "staging_samples_pandaset"   # git-ignored, not shipped (DATASETS.md N5)
SHIPPED = ("pandaset_019_f40", "test_pcd")


def sample_home(name: str) -> Path:
    """Where a sample lives: ``samples/`` (shipped), or the PandaSet staging copy until it ships."""
    if name.startswith("pandaset") and not (SAMPLES / f"{name}.npz").is_file():
        return PANDASET_STAGING
    return SAMPLES


def small_golden_home(name: str) -> Path:
    return PANDASET_STAGING / "goldens" if sample_home(name) == PANDASET_STAGING else BUNDLE_GOLDENS / "detections"
PUBLIC = ("ps019_f10", "ps019_f40", "ps019_f70", "ps005_f40", "ps090_f10", "ps090_f40", "ps090_f70", "ps029_f40",
          "ns0103_k10", "ns0103_k20", "ns0103_k30", "ns0916_k10", "ns0916_k25", "av2_280269f9_s080",
          "av2_51bbdd4d_s080", "av2_e13c06cb_s080")
ROSBAG = ("frame_0050", "frame_0150", "frame_0250")
FIELDS5 = ("x", "y", "z", "intensity", "time_lag")
COMPAT = {"mode": "autoware_compat", "offset": 12345}
TAP_SET = tuple(MODULE_TAPS) + ("pfn.layer0", "decoder.key_pos_embed")
STRONG = 0.1


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


def file_sha(p: Path) -> str:
    return hashlib.sha256(p.read_bytes()).hexdigest()


def frames(research: Path):
    """(name, source, sampling, license, shippable)."""
    for name in SHIPPED:
        lic = "CC BY 4.0 + PandaSet Dataset Terms" if name.startswith("pandaset") else "Apache-2.0 (autoware_universe)"
        yield name, sample_home(name) / f"{name}.npz", None, lic, True
    for name in PUBLIC:
        path = research / "public_data" / "inputs" / f"{name}.npy"
        if not path.is_file():
            print(f"skip {name}: {path} missing")
            continue
        meta = json.loads(path.with_suffix(".json").read_text())
        lic = meta.get("license") or ("CC BY 4.0 + PandaSet Dataset Terms" if name.startswith("ps") else
                                      "CC BY-NC-SA 4.0 (research only)")
        yield f"public_{name}", tio.PointCloud(np.load(path, allow_pickle=False), FIELDS5), None, lic, False
    for name in ROSBAG:
        path = research / "sample_frames" / f"{name}_sweeps2.npy"
        if not path.is_file():
            print(f"skip {name}: {path} missing")
            continue
        cloud = tio.PointCloud(np.load(path, allow_pickle=False), FIELDS5)
        lic = "Autoware sample rosbag (license unstated): research only"
        yield f"rosbag_{name[-4:]}", cloud, None, lic, False
        if name == "frame_0150":
            yield "rosbag_0150_compat", cloud, COMPAT, lic, False


def agreement(rt, ro_tapped, run, run_o, tt: TapRegistry, to: TapRegistry) -> dict:
    """Pre-selection PCC per tap, top-k set agreement, teacher-forced PCC after the selection, detection match."""
    pre = {k: float(pcc(tt[k], to[k].reshape(tt[k].shape))) for k in to.names()
           if k in tt and k not in SELECTION_TAPS}
    mine, ref = tt["head.topk_indices"].ravel(), to["head.topk_indices"].ravel()
    strong = ref[to["head.topk_values"].ravel() >= STRONG]
    tf = TapRegistry(include=SELECTION_TAPS)
    rt.network(run.prepared, tf, topk_indices=to["head.topk_indices"])
    post = {k: float(pcc(tf[k], to[k].reshape(tf[k].shape))) for k in SELECTION_TAPS
            if k != "head.topk_indices" and k in to}
    a, b = run.output, run_o.output
    m = match_detections(a.boxes, a.label_ids, a.scores, b.boxes, b.label_ids, b.scores, max_dist=0.01)
    return {"pre_selection_min_pcc": min(pre.values()), "pre_selection": pre,
            "topk": {"shared": int(len(set(mine.tolist()) & set(ref.tolist()))), "k": int(ref.size),
                     "positional_equal": float(np.mean(mine == ref)), "strong": int(len(strong)),
                     "strong_covered": bool(set(strong.tolist()) <= set(mine.tolist()))},
            "teacher_forced_min_pcc": min(post.values()), "teacher_forced": post,
            "detections": {"counts_equal": run.counts == run_o.counts, "ort_counts": run_o.counts,
                           "recall": m.recall, "precision": m.precision,
                           "max_score_err": max(m.score_errors, default=0.0)}}


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--weights", type=Path, required=True)
    ap.add_argument("--research", type=Path, required=True)
    ap.add_argument("--out", type=Path, required=True)
    ap.add_argument("--taps", default="pandaset_019_f40,public_ns0103_k20,rosbag_0150",
                    help="frames whose module taps are saved")
    ap.add_argument("--threads", type=int, default=4)
    ap.add_argument("--only", default="", help="comma list of frame names (default: all)")
    ap.add_argument("--set", default="dev", choices=("dev", "selection", "heldout"),
                    help="dev: the development goldens (default); selection: add <frame>.selection.npz to them; "
                         "heldout: the compact held-out goldens in <out>/heldout")
    a = ap.parse_args()
    if a.set == "selection":
        return selection_main(a)
    if a.set == "heldout":
        return heldout_main(a)
    import onnxruntime
    import torch

    torch.set_num_threads(a.threads)
    taps_for = set(filter(None, a.taps.split(",")))
    only = set(filter(None, a.only.split(",")))
    a.out.mkdir(parents=True, exist_ok=True)
    (BUNDLE_GOLDENS / "detections").mkdir(parents=True, exist_ok=True)
    manifest_path = a.out / "manifest.json"
    manifest = json.loads(manifest_path.read_text()) if manifest_path.is_file() else {"files": {}}
    rt = TransFusionReference(a.weights, backend="torch", threads=a.threads)
    ro = TransFusionReference(a.weights, backend="ort", threads=a.threads, intermediates=TAP_SET)
    plain = OrtTransFusion(a.weights, threads=a.threads, weights=rt.weights)
    manifest.update({"bundle_version": __version__, "ttaw": ttaw_version, "torch": torch.__version__,
                     "onnxruntime": onnxruntime.__version__, "numpy": np.__version__,
                     "python": platform.python_version(), "onnx_sha256": rt.weights.sha256,
                     "weights": "AutowareFoundation/lidar_transfusion@8f07eb93 (v2.1)",
                     "config": rt.config.describe(), "written": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
                     "command": " ".join(sys.argv)})
    voxels = {}
    if (BUNDLE_GOLDENS / "voxels.json").is_file():
        voxels = json.loads((BUNDLE_GOLDENS / "voxels.json").read_text())
    for name, src, sampling, lic, shippable in frames(a.research):
        if only and name not in only:
            continue
        t0 = time.perf_counter()
        tt, to = TapRegistry(include=TAP_SET), TapRegistry(include=TAP_SET)
        run = rt.run(src, sampling=sampling, taps=tt)
        run_o = ro.run(src, sampling=sampling, taps=to)
        p = run.prepared
        n = p.num_pillars
        assert np.array_equal(p.features, run_o.prepared.features), name
        inputs = p.onnx_inputs()
        ort_out = plain.run(inputs)
        agree = agreement(rt, ro, run, run_o, tt, to)
        t1 = time.perf_counter()
        io_path = a.out / f"{name}.io.npz"
        save_goldens(io_path, {**inputs, "features": p.features[:n], **run.outputs,
                               **{f"ort_{k}": v for k, v in ort_out.items()}},
                     meta={"frame": name, "num_pillars": n, "license": lic, "shippable": shippable,
                           "sampling": p.sampling, "points_sha256": sha(p.points)}, compress=True)
        files = [io_path]
        if name in taps_for:
            taps_path = a.out / f"{name}.taps.npz"
            tt.save(taps_path, meta={"frame": name, "taps": tt.names(), "sampling": p.sampling}, compress=True)
            files.append(taps_path)
        body = run.output.to_dict()
        rec = {"frame": name, "license": lic, "shippable": shippable,
               "reference": "tt_transfusion.reference torch fp32 (ONNX initializers, BN folded)",
               "counts": run.counts, "pillars": {k: v for k, v in p.stats().items() if k != "densify"},
               "densify": p.densify, "ort_agreement": agree,
               "proposals": [int(v) for v in run.objects.extra["proposal"]], "output": body}
        det_path = a.out / f"{name}.detections.json"
        det_path.write_text(json.dumps(rec, indent=1))
        files.append(det_path)
        for f in files:
            manifest["files"][f.name] = {"sha256": file_sha(f), "bytes": f.stat().st_size}
        if name in SHIPPED:
            small = {k: rec[k] for k in ("frame", "license", "counts", "pillars", "proposals", "output")}
            small["ort_agreement"] = {"pre_selection_min_pcc": agree["pre_selection_min_pcc"],
                                      "topk": agree["topk"], "teacher_forced_min_pcc": agree["teacher_forced_min_pcc"],
                                      "detections": agree["detections"]}
            small_golden_home(name).mkdir(parents=True, exist_ok=True)
            (small_golden_home(name) / f"{name}.json").write_text(json.dumps(small, indent=1))
            (sample_home(name) / f"{name}.reference.json").write_text(json.dumps(body, indent=1))
            voxels[f"{name}.npz"] = {"num_pillars": n, "num_pillars_total": p.pillars.num_pillars_total,
                                     "points_sha256": sha(p.points), "features_sha256": sha(p.features[:n]),
                                     "coords_sha256": sha(p.coords[:n]),
                                     "num_points_sha256": sha(p.pillars.num_points[:n]),
                                     "canvas_index_sha256": sha(p.canvas_index)}
        print(f"{name}: pillars {n}/{p.pillars.num_pillars_total} counts {run.counts} | ort pre-sel min pcc "
              f"{agree['pre_selection_min_pcc']:.10f}, top-k shared {agree['topk']['shared']}/{agree['topk']['k']} "
              f"strong {agree['topk']['strong']} covered {agree['topk']['strong_covered']}, teacher-forced min pcc "
              f"{agree['teacher_forced_min_pcc']:.10f}, detections recall/precision "
              f"{agree['detections']['recall']:.3f}/{agree['detections']['precision']:.3f} "
              f"({t1 - t0:.1f} s + {time.perf_counter() - t1:.1f} s io)", flush=True)
        manifest_path.write_text(json.dumps(manifest, indent=1))
    if voxels:
        (BUNDLE_GOLDENS / "voxels.json").write_text(json.dumps(voxels, indent=1))
    return 0


SEL_TAPS = ("head.topk_indices", "head.topk_values")


def _manifest(path: Path, a, rt, extra: dict) -> dict:
    import onnxruntime
    import torch

    m = json.loads(path.read_text()) if path.is_file() else {"files": {}}
    m.setdefault("files", {})
    m.update({"bundle_version": __version__, "ttaw": ttaw_version, "torch": torch.__version__,
              "onnxruntime": onnxruntime.__version__, "numpy": np.__version__, "python": platform.python_version(),
              "onnx_sha256": rt.weights.sha256, "weights": "AutowareFoundation/lidar_transfusion@8f07eb93 (v2.1)",
              "config": rt.config.describe(), "written": time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()),
              "command": " ".join(sys.argv), **extra})
    return m


def selection_main(a) -> int:
    """``--set selection``: ``<frame>.selection.npz`` (the reference top-500) next to every development golden; the
    fresh reference outputs must equal the stored ones bit for bit."""
    import torch

    torch.set_num_threads(a.threads)
    only = set(filter(None, a.only.split(",")))
    rt = TransFusionReference(a.weights, backend="torch", threads=a.threads)
    manifest_path = a.out / "manifest.json"
    manifest = json.loads(manifest_path.read_text())
    for name, src, sampling, _lic, _ship in frames(a.research):
        if only and name not in only:
            continue
        io_path = a.out / f"{name}.io.npz"
        if not io_path.is_file():
            print(f"skip {name}: {io_path} missing")
            continue
        taps = TapRegistry(include=SEL_TAPS)
        run = rt.run(src, sampling=sampling, taps=taps)
        with np.load(io_path, allow_pickle=False) as z:
            for k, v in run.outputs.items():
                if not np.array_equal(np.asarray(v), z[k]):
                    raise SystemExit(f"{name}: the reference output {k} differs from the stored golden")
        path = a.out / f"{name}.selection.npz"
        save_goldens(path, {k: taps[k] for k in SEL_TAPS}, meta={"frame": name, "io_sha256": file_sha(io_path)},
                     compress=True)
        manifest["files"][path.name] = {"sha256": file_sha(path), "bytes": path.stat().st_size}
        print(f"{name}: selection.npz written (strong {int((taps['head.topk_values'] >= STRONG).sum())}), outputs "
              f"== stored golden", flush=True)
    manifest["selection_written"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
    manifest_path.write_text(json.dumps(manifest, indent=1))
    return 0


def heldout_main(a) -> int:
    """``--set heldout``: compact goldens of the held-out frames (module docstring)."""
    import torch

    sys.path.insert(0, str(CODE))
    from tt_transfusion.tests._data import HELDOUT

    torch.set_num_threads(a.threads)
    only = set(filter(None, a.only.split(",")))
    out = a.out / "heldout"
    out.mkdir(parents=True, exist_ok=True)
    rt = TransFusionReference(a.weights, backend="torch", threads=a.threads)
    ro = TransFusionReference(a.weights, backend="ort", threads=a.threads, intermediates=TAP_SET)
    plain = OrtTransFusion(a.weights, threads=a.threads, weights=rt.weights)
    inputs_dir = a.research / "heldout" / "inputs"
    manifest_path = out / "manifest.json"
    manifest = _manifest(manifest_path, a, rt, {"set": "heldout", "inputs": str(inputs_dir)})
    for fid in HELDOUT:
        name = f"heldout_{fid}"
        if only and name not in only and fid not in only:
            continue
        path = inputs_dir / f"{fid}.npy"
        if not path.is_file():
            print(f"skip {fid}: {path} missing")
            continue
        t0 = time.perf_counter()
        cloud = tio.PointCloud(np.load(path, allow_pickle=False), FIELDS5)
        taps, to = TapRegistry(include=TAP_SET), TapRegistry(include=TAP_SET)
        run = rt.run(cloud, taps=taps)
        run_o = ro.run(cloud, taps=to)
        p = run.prepared
        assert np.array_equal(p.features, run_o.prepared.features), name
        inputs = p.onnx_inputs()
        ort_out = plain.run(inputs)
        agree = agreement(rt, ro, run, run_o, taps, to)     # the development goldens' ORT agreement record
        m = agree["detections"]
        io_path = out / f"{name}.io.npz"
        save_goldens(io_path, {**run.outputs, **{k: taps[k] for k in SEL_TAPS},
                               **{f"ort_{k}": v for k, v in ort_out.items()}},
                     meta={"frame": name, "num_pillars": p.num_pillars, "points_sha256": sha(p.points),
                           "input_sha256": file_sha(path)}, compress=True)
        meta = json.loads(path.with_suffix(".json").read_text())
        lic = meta.get("license") or ("CC BY 4.0 + PandaSet Dataset Terms" if fid.startswith("ps") else
                                      "CC BY-NC-SA 4.0 (research only)")
        rec = {"frame": name, "license": lic, "shippable": False, "set": "heldout",
               "reference": "tt_transfusion.reference torch fp32 (ONNX initializers, BN folded)",
               "counts": run.counts, "pillars": {k: v for k, v in p.stats().items() if k != "densify"},
               "ort_agreement": agree, "proposals": [int(v) for v in run.objects.extra["proposal"]],
               "output": run.output.to_dict()}
        det_path = out / f"{name}.detections.json"
        det_path.write_text(json.dumps(rec, indent=1))
        for f in (io_path, det_path):
            manifest["files"][f.name] = {"sha256": file_sha(f), "bytes": f.stat().st_size}
        manifest_path.write_text(json.dumps(manifest, indent=1))
        print(f"{name}: pillars {p.num_pillars} counts {run.counts} | ort pre-sel min pcc "
              f"{agree['pre_selection_min_pcc']:.10f}, top-k shared {agree['topk']['shared']}/{agree['topk']['k']} "
              f"strong {agree['topk']['strong']} covered {agree['topk']['strong_covered']}, teacher-forced min pcc "
              f"{agree['teacher_forced_min_pcc']:.10f}, detections recall/precision {m['recall']:.3f}/"
              f"{m['precision']:.3f} counts equal {m['counts_equal']} ({time.perf_counter() - t0:.1f} s)", flush=True)
    return 0


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