# SPDX-License-Identifier: Apache-2.0 """Host post-processing of the SceneSeg node, exact (research/sceneseg/SPEC.md section 5). ``MR/ROS2/models/src/run_model_node.cpp:117-188`` (the CUDA path ``cuda_visualization_kernels.cu:13-42`` is equivalent): for every pixel of the ``[1, 3, 320, 640]`` fp32 logits, scan the classes c = 0, 1, 2 with ``max_score = -1e9`` and a strict ``>`` (ties go to the lowest index), publish ``255`` where the winner is class 1 (foreground) and ``0`` elsewhere, and resize that mono8 mask back to the input size with ``cv::resize(INTER_NEAREST)``. No threshold, no softmax, no morphology, no temporal filter. The same decision written on class margins (``m01 = l1 - l0``, ``m12 = l1 - l2``, ``m20 = l2 - l0``) is what the TT graph evaluates with the margins folded into its last conv (PLAN 2.2): foreground iff ``m01 > 0 and m12 >= 0``, road iff ``m20 > 0 and m12 < 0``, background otherwise. For float logits the two forms agree exactly (the sign of a float difference is the comparison), which :func:`class_map_from_margins` and the tests rely on. The 3-class map (0 background incl. sky, 1 foreground, 2 drivable road) is a by-product the node does not publish; the API returns it next to the mask, resized with the same nearest-neighbour rule. Two bit-identical implementations of the mask + resize step (knob ``SCENESEG_HOST_POST``, OPT_PLAN.md H2): ``numba`` (default) is the one-pass gather of ``host.fast.post_into``; ``numpy`` is :func:`foreground_mask` + :func:`to_source` (the oracle). """ from __future__ import annotations from dataclasses import dataclass from typing import Optional, Tuple import numpy as np from ..reference import config from ..ttaw.image_linear import opencv_resize_nearest __all__ = ["post_method", "SceneSegResult", "class_map_from_logits", "class_map_from_margins", "margins_from_logits", "foreground_mask", "to_source", "postprocess"] POST_METHODS = ("numba", "numpy") _POST_METHOD: Optional[str] = None def post_method() -> str: """The post-processing implementation of this process: knob ``SCENESEG_HOST_POST`` (read once), ``numpy`` when numba is missing.""" global _POST_METHOD if _POST_METHOD is None: from ..tt.knobs import KNOBS from . import fast m = KNOBS.read().HOST_POST _POST_METHOD = m if (m != "numba" or fast.HAVE_NUMBA) else "numpy" return _POST_METHOD def _planes(logits: np.ndarray) -> np.ndarray: a = np.asarray(logits) if a.ndim == 4: if a.shape[0] != 1: raise ValueError(f"batch 1 only, got {a.shape}") a = a[0] if a.ndim != 3 or a.shape[0] != config.NUM_CLASSES: raise ValueError(f"expected [1, {config.NUM_CLASSES}, H, W] or [{config.NUM_CLASSES}, H, W] logits, got " f"{np.shape(logits)}") return a def class_map_from_logits(logits: np.ndarray) -> np.ndarray: """The node's per-pixel decision (``run_model_node.cpp:148-163``), literally: best = 0, max = -1e9; for c in 0, 1, 2: if score > max: max = score, best = c. -> uint8 ``(H, W)``. Equals ``argmax`` (first maximum) for any finite logits above -1e9.""" planes = _planes(logits) best = np.zeros(planes.shape[1:], np.uint8) mx = np.full(planes.shape[1:], np.float32(config.ARGMAX_INIT), dtype=np.float32) for c in range(planes.shape[0]): score = planes[c].astype(np.float32, copy=False) upd = score > mx mx = np.where(upd, score, mx) best[upd] = c return best def margins_from_logits(logits: np.ndarray) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """``(m01, m12, m20) = (l1 - l0, l1 - l2, l2 - l0)`` in float32 (what the TT graph's folded last conv emits).""" p = _planes(logits).astype(np.float32, copy=False) return p[1] - p[0], p[1] - p[2], p[2] - p[0] def class_map_from_margins(m01: np.ndarray, m12: np.ndarray, m20: Optional[np.ndarray] = None) -> np.ndarray: """The node's decision on class margins: 1 where ``m01 > 0 and m12 >= 0`` (foreground, lowest-index ties), 2 where ``m20 > 0 and m12 < 0`` (road; needs ``m20``, else road is reported as background), 0 elsewhere.""" m01, m12 = np.asarray(m01), np.asarray(m12) out = np.zeros(m01.shape, np.uint8) out[(m01 > 0) & (m12 >= 0)] = 1 if m20 is not None: out[(np.asarray(m20) > 0) & (m12 < 0)] = 2 return out def foreground_mask(class_map: np.ndarray, *, foreground_class: int = config.FOREGROUND_CLASS, value: int = config.MASK_VALUE) -> np.ndarray: """``mask = (best_class == 1) ? 255 : 0`` (uint8).""" return np.where(np.asarray(class_map) == foreground_class, np.uint8(value), np.uint8(0)).astype(np.uint8) def to_source(image: np.ndarray, src_hw: Tuple[int, int]) -> np.ndarray: """``cv::resize(mask, src_size, INTER_NEAREST)`` (``ttaw.image_linear.opencv_resize_nearest``, bit-exact).""" return opencv_resize_nearest(image, src_hw) @dataclass(frozen=True) class SceneSegResult: """``class_map``: uint8 ``(320, 640)`` at network resolution; ``mask`` / ``class_map_src``: uint8 ``(H, W)`` at the source size (``mask`` is the published mono8 topic: 255 = foreground).""" class_map: np.ndarray mask: np.ndarray class_map_src: np.ndarray @property def foreground_fraction(self) -> float: return float((self.mask == config.MASK_VALUE).mean()) def postprocess(src_hw: Tuple[int, int], *, logits: Optional[np.ndarray] = None, class_map: Optional[np.ndarray] = None, method: Optional[str] = None) -> SceneSegResult: """Logits (CPU reference, or a TT debug readback) or the device's class map -> :class:`SceneSegResult`. ``method``: ``numba`` / ``numpy`` (default: :func:`post_method`), bit-identical.""" if (logits is None) == (class_map is None): raise ValueError("give exactly one of logits / class_map") cmap = class_map_from_logits(logits) if logits is not None else np.asarray(class_map, np.uint8) if cmap.shape != tuple(config.INPUT_HW): raise ValueError(f"class map must be {config.INPUT_HW}, got {cmap.shape}") if (method or post_method()) == "numba": from . import fast mask, cls = fast.post_into(cmap, src_hw, foreground_class=config.FOREGROUND_CLASS, value=config.MASK_VALUE) return SceneSegResult(cmap, mask, cls) return SceneSegResult(cmap, to_source(foreground_mask(cmap), src_hw), to_source(cmap, src_hw))