code/scripts/make_goldens.py

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

    cd bundles/bevfusion-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/bevfusion/hf_bevfusion \\
        --research /home/ubuntu/experiments/tt-models/research/bevfusion \\
        --out /home/ubuntu/experiments/tt-models/research/bevfusion/goldens

For every development frame the fp32 torch reference (``reference/pipeline.py``: host pre-processing, the deployed
network rebuilt from the ONNX initializers, host post-processing) runs once with its module taps; the hybrid
reference (the same torch sparse encoder, then the deployed ONNX's plugin-free dense sub-graph under ONNX Runtime,
with its intermediates exposed) runs on the same tensors; the torch reference runs once more teacher-forced with
ORT's top-k indices; the research reference's independent sparse encoder (``research/bevfusion/scripts/
bevfusion_ref.py SparseEncoderRef``: scatter form, BN unfolded) runs on the same voxel features. The agreement (sparse
BEV PCC vs the research encoder, PCC per dense 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 (``coors``, ``num_points_per_voxel``; ``voxels`` are regenerated
  bit-exactly by ``host.preprocess.hard_voxelize``), the voxel features (``vfe``), the per-level active counts, the
  reference outputs (``bbox_pred``, ``score``, ``label_pred``), the top-500 (``topk_indices`` int64,
  ``topk_values``) and the hybrid (ORT dense) outputs (``ort_bbox_pred``, ...);
- ``<out>/<frame>.detections.json``: the reference output (``Detections3D.to_dict``), per-stage counts, voxel and
  sparse statistics, the agreement record;
- ``<out>/<frame>.taps.npz`` (``TAP_FRAMES`` only; ~100-200 MB each): every tap of ``reference.model.MODULE_TAPS``
  plus ``neck`` and ``decoder.key_pos_embed`` (float32; ``head.topk_indices`` / ``label_pred`` 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; the PandaSet sample's files go to the git-ignored ``staging_samples_pandaset/``):
  ``tests/goldens/voxels.json`` (voxel, VFE
  and rulebook hashes of the shipped samples), ``tests/goldens/detections/<sample>.json`` and the smoke reference
  ``samples/<stem>.reference.json``.

Frames: the shipped samples (``pandaset_019_f40`` = public ps019_f40, ``test_pcd``), the Autoware sample-rosbag top
scan #144 (single sweep, ground-levelled, the SPEC's frame; license unstated: research only) and the public frames
ps090_f40 (PandaSet), ns0103_k20, ns0916_k10 and ns0103_k00 (nuScenes mini_val, single-sweep first key-frame; CC
BY-NC-SA: research only). The held-out frames of the device e2e test (``tests/_data.py HELDOUT``) use the research
public-data goldens (the research reference, SPEC 6.2) and are never used for any choice.
"""
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_bevfusion import __version__  # noqa: E402
from tt_bevfusion import io as tio  # noqa: E402
from tt_bevfusion.reference.model import MODULE_TAPS, SELECTION_TAPS, SPARSE_TAPS  # noqa: E402
from tt_bevfusion.reference.ort import DENSE_ORT_TAPS  # noqa: E402
from tt_bevfusion.reference.pipeline import BEVFusionReference  # noqa: E402
from tt_bevfusion.ttaw import __version__ as ttaw_version  # noqa: E402
from tt_bevfusion.ttaw.golden import TapRegistry, save_goldens  # noqa: E402
from tt_bevfusion.ttaw.metrics import match_detections, pcc  # noqa: E402

SAMPLES = CODE / "tt_bevfusion" / "samples"
# The PandaSet sample and its small goldens live in the git-ignored staging directory until PandaSet samples may ship
# (research/DATASETS.md N5; staging_samples_pandaset/README.md); test.pcd ships.
STAGING = CODE.parent / "staging_samples_pandaset"
BUNDLE_GOLDENS = CODE / "tt_bevfusion" / "tests" / "goldens"
SHIPPED = ("pandaset_019_f40", "test_pcd")
PUBLIC = ("ps090_f40", "ns0103_k20", "ns0916_k10", "ns0103_k00")
TAP_FRAMES = ("pandaset_019_f40", "rosbag_0144", "public_ns0103_k20")
TAP_SET = tuple(MODULE_TAPS) + ("neck", "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 load_public(path: Path) -> tio.PointCloud:
    with np.load(path, allow_pickle=False) as z:
        pts = z["points"]
    return tio.PointCloud(pts, ("x", "y", "z", "time_lag"))


def frames(research: Path):
    """(name, source, 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, (STAGING if name.startswith("pandaset") else SAMPLES) / f"{name}.npz", lic, True
    rosbag = research / "sample_inputs" / "rosbag_top144_concat_leveled.npz"
    if rosbag.is_file():
        with np.load(rosbag, allow_pickle=False) as z:
            pts = z["points"][:, :3]
        yield "rosbag_0144", tio.PointCloud(pts, ("x", "y", "z")), \
            "Autoware sample rosbag (license unstated): research only", False
    for name in PUBLIC:
        path = research / "public_data" / "inputs" / f"{name}.npz"
        if not path.is_file():
            print(f"skip {name}: {path} missing")
            continue
        lic = "CC BY 4.0 + PandaSet Dataset Terms" if name.startswith("ps") else "CC BY-NC-SA 4.0 (research only)"
        yield f"public_{name}", load_public(path), lic, False


def research_bev(research: Path, weights: Path, vfe: np.ndarray, coors: np.ndarray) -> np.ndarray:
    scripts = str(research / "scripts")
    if scripts not in sys.path:
        sys.path.insert(0, scripts)
    import bevfusion_ref as R  # noqa: E402  (workspace research code, not a download)

    enc = R.SparseEncoderRef(R.load_initializers(str(weights / "bevfusion_lidar.onnx")))
    return enc.forward(vfe, coors)


def agreement(rt: BEVFusionReference, rh: BEVFusionReference, run, run_h, tt: TapRegistry, th: TapRegistry,
              bev_research: np.ndarray) -> dict:
    rec: dict = {"sparse_bev_pcc_vs_research": pcc(tt["sparse.bev"], bev_research),
                 "sparse_bev_max_abs_vs_research": float(np.abs(tt["sparse.bev"] - bev_research).max())}
    pre = {}
    for name in DENSE_ORT_TAPS:
        if name in SELECTION_TAPS or name not in th.names() or name not in tt.names():
            continue
        pre[name] = pcc(tt[name], th[name].reshape(tt[name].shape))
    rec["dense_pre_selection_pcc_min"] = min(pre.values())
    rec["dense_pre_selection_pcc"] = pre
    mine, ref = tt["head.topk_indices"].ravel(), th["head.topk_indices"].ravel()
    strong = ref[th["head.topk_values"].ravel() >= STRONG]
    rec["topk"] = {"k": int(ref.size), "shared": len(set(mine.tolist()) & set(ref.tolist())),
                   "positional": float(np.mean(mine == ref)), "strong": int(len(strong)),
                   "strong_kept": int(len(set(strong.tolist()) & set(mine.tolist())))}
    tf = TapRegistry(include=SELECTION_TAPS)
    rt.network(run.prepared, tf, topk_indices=th["head.topk_indices"])
    post = {}
    for name in SELECTION_TAPS:
        if name in ("head.topk_indices", "label_pred"):
            continue
        post[name] = pcc(tf[name], th[name].reshape(tf[name].shape))
    rec["teacher_forced_pcc_min"] = min(post.values())
    rec["teacher_forced_pcc"] = post
    rec["teacher_forced_labels_equal"] = bool(np.array_equal(tf["label_pred"], th["label_pred"]))
    da, db = run.output, run_h.output
    m = match_detections(da.boxes, da.label_ids, da.scores, db.boxes, db.label_ids, db.scores, max_dist=0.01)
    rec["detections"] = {"reference": len(da), "hybrid": len(db), **m.to_dict()}
    rec["counts_equal"] = run.counts == run_h.counts
    return rec


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("--threads", type=int, default=4)
    ap.add_argument("--frames", default=None, help="comma-separated subset of frame names")
    a = ap.parse_args()
    import onnxruntime
    import torch

    torch.set_num_threads(a.threads)
    a.out.mkdir(parents=True, exist_ok=True)
    (BUNDLE_GOLDENS / "detections").mkdir(parents=True, exist_ok=True)
    rt = BEVFusionReference(a.weights, backend="torch", threads=a.threads)
    rh = BEVFusionReference(a.weights, backend="hybrid", threads=a.threads, intermediates=True)
    manifest = {"bundle_version": __version__, "ttaw": ttaw_version, "torch": torch.__version__,
                "onnxruntime": onnxruntime.__version__, "numpy": np.__version__, "python": platform.python_version(),
                "onnx_sha256": rt.weights.sha256, "frames": {}}
    voxels = {}
    want = set(a.frames.split(",")) if a.frames else None
    for name, src, lic, shippable in frames(a.research):
        if want and name not in want:
            continue
        t0 = time.time()
        tt = TapRegistry(include=TAP_SET)
        run = rt.run(src, taps=tt)
        th = TapRegistry()
        run_h = rh.run(src, taps=th)
        prep = run.prepared
        assert np.array_equal(prep.vfe, run_h.prepared.vfe)
        bev_r = research_bev(a.research, a.weights, prep.vfe, prep.coors)
        rec = agreement(rt, rh, run, run_h, tt, th, bev_r)
        plan = prep.plan
        io = {"coors": prep.coors, "num_points_per_voxel": prep.voxel_set.num_points, "vfe": prep.vfe,
              "level_counts": np.array([plan.levels[k].num_active for k in ("l1", "l2", "l3", "l4", "out")]),
              "topk_indices": tt["head.topk_indices"], "topk_values": tt["head.topk_values"],
              **{k: v for k, v in run.outputs.items()}, **{f"ort_{k}": v for k, v in run_h.outputs.items()}}
        io_path = a.out / f"{name}.io.npz"
        save_goldens(io_path, io, meta={"frame": name, "license": lic, "stats": prep.stats()}, compress=True)
        body = run.output.to_dict()
        det = {"frame": name, "license": lic, "output": body, "counts": run.counts, "stats": prep.stats(),
               "agreement": rec}
        (a.out / f"{name}.detections.json").write_text(json.dumps(det, indent=1, default=float) + "\n")
        files = [io_path.name, f"{name}.detections.json"]
        if name in TAP_FRAMES:
            tp = a.out / f"{name}.taps.npz"
            save_goldens(tp, tt.to_dict(), meta={"frame": name, "license": lic, "taps": sorted(tt.names())},
                         compress=True)
            files.append(tp.name)
        manifest["frames"][name] = {"license": lic, "files": {f: file_sha(a.out / f) for f in files},
                                    "num_voxels": prep.num_voxels, "levels": plan.counts(), "bucket": plan.bucket,
                                    "detections": len(run.output), "agreement_summary": {
                                        "sparse_bev_pcc_vs_research": rec["sparse_bev_pcc_vs_research"],
                                        "dense_pre_selection_pcc_min": rec["dense_pre_selection_pcc_min"],
                                        "topk": rec["topk"], "teacher_forced_pcc_min": rec["teacher_forced_pcc_min"],
                                        "detections_recall": rec["detections"]["recall"],
                                        "detections_precision": rec["detections"]["precision"]}}
        if shippable:
            staged = name.startswith("pandaset")
            det_dir = (STAGING / "goldens" if staged else BUNDLE_GOLDENS) / "detections"
            det_dir.mkdir(parents=True, exist_ok=True)
            (det_dir / f"{name}.json").write_text(json.dumps(det, indent=1, default=float) + "\n")
            ((STAGING if staged else SAMPLES) / f"{name}.reference.json").write_text(json.dumps(body, indent=1) + "\n")
            voxels[name] = {"input_sha256": file_sha(src), "num_voxels": prep.num_voxels,
                            "coors_sha256": sha(prep.coors), "num_points_sha256": sha(prep.voxel_set.num_points),
                            "voxels_sha256": sha(prep.voxel_set.voxels), "vfe_sha256": sha(prep.vfe),
                            "levels": plan.counts(), "pairs": plan.pairs(), "bucket": plan.bucket,
                            "maps_sha256": {k: sha(v) for k, v in plan.maps.items()},
                            "stats": prep.voxel_set.stats}
        print(f"{name}: {prep.num_voxels} voxels, levels {plan.counts()}, bucket {plan.bucket}, "
              f"{len(run.output)} det, bev vs research {rec['sparse_bev_pcc_vs_research']:.10f}, "
              f"dense min {rec['dense_pre_selection_pcc_min']:.10f}, topk {rec['topk']}, "
              f"teacher min {rec['teacher_forced_pcc_min']:.10f}, det recall {rec['detections']['recall']} "
              f"precision {rec['detections']['precision']}, {time.time() - t0:.1f} s", flush=True)
    (a.out / "manifest.json").write_text(json.dumps(manifest, indent=1, default=float) + "\n")
    if voxels:
        (BUNDLE_GOLDENS / "voxels.json").write_text(json.dumps(voxels, indent=1) + "\n")
    return 0


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