code/tt_transfusion/tests/test_reference_structure_host.py

# SPDX-License-Identifier: Apache-2.0
"""The CPU reference reads the weights the way the spec describes them (S:transfusion:57-208): the Autoware
parameters, and the module structure checked from the ONNX graph (layer counts, strides, channels, the decoder's
heads / scale / LayerNorm eps, the in-graph decoration constants, ``bev_pos``, parameter counts after the BN folds),
without onnxruntime or torch.

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

import numpy as np
import pytest

pytest.importorskip("onnx")
pytest.importorskip("yaml")

from tt_transfusion.reference.config import (N_MAX, NODE_PARAMS, OUT_SIZE_FACTOR, POINTS_PER_VOXEL,  # noqa: E402
                                             TransFusionConfig)
from tt_transfusion.reference.weights import PREDICTION_HEADS, load_weights  # noqa: E402
from tt_transfusion.tests import _data  # noqa: E402


@pytest.fixture(scope="module")
def config():
    return TransFusionConfig.from_weights(_data.weights_dir())


@pytest.fixture(scope="module")
def weights(config):
    return load_weights(_data.weights_dir(), config)


def test_autoware_parameters(config):
    c = config
    assert c.version == "v2.1" and c.class_names == ("CAR", "TRUCK", "BUS", "BICYCLE", "PEDESTRIAN")
    assert c.voxels_num == (5000, 30000, 60000) and (c.min_voxels, c.max_voxels) == (5000, 60000)
    assert c.point_cloud_range == (-92.16, -92.16, -3.0, 92.16, 92.16, 7.0) and c.voxel_size == (0.24, 0.24, 10.0)
    assert (c.grid_x, c.grid_y, c.grid_z) == (768, 768, 1)                      # float32 184.32 / 0.24 = 768.00006
    assert (c.feature_x, c.feature_y) == (192, 192) and OUT_SIZE_FACTOR == 4
    assert c.num_proposals == 500 and POINTS_PER_VOXEL == 20 and N_MAX == 65536
    spec = c.grid_spec
    assert spec.order == "raster" and spec.max_pillars == N_MAX and spec.max_points_per_pillar == 20
    assert c.score_threshold == NODE_PARAMS["score_threshold"] == 0.1
    assert c.yaw_norm_thresholds.tolist() == pytest.approx([0.3, 0.3, 0.3, 0.3, 0.0])
    assert c.with_node(yaw_norm_thresholds=(1.5, -1, 0.3, 0.999, 1.0)).yaw_norm_thresholds.tolist() == pytest.approx(
        [0.0, 0.0, 0.3, 0.999, 0.0])                                             # transfusion_config.hpp:79-81
    assert TransFusionConfig.sanitized_score_threshold(0.0) == 0.2 and TransFusionConfig.sanitized_score_threshold(-1) == 0.2
    assert TransFusionConfig.sanitized_circle_nms(0.0) == 0.5 and TransFusionConfig.sanitized_circle_nms(1.2) == 1.2
    rm = c.remapper
    assert rm.num_labels == 8 and rm.allow[1, 2] and rm.allow[1, 4] and rm.allow[2, 4] and rm.allow[3, 4]
    assert int(rm.allow.sum()) == 4
    assert (rm.min_area[1, 2], rm.max_area[1, 2], rm.min_area[3, 4], rm.max_area[3, 4]) == (12.1, 36.0, 36.0, 999.999)
    assert c.decoration_scale_zyx.tolist() == [np.float32(10.0), np.float32(0.24), np.float32(0.24)]
    assert c.decoration_offset_zyx.tolist() == [np.float32(2.0), np.float32(-92.04), np.float32(-92.04)]


def test_refuses_another_release(tmp_path, config):
    import shutil

    for f in ("transfusion_ml_package.param.yaml", "detection_class_remapper.param.yaml"):
        shutil.copy(_data.weights_dir() / f, tmp_path / f)
    (tmp_path / "deploy_metadata.yaml").write_text("version: v2\n")
    with pytest.raises(ValueError, match="not supported"):
        TransFusionConfig.from_model_dir(tmp_path)


def test_graph_structure_and_parameter_counts(weights):
    w = weights
    assert [len(b) for b in w.blocks] == [4, 6, 6]
    assert [(b[0].in_channels, b[-1].out_channels, b[0].stride) for b in w.blocks] == [(64, 64, 2), (64, 128, 2),
                                                                                       (128, 256, 2)]
    assert all(layer.kernel == 3 and layer.padding == 1 and layer.relu for b in w.blocks for layer in b)
    assert [(d.kind, d.in_channels, d.out_channels, d.kernel, d.stride, d.padding) for d in w.deblocks] == [
        ("conv", 64, 128, 2, 2, 0), ("conv", 128, 128, 1, 1, 0), ("conv_transpose", 256, 128, 2, 2, 0)]
    assert [d.bn_node is not None for d in w.deblocks] == [False, False, True]
    h = w.head
    assert (h.shared.in_channels, h.shared.out_channels, h.shared.kernel, h.shared.relu) == (384, 128, 3, False)
    assert (h.heatmap_hidden.out_channels, h.heatmap_hidden.relu, h.heatmap_out.out_channels) == (128, True, 5)
    assert h.nms_kernel == 3 and h.num_proposals == 500
    d = h.decoder
    for a in (d.self_attn, d.cross_attn):
        assert (a.num_heads, a.head_dim, a.q_divisor, a.scale) == (8, 16, 4.0, 0.25)
        assert a.wq.shape == a.wk.shape == a.wv.shape == a.wo.shape == (128, 128)
    assert [eps for _, _, eps in d.norms] == [pytest.approx(1e-5)] * 3
    assert d.ffn_w1.shape == (128, 256) and d.ffn_w2.shape == (256, 128)
    assert [(n, fin.out_channels) for n, _, fin in h.prediction_heads] == list(PREDICTION_HEADS)
    assert w.pfn.w0.shape == (11, 32) and w.pfn.w1.shape == (64, 64) and w.grid == (768, 768)
    # 5,270,580 learned ONNX parameters (S:transfusion:149); the folds turn the PFN BNs (32 + 64 channels x 4) into
    # 96 biases and the deblock2 BN (128 x 4) into 128 biases
    assert w.num_params() == 5_270_580 - (4 * 96 - 96) - (4 * 128 - 128)


def test_constants_and_tables(weights):
    h = weights.head
    k = np.arange(192 * 192)
    assert h.bev_pos.shape == (36864, 2)
    assert np.array_equal(h.bev_pos[:, 0], (k % 192).astype(np.float32) + 0.5)       # k = y * 192 + x
    assert np.array_equal(h.bev_pos[:, 1], (k // 192).astype(np.float32) + 0.5)
    assert h.bev_pos[1].tolist() == [1.5, 0.5] and h.bev_pos[192].tolist() == [0.5, 1.5]
    t = h.class_table
    assert t.shape == (5, 128) and np.array_equal(t[3], (h.class_encoding_w[:, 3] + h.class_encoding_b).astype(
        np.float32))
    p = weights.pfn
    assert np.array_equal(p.scale_zyx, np.array([10.0, 0.24, 0.24], np.float32))
    assert np.array_equal(p.offset_zyx, np.array([2.0, -92.04, -92.04], np.float32))


def test_onnx_tap_map_covers_the_reference_taps(weights):
    from tt_transfusion.reference.model import EXTRA_TAPS, MODULE_TAPS

    missing = [t for t in MODULE_TAPS + EXTRA_TAPS if t not in weights.onnx_taps]
    assert not missing, missing
    assert weights.onnx_taps["head.topk_indices"].endswith("TopK_output_1")
    assert weights.onnx_taps["pfn.input"] == "/pts_voxel_encoder/Mul_1_output_0"