code/tt_transfusion/tests/test_api_host.py

# SPDX-License-Identifier: Apache-2.0
"""Host tests of the API wiring (no device): identity and runtime parameters, the host half of the hooks
(``_build_host``, ``_prepare``, ``_postprocess``) with the fp32 torch reference network in the device's place, against
the stored CPU reference of the shipped samples (compared with the container smoke's own comparator,
``ttaw.server.smoke``); the ``sampling`` input keyword; a frame Autoware skips; the server's ``sampling`` field.

    TT_VISIBLE_DEVICES=none python -m pytest -q -p fake_ttnn_plugin code/tt_transfusion/tests/test_api_host.py
"""
from __future__ import annotations

import base64
import json

import numpy as np
import pytest

from tt_transfusion import io as tio
from tt_transfusion.api import TransFusion
from tt_transfusion.host.postprocess import RUNTIME_PARAMS
from tt_transfusion.reference.config import NODE_PARAMS
from tt_transfusion.tests import _data
from tt_transfusion.ttaw.decode import AUTOWARE_LABELS
from tt_transfusion.ttaw.server import smoke


def test_identity_and_runtime_params():
    assert tuple(TransFusion.VARIANTS) == ("default",) and TransFusion.DEFAULT_VARIANT == "default"
    assert tuple(TransFusion.LABELS) == AUTOWARE_LABELS and TransFusion.LABELS[7] == "PEDESTRIAN"
    assert TransFusion.RUNTIME_PARAMS is RUNTIME_PARAMS and TransFusion.EXTRA_INPUTS == ("sampling",)
    assert TransFusion.DEFAULT_REVISION == "8f07eb93038bfd62e35d1f78e56a6b3fb76b12aa"
    p = TransFusion.validate_params({})
    assert p["score_threshold"] == 0.1 and p["remap_classes"] is True and p["max_detections"] is None
    for k in ("circle_nms_dist_threshold", "iou_nms_threshold", "iou_nms_search_distance_2d"):
        assert p[k] == NODE_PARAMS[k]                       # Autoware's defaults
    assert TransFusion.validate_params({"score_threshold": "0.5"})["score_threshold"] == 0.5
    for bad in ({"score_threshold": 1.5}, {"remap_classes": 1}, {"max_detections": 0}, {"iou_nms_threshold": 1.5},
                {"sampling": "clean"}):                     # sampling is an input keyword, not a runtime param
        with pytest.raises(tio.InputError):
            TransFusion.validate_params(bad)


def _host_model() -> TransFusion:
    m = TransFusion.__new__(TransFusion)                    # no device: the host half only
    m.weights_path = _data.weights_dir()
    m.variant = "default"
    m._build_host()
    return m


@pytest.fixture(scope="module")
def net():
    torch = pytest.importorskip("torch")
    pytest.importorskip("onnx")
    from tt_transfusion.reference.model import TransFusionTorch
    from tt_transfusion.reference.weights import load_weights

    torch.set_num_threads(2)
    return TransFusionTorch(load_weights(_data.weights_dir()), threads=2)


@pytest.mark.parametrize("sample", ["test_pcd", "pandaset_019_f40"])
def test_hooks_reproduce_the_stored_reference(net, sample):
    m = _host_model()
    prep = m._prepare(points=_data.sample_path(f"{sample}.npz"))
    raw = net.forward(prep.features, prep.coords, prep.num_pillars)
    out = m._postprocess(raw, prep, m.validate_params({}))
    body = out.to_dict()
    ref = json.loads(_data.sample_path(f"{sample}.reference.json").read_text())
    metrics, fails = smoke.compare_with_reference(body, ref, gates={"max_center_dist": 0.01, "min_recall": 1.0,
                                                                    "min_precision": 1.0, "max_score_err": 1e-4})
    assert fails == [], (fails, metrics)
    assert body["meta"]["counts"] == ref["meta"]["counts"] and body["meta"]["variant"] == "default"
    assert body["meta"]["num_pillars"] == ref["meta"]["num_pillars"]
    assert len(body["meta"]["orientation_availability"]) == body["num_detections"] == len(body["meta"]["proposal"])


