code/tt_bevformer/tests/test_reference_cpu.py

# SPDX-License-Identifier: Apache-2.0
"""The CPU fp32 reference (``tt_bevformer.reference``) against the oracle (PLAN.md 0.3 item 1).

Autoware's ``bevformer_small.onnx`` is not obtainable and ORT cannot run a graph with its TensorRT plugins, so the
oracle is the SPEC's (research/bevformer/SPEC.md section 6): ONNX Runtime on the DERIVED pure ONNX of the research
re-implementation of the deployed ``forward_trt`` graph (``reference.ort``), plus the research goldens. Gates:

- before the data-dependent selection: PCC >= 0.9999 on ``bev_embed``, ``outputs_classes`` and ``outputs_coords``
  (all six decoder layers) and on the research golden taps;
- after it (top-300 + threshold 0.2 + range filter): set-based detection agreement (same label, BEV centre within
  0.1 m, recall and precision >= 0.97).

Both forms of the reference are checked: the deployed graph (``forward``, geometry inside, BatchNorm folded where the
port folds it) and the port's decomposition (``forward_port``: host float64 geometry, rebatched SCA, rotation by
gather). The shipped PandaSet frames run end to end through the host path (manifest -> JPEG decode -> the node's
resize -> the C14 preset -> calibration -> TemporalState -> network -> post-processing).

Runs in the research venv (onnxruntime); skipped where onnxruntime, the weights or the research data are absent::

    cd bundles/bevformer-p150 && OMP_NUM_THREADS=4 PYTHONPATH=$PWD/code \\
        /home/ubuntu/experiments/tt-models/tools/research-venv/bin/python -m pytest -q -s -p no:cacheprovider \\
        code/tt_bevformer/tests/test_reference_cpu.py
"""
from __future__ import annotations

import json
import os

import numpy as np
import pytest

ort = pytest.importorskip("onnxruntime")
torch = pytest.importorskip("torch")

from tt_bevformer.host import postprocess as hpost  # noqa: E402
from tt_bevformer.host.inputs import load_frame  # noqa: E402
from tt_bevformer.reference import config as C  # noqa: E402
from tt_bevformer.tests import paths  # noqa: E402
from tt_bevformer.ttaw.metrics import match_detections, pcc  # noqa: E402

REF_PCC = 0.9999                     # PLAN.md 0.3 item 1
SET_MIN = 0.97                       # set-based agreement after the top-k selection
THREADS = int(os.environ.get("BEVFORMER_TEST_THREADS", "4"))   # workspace rule: <= 4 CPU threads
OUTS = ("bev_embed", "outputs_classes", "outputs_coords")
torch.set_num_threads(THREADS)


def _need(x, what):
    if x is None:
        pytest.skip(f"{what} not in this workspace")
    return x


@pytest.fixture(scope="module")
def reference():
    from tt_bevformer.reference.pipeline import BEVFormerReference

    return BEVFormerReference(_need(paths.weights_file(), "BEVFormer weights (BEVFORMER_WEIGHTS_DIR)"))


@pytest.fixture(scope="module")
def oracle():
    from tt_bevformer.reference.ort import OrtOracle

    return OrtOracle(_need(paths.derived_pure_onnx(), "research/bevformer/derived_onnx"), threads=THREADS)


@pytest.fixture(scope="module")
def sample0():
    golden = _need(paths.research_file("golden", "sample0_golden.npz"), "research/bevformer golden")
    with np.load(golden) as g:
        return {k: g[k] for k in g.files}


def _inputs(g, f):
    prev = g["f1_prev_bev"] if f == 1 else g[f"f{f - 1}_bev_embed"]
    return (g["image"], prev, g[f"f{f}_use_prev_bev"], g[f"f{f}_can_bus"], g["lidar2img"])


def _check(name, ours, ref, gate=REF_PCC):
    p = pcc(ours, ref)
    print(f"  {name:28s} pcc={p:.10f} max_abs={float(np.abs(np.asarray(ours) - np.asarray(ref)).max()):.3e}")
    assert p >= gate, (name, p)
    return p


def _set_agreement(ours: hpost.Detections, ref_boxes: list, name: str):
    """Same-label BEV-centre matching of the node's detections (score > 0.2) against a reference list."""
    rc = np.array([[d["x"], d["y"]] for d in ref_boxes], np.float64).reshape(-1, 2)
    rl = np.array([d["label"] for d in ref_boxes], np.int64)
    rs = np.array([d["score"] for d in ref_boxes], np.float64)
    m = match_detections(np.stack([ours.x, ours.y], 1), ours.label, ours.score, rc, rl, rs, max_dist=0.1)
    print(f"  {name:28s} dets ours={len(ours)} ref={len(ref_boxes)} recall={m.recall:.4f} "
          f"precision={m.precision:.4f}")
    assert m.recall >= SET_MIN and m.precision >= SET_MIN
    return m


@pytest.fixture(scope="module")
def sample0_runs(reference, sample0):
    """Each sample0 frame through the deployed form (frames 1-3) and the port form (frames 2-3), with taps on 1."""
    out = {}
    for f in (1, 2, 3):
        out[("graph", f)] = reference.run_frame(*_inputs(sample0, f), taps=(f == 1))
        if f > 1:
            out[("port", f)] = reference.run_frame(*_inputs(sample0, f), port=True)
    return out


