| |
| """Host post-processing (``tt_sceneseg.host.postprocess`` / ``host.output``) against the deployed node's rule and the |
| research goldens (no device, no weights). |
| |
| TT_VISIBLE_DEVICES=none python -m pytest -q code/tt_sceneseg/tests/test_host_postprocess.py |
| |
| Oracles: a literal transliteration of ``run_model_node.cpp:148-188`` (here) and of the research's independent one |
| (``research/sceneseg/scripts/verify_ort.py`` ``post_cpp``, when onnxruntime imports); ``numpy.argmax`` + ``cv2.resize`` |
| (INTER_NEAREST); the SPEC's deployed-path goldens (``ref_runs/<image>/cpp_{bgr,rgb}_mask_srcres.png`` from the stored |
| ORT logits, bit for bit) and the public-data goldens (``golden/<id>/ort_fp32.npz`` -> ``cls_*`` and ``mask_*.png``, |
| bit for bit). |
| """ |
| from __future__ import annotations |
|
|
| import base64 |
| import io |
| import json |
|
|
| import numpy as np |
| import pytest |
|
|
| from tt_sceneseg.host.output import SceneSegMask, to_output |
| from tt_sceneseg.host.postprocess import (class_map_from_logits, class_map_from_margins, foreground_mask, |
| margins_from_logits, postprocess, to_source) |
| from tt_sceneseg.reference import config |
| from tt_sceneseg.tests.research import PUBLIC, SPEC_REF_RUNS, TUTORIAL, public_ids, research_module |
|
|
|
|
| def _scalar_node(logits: np.ndarray) -> np.ndarray: |
| """run_model_node.cpp:148-163, one pixel at a time (small arrays only).""" |
| c, h, w = logits[0].shape |
| out = np.zeros((h, w), np.uint8) |
| for y in range(h): |
| for x in range(w): |
| best, mx = 0, np.float32(-1e9) |
| for k in range(c): |
| s = np.float32(logits[0, k, y, x]) |
| if s > mx: |
| mx, best = s, k |
| out[y, x] = best |
| return out |
|
|
|
|
| def _logits(rng, shape=(1, 3, 32, 48), ties=True): |
| lg = rng.normal(size=shape).astype(np.float32) * 4 |
| if ties: |
| lg[0, 1, :4] = lg[0, 0, :4] |
| lg[0, 2, 4:8] = lg[0, 1, 4:8] |
| lg[0, 2, 8:12] = lg[0, 0, 8:12] |
| lg[0, :, 12:14] = 1.5 |
| lg[0, 1, 14:16] = np.maximum(lg[0, 0, 14:16], lg[0, 2, 14:16]) |
| return lg |
|
|
|
|
| def test_class_decision_matches_the_node_and_argmax(): |
| rng = np.random.default_rng(0) |
| for _ in range(3): |
| lg = _logits(rng) |
| want = _scalar_node(lg) |
| got = class_map_from_logits(lg) |
| np.testing.assert_array_equal(got, want) |
| np.testing.assert_array_equal(got, np.argmax(lg[0], axis=0)) |
| np.testing.assert_array_equal(class_map_from_margins(*margins_from_logits(lg)), want) |
| m01, m12, _ = margins_from_logits(lg) |
| np.testing.assert_array_equal(class_map_from_margins(m01, m12) == 1, want == 1) |
| assert got.dtype == np.uint8 and set(np.unique(got)) <= {0, 1, 2} |
| |
| lg = np.full((1, 3, 2, 2), -2e9, np.float32) |
| lg[0, 2, 0, 0] = -1.5e9 |
| assert class_map_from_logits(lg)[0, 0] == 0 == _scalar_node(lg)[0, 0] |
|
|
|
|
| def test_research_post_cpp(): |
| """research verify_ort.post_cpp (the verification pass's independent C++ re-implementation).""" |
| pytest.importorskip("onnxruntime") |
| cv2 = pytest.importorskip("cv2") |
| ver = research_module("verify_ort") |
| rng = np.random.default_rng(1) |
| for hw in ((414, 727), (1080, 1920), (100, 150)): |
| lg = _logits(rng, (1, 3, 320, 640)) |
| best, mask = ver.post_cpp(lg, hw) |
| res = postprocess(hw, logits=lg) |
| np.testing.assert_array_equal(res.class_map, best) |
| np.testing.assert_array_equal(res.mask, mask) |
| np.testing.assert_array_equal(res.class_map_src, cv2.resize(best, (hw[1], hw[0]), |
| interpolation=cv2.INTER_NEAREST)) |
|
|
|
|
| def test_postprocess_equals_opencv_nearest(): |
| cv2 = pytest.importorskip("cv2") |
| rng = np.random.default_rng(2) |
| lg = _logits(rng, (1, 3, 320, 640)) |
| cls = np.argmax(lg[0], axis=0).astype(np.uint8) |
| for hw in ((414, 727), (1080, 1920), (331, 681), (320, 640), (2160, 3840), (17, 29)): |
| res = postprocess(hw, logits=lg) |
| want = cv2.resize(np.where(cls == 1, 255, 0).astype(np.uint8), (hw[1], hw[0]), interpolation=cv2.INTER_NEAREST) |
| np.testing.assert_array_equal(res.mask, want) |
| assert res.mask.dtype == np.uint8 and res.mask.shape == hw and set(np.unique(res.mask)) <= {0, 255} |
| np.testing.assert_array_equal(res.mask == 255, res.class_map_src == 1) |
| np.testing.assert_array_equal(postprocess((414, 727), class_map=cls).mask, postprocess((414, 727), logits=lg).mask) |
| with pytest.raises(ValueError): |
| postprocess((10, 10)) |
| with pytest.raises(ValueError): |
| postprocess((10, 10), class_map=cls[:100]) |
| with pytest.raises(ValueError): |
| class_map_from_logits(np.zeros((2, 3, 4, 4), np.float32)) |
| np.testing.assert_array_equal(foreground_mask(np.array([[0, 1, 2]])), [[0, 255, 0]]) |
| np.testing.assert_array_equal(to_source(np.arange(6, dtype=np.uint8).reshape(2, 3), (2, 3)), |
| np.arange(6).reshape(2, 3)) |
|
|
|
|
| @pytest.mark.parametrize("name", ["highway_normal_1", "highway_normal_2"]) |
| def test_spec_deployed_masks(name): |
| """The SPEC's deployed-path goldens: the stored ORT logits of both orders -> the stored source-size masks.""" |
| from PIL import Image |
|
|
| d = SPEC_REF_RUNS / name |
| if not (d / "cpp_bgr_logits_ort_f32.npy").is_file(): |
| pytest.skip("research/sceneseg/ref_runs not present") |
| for tag in ("cpp_bgr", "cpp_rgb"): |
| lg = np.load(d / f"{tag}_logits_ort_f32.npy") |
| with Image.open(d / f"{tag}_mask_srcres.png") as im: |
| want = np.asarray(im) |
| res = postprocess(want.shape, logits=lg) |
| np.testing.assert_array_equal(res.mask, want, err_msg=f"{name} {tag}") |
| with Image.open(d / f"{tag}_cls.png") as im: |
| cls = np.asarray(im) // 100 |
| assert (res.class_map == cls).mean() >= 0.99999 |
|
|
|
|
| def test_public_goldens(): |
| """research/sceneseg/public_data/golden: ORT logits (both orders) of 23 frames -> the stored class maps and |
| published masks, bit for bit.""" |
| from PIL import Image |
|
|
| for fid in public_ids(): |
| d = PUBLIC / "golden" / fid |
| if not (d / "ort_fp32.npz").is_file(): |
| pytest.skip("research/sceneseg/public_data/golden not present") |
| with np.load(d / "ort_fp32.npz") as z: |
| for order in config.CHANNEL_ORDERS: |
| with Image.open(d / f"mask_{order}.png") as im: |
| want = np.asarray(im) |
| res = postprocess(want.shape, logits=z[f"logits_{order}"]) |
| np.testing.assert_array_equal(res.class_map, z[f"cls_{order}"], err_msg=f"{fid} {order}") |
| np.testing.assert_array_equal(res.mask, want, err_msg=f"{fid} {order}") |
|
|
|
|
| def test_output_body(): |
| rng = np.random.default_rng(4) |
| res = postprocess((90, 160), logits=_logits(rng, (1, 3, 320, 640), ties=False)) |
| out = to_output(res, model="sceneseg-p150", channel_order="rgb") |
| assert isinstance(out, SceneSegMask) |
| body = json.loads(json.dumps(out.to_dict())) |
| assert body["model"] == "sceneseg-p150" and body["frame_id"] == "camera" |
| assert body["mask"]["shape"] == body["class_map"]["shape"] == [90, 160] |
| assert body["class_names"] == list(config.CLASS_NAMES) and sum(body["class_counts"].values()) == 90 * 160 |
| assert body["meta"]["channel_order"] == "rgb" and body["meta"]["foreground_fraction"] == pytest.approx( |
| float((res.mask == 255).mean()), abs=1e-6) |
| from PIL import Image |
|
|
| mask = np.asarray(Image.open(io.BytesIO(base64.b64decode(body["mask"]["data"])))) |
| cmap = np.asarray(Image.open(io.BytesIO(base64.b64decode(body["class_map"]["data"])))) |
| np.testing.assert_array_equal(mask, res.mask) |
| np.testing.assert_array_equal(cmap, res.class_map_src) |
| assert [d["label"] for d in out.to_dicts()] == [config.CLASS_NAMES[i] for i in np.unique(res.class_map_src)] |
| no_map = to_output(res, model="m", channel_order="bgr", include_class_map=False).to_dict() |
| assert "class_map" not in no_map and no_map["class_counts"] == body["class_counts"] |
| npz = out.to_dict("npz") |
| assert npz["mask"]["format"] == npz["class_map"]["format"] == "npz" |
| with pytest.raises(ValueError): |
| SceneSegMask(res.mask, class_map=res.class_map) |
| assert SceneSegMask(res.mask).class_counts().keys() <= {"0", "1"} | set(config.CLASS_NAMES) |
|
|
|
|
| def test_tutorial_class_maps_are_stable(): |
| """The stored small goldens of the 9 shipped samples (code/scripts/ref_golden.py): the host post-processing of |
| the stored logits samples reproduces the stored class-map pixels.""" |
| from tt_sceneseg.tests.research import GOLDENS |
|
|
| for name in TUTORIAL: |
| path = GOLDENS / f"{name}.npz" |
| if not path.is_file(): |
| pytest.skip("tests/goldens not generated (code/scripts/ref_golden.py)") |
| with np.load(path) as z: |
| for order in config.CHANNEL_ORDERS: |
| yx = z["sample_yx"] |
| lg = z[f"{order}.logits_samples"] |
| cls = class_map_from_logits(lg[None, :, :, None])[:, 0] |
| np.testing.assert_array_equal(cls, z[f"{order}.class_map"][yx[0], yx[1]]) |
|
|
