| |
| """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)) |
|
|