# ---------------------------------------------------------------------------------------------- the oracle

@pytest.mark.parametrize("frame", [1, 2])
def test_reference_vs_onnxruntime(oracle, sample0, sample0_runs, frame):
    """sample0 frame 1 (reset, random prev_bev masked out) and frame 2 (history): both forms vs ORT."""
    ort_out = oracle(*_inputs(sample0, frame))
    print(f"\nsample0 frame {frame} vs ONNX Runtime (derived pure ONNX)")
    forms = ["graph"] + (["port"] if frame > 1 else [])
    for form in forms:
        r = sample0_runs[(form, frame)]
        for name in OUTS:
            _check(f"{form}.{name}", getattr(r, name), ort_out[name])
        ort_dets = hpost.decode(ort_out["outputs_classes"], ort_out["outputs_coords"])
        _set_agreement(r.detections, ort_dets.to_records(), f"{form}.detections")


@pytest.mark.parametrize("frame", [1, 2, 3])
def test_reference_vs_research_goldens(sample0, sample0_runs, frame):
    """The research run's outputs (frame 3: +2 m / +3 deg synthetic ego motion, so the shift and the rotation of
    prev_bev run) and frame-1 taps; the detections against the research detection lists (SPEC section 6: 49 / 37 /
    39 boxes)."""
    print(f"\nsample0 frame {frame} vs research golden")
    for form in ("graph", "port"):
        r = sample0_runs.get((form, frame))
        if r is None:
            continue
        for name in OUTS:
            _check(f"{form}.{name}", getattr(r, name), sample0[f"f{frame}_{name}"])
        dets = json.loads(_need(paths.research_file("golden", f"sample0_f{frame}_dets.json"), "dets").read_text())
        _set_agreement(r.detections, dets, f"{form}.detections")
        assert abs(len(r.detections) - {1: 49, 2: 37, 3: 39}[frame]) <= 1
    if frame == 1:
        taps = sample0_runs[("graph", 1)].taps
        for k in ("backbone.layer4", "neck", "ref_cam", "bev_mask_n", "encoder.layer0", "encoder.layer1",
                  "encoder.layer2", "decoder.layer0", "decoder.layer5"):
            gold = sample0[f"f1_tap_{k}"]
            _check(f"tap.{k}", taps[k].reshape(gold.shape), gold)


# ------------------------------------------------------------------------------- public data (end to end)

def test_shipped_pandaset_sequence_end_to_end(reference):
    """The shipped sample through the host path and the port form, as a stateful sequence (reset at frame 30),
    against the public-data workflow's CPU goldens (research reference, deployed form) and its detections."""
    paths.need_pandaset()
    pub = _need(paths.public_data(), "research/bevformer/public_data")
    gdir = pub / "pandaset" / "goldens" / "ps019_f30-40"
    exp = json.loads((_need(paths.shippable(), "shippable") / "expected_cpu_detections.json").read_text())
    reference.reset()
    print("\nPandaSet 019 f30-40 (shipped sample) end to end, port form")
    for i, f in enumerate(paths.SAMPLE_FRAMES):
        fr = load_frame(paths.PANDASET_SAMPLES / f"pandaset_019_f{f}.json")
        r = reference.step(fr, port=True, workers=4)
        with np.load(gdir / f"f{f}.npz") as g:
            assert np.array_equal(r.inputs["can_bus"], g["can_bus"]) and np.array_equal(r.inputs["lidar2img"],
                                                                                         g["lidar2img"])
            for name in OUTS:
                if name in g.files:
                    _check(f"f{f}.{name}", getattr(r, name), g[name])
        _set_agreement(r.detections, exp["frames"][i]["detections"], f"f{f}.detections")


def test_nuscenes_teacher_forced_taps(reference):
    """nuScenes mini_val scene-0103 kf01 (history, 4.3 m / -1.0 deg ego motion), teacher-forced with the golden
    kf00 bev_embed: the deployed form's taps and outputs against the public-data golden taps."""
    P = _need(paths.import_public_common(), "research/bevformer/public_data")
    pub = paths.public_data()
    x = P.load_frame_inputs(str(pub / "nuscenes" / "inputs" / "scene-0103" / "kf01.npz"))
    assert x["image_bgr_sha256_ok"]
    with np.load(pub / "nuscenes" / "goldens" / "scene-0103" / "kf00.npz") as g0:
        prev = g0["bev_embed"]
    r = reference.run_frame(x["image"], prev, x["use_prev_bev"], x["can_bus"], x["lidar2img"], taps=True)
    print("\nnuScenes scene-0103 kf01 (teacher-forced) vs public-data golden")
    with np.load(pub / "nuscenes" / "goldens" / "scene-0103" / "kf01.npz") as g:
        for name in OUTS:
            if name in g.files:
                _check(name, getattr(r, name), g[name])
        for k in [f[4:] for f in g.files if f.startswith("tap_")]:
            gold = g[f"tap_{k}"]
            _check(f"tap.{k}", r.taps[k].reshape(gold.shape), gold)