def test_prepare_contract():
    m = _host_model()
    with pytest.raises(tio.InputError):
        m._prepare(points=None)
    pts = np.random.default_rng(0).uniform(-40, 40, (2000, 4)).astype(np.float32)
    with pytest.raises(tio.InputError, match="point cloud only"):
        m._prepare(points=pts, images=[{"camera": "front"}])
    with pytest.raises(tio.InputError, match="sampling"):
        m._prepare(points=pts, sampling="random")
    clean = m._prepare(points=pts)
    assert clean.sampling == {"mode": "clean", "offset": 0, "voxelizer_points": 2000}
    assert clean.skipped is None and clean.autoware_would_skip == "too_few_voxels"


def test_a_frame_autoware_skips_gives_no_objects():
    """``autoware_compat`` + fewer than 5,000 pillars: no network run (``_forward`` returns None without touching a
    device), an empty output carrying the reason."""
    m = _host_model()
    pts = np.random.default_rng(1).uniform(-40, 40, (2000, 4)).astype(np.float32)
    prep = m._prepare(points=pts, sampling={"mode": "autoware_compat", "offset": 7})
    assert prep.skipped == "too_few_voxels"
    raw = m._forward(prep)
    out = m._postprocess(raw, prep, m.validate_params({}))
    body = out.to_dict()
    assert raw is None and body["num_detections"] == 0 and body["meta"]["skipped"] == "too_few_voxels"
    assert body["meta"]["sampling"]["mode"] == "autoware_compat"


def test_server_accepts_the_sampling_field(monkeypatch):
    pytest.importorskip("fastapi")
    from fastapi.testclient import TestClient

    from tt_transfusion.server import app as server
    from tt_transfusion.tests.stubs import StubModel

    seen = {}

    class Recorder(StubModel):
        def __call__(self, points=None, **kw):
            seen.update(kw)
            return super().__call__(points, **kw)

    monkeypatch.setenv("TT_MESH_SHAPE", "1x1")
    monkeypatch.setattr(server.app.state.ttaw, "model_factory", Recorder)
    pts = np.random.default_rng(0).normal(size=(10, 4)).astype(np.float32)
    data = base64.b64encode(pts.tobytes()).decode()
    with TestClient(server.app) as c:
        r = c.post("/predict", json={"points": {"format": "bin", "data": data},
                                     "sampling": {"mode": "autoware_compat", "offset": 5}})
        assert r.status_code == 200, r.text
        assert seen["sampling"] == {"mode": "autoware_compat", "offset": 5}
        assert c.post("/predict", json={"points": {"format": "bin", "data": data}, "samplin": "x"}).status_code == 422


def test_a_frame_without_pillars_runs(net):
    """No point inside the range: zero pillars, an all-zero canvas, the network still runs (Autoware would skip the
    frame: reported as ``autoware_would_skip``) and the post-processing returns a valid, possibly empty, output."""
    m = _host_model()
    pts = np.array([[200.0, 0.0, 0.0, 10.0], [0.0, 0.0, 9.0, 10.0]], np.float32)
    prep = m._prepare(points=pts)
    assert prep.num_pillars == 0 and prep.autoware_would_skip == "too_few_voxels"
    assert (prep.canvas_index == prep.features.shape[0]).all() and not prep.features.any()
    raw = net.forward(prep.features, prep.coords, prep.num_pillars)
    assert raw["cls_score0"].shape == (1, 5, 500) and np.isfinite(raw["bbox_pred0"]).all()
    body = m._postprocess(raw, prep, m.validate_params({})).to_dict()
    assert body["meta"]["num_pillars"] == 0 and body["num_detections"] == len(body["meta"]["proposal"])