code/tt_sceneseg/tests/test_host_preprocess.py

# SPDX-License-Identifier: Apache-2.0
"""Host pre-processing (``tt_sceneseg.host.preprocess``) against the deployed C++ semantics and the research
references (no device, no weights).

    TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_sceneseg/tests/test_host_preprocess.py

Oracles: OpenCV itself running the C++ backend's calls (``cv2.imread`` / ``cv2.resize`` / ``cv2.subtract`` /
``cv2.divide`` with the BGR-ordered Scalars: tt-env has cv2 4.8.1, the research venv 4.11.0); the research
public-data device inputs (``input_u8_bgr_640x320.png``); the SPEC reference (``research/sceneseg/scripts/
sceneseg_ref.py``: float32 division, so within one ulp) and the SPEC's stored golden inputs
(``ref_runs/highway_normal_{1,2}/cpp_{bgr,rgb}_input_f32.npy``: ``x / 255``, within 1e-6).
"""
from __future__ import annotations

import numpy as np
import pytest

from tt_sceneseg.host.preprocess import PreparedImage, check_channel_order, prepare_bgr, prepare_image, rgb_to_bgr
from tt_sceneseg.reference import config
from tt_sceneseg.tests.research import PUBLIC, SAMPLES, SPEC_REF_RUNS, TUTORIAL, research_module

F32 = np.float32


def _cv2():
    cv2 = pytest.importorskip("cv2")
    cv2.setNumThreads(2)
    return cv2


def _cpp_backend(cv2, bgr, order):
    """onnx_runtime_backend.cpp:41-60 with OpenCV's own calls (the rgb order: the same on cvtColor(BGR2RGB))."""
    img = bgr if order == "bgr" else cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
    mean, std = ((0.406, 0.456, 0.485), (0.225, 0.224, 0.229)) if order == "bgr" else \
        ((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
    r = cv2.resize(img, (640, 320))
    f = r.astype(F32) * F32(1.0 / 255.0)                       # convertTo(CV_32FC3, 1/255): one float32 product
    f = cv2.subtract(f, (*mean, 0.0))
    f = cv2.divide(f, (*std, 0.0))
    return r, np.ascontiguousarray(f.transpose(2, 0, 1)[None])


@pytest.mark.parametrize("name", TUTORIAL)
def test_samples_equal_the_cpp_backend(name):
    """The shipped samples (RGBA PNGs: ``IMREAD_COLOR`` drops the alpha) through the API decoding: the uint8 squash
    and the network input of both orders equal OpenCV running the C++ backend, bit for bit."""
    cv2 = _cv2()
    path = SAMPLES / f"{name}.png"
    bgr = cv2.imread(str(path), cv2.IMREAD_COLOR)
    for order in config.CHANNEL_ORDERS:
        prep = prepare_image(str(path), channel_order=order)
        r, x = _cpp_backend(cv2, bgr, order)
        assert prep.src_hw == bgr.shape[:2] and prep.channel_order == order
        np.testing.assert_array_equal(prep.u8_bgr, cv2.resize(bgr, (640, 320)))
        np.testing.assert_array_equal(prep.u8_bgr if order == "bgr" else prep.u8_bgr[:, :, ::-1], r)
        np.testing.assert_array_equal(prep.network_input(), x)


def test_random_sizes_equal_the_cpp_backend():
    cv2 = _cv2()
    rng = np.random.default_rng(7)
    for hw in ((1080, 1920), (900, 1600), (375, 1242), (640, 1280), (320, 640), (77, 99), (2000, 600)):
        bgr = rng.integers(0, 256, (*hw, 3), dtype=np.uint8)
        for order in config.CHANNEL_ORDERS:
            prep = prepare_bgr(bgr, channel_order=order)
            _, x = _cpp_backend(cv2, bgr, order)
            np.testing.assert_array_equal(prep.network_input(), x, err_msg=f"{hw} {order}")


def test_channel_orders_are_a_permutation_and_device_input():
    rng = np.random.default_rng(3)
    bgr = rng.integers(0, 256, (414, 727, 3), dtype=np.uint8)
    pb, pr = prepare_bgr(bgr, channel_order="bgr"), prepare_bgr(bgr, channel_order="rgb")
    np.testing.assert_array_equal(pb.u8_bgr, pr.u8_bgr)                    # one device input for both variants
    np.testing.assert_array_equal(pr.network_input()[0], pb.network_input()[0, ::-1])
    d = pb.device_input()
    assert d.shape == (1, 1, 320 * 640, 3) and d.dtype == np.uint8 and d.flags["C_CONTIGUOUS"]
    np.testing.assert_array_equal(d.reshape(320, 640, 3), pb.u8_bgr)
    np.testing.assert_array_equal(rgb_to_bgr(bgr[:, :, ::-1]), bgr)


def test_input_errors():
    with pytest.raises(ValueError):
        check_channel_order("BGR")
    with pytest.raises(ValueError):
        prepare_bgr(np.zeros((10, 10, 3), np.float32))
    with pytest.raises(ValueError):
        prepare_bgr(np.zeros((10, 10), np.uint8))
    with pytest.raises(ValueError):
        rgb_to_bgr(np.zeros((10, 10, 4), np.uint8))
    assert isinstance(prepare_bgr(np.zeros((1, 1, 3), np.uint8)), PreparedImage)


def test_public_device_inputs():
    """research/sceneseg/public_data: the converter's ``input_u8_bgr_640x320.png`` of all 23 frames (cv2 of the
    decoded source) bit for bit; its ``input_f32.npz`` (float32 division) within one ulp."""
    from PIL import Image

    from tt_sceneseg.tests.research import public_ids

    for fid in public_ids():
        d = PUBLIC / "inputs" / fid
        with Image.open(d / "source_bgr.png") as im:                         # cv2.imwrite: true colours
            bgr = np.ascontiguousarray(np.asarray(im.convert("RGB"))[:, :, ::-1])
        with Image.open(d / "input_u8_bgr_640x320.png") as im:
            want = np.ascontiguousarray(np.asarray(im.convert("RGB"))[:, :, ::-1])
        prep = prepare_bgr(bgr)
        np.testing.assert_array_equal(prep.u8_bgr, want, err_msg=fid)
        with np.load(d / "input_f32.npz") as z:
            for order in config.CHANNEL_ORDERS:
                x = prepare_bgr(bgr, channel_order=order).network_input()
                ulp = np.abs(x.view(np.int32).astype(np.int64) - z[order].view(np.int32).astype(np.int64))
                assert ulp.max() <= 1, (fid, order)


def test_spec_reference_and_spec_goldens():
    """The SPEC section 6.2 reference functions (cv2 + float32 normalisation: one ulp) and the SPEC's stored golden
    inputs of highway_normal_1 / 2 (``x / 255``: within 1e-6)."""
    ref = research_module("sceneseg_ref")
    cv2 = _cv2()
    for name in ("highway_normal_1", "highway_normal_2"):
        bgr = cv2.imread(str(SAMPLES / f"{name}.png"), cv2.IMREAD_COLOR)
        for order, fn, tag in (("bgr", ref.preprocess_cpp_bgr, "cpp_bgr"), ("rgb", ref.preprocess_cpp_rgb, "cpp_rgb")):
            x = prepare_image(str(SAMPLES / f"{name}.png"), channel_order=order).network_input()
            want = fn(bgr)
            ulp = np.abs(x.view(np.int32).astype(np.int64) - want.view(np.int32).astype(np.int64))
            assert ulp.max() <= 1, (name, order)
            stored = SPEC_REF_RUNS / name / f"{tag}_input_f32.npy"
            if stored.is_file():
                assert np.abs(x - np.load(stored)).max() <= 1e-6, (name, order)