| |
| """Host pre-processing of the SceneSeg node, bit-exact (research/sceneseg/SPEC.md section 3; PLAN.md 2.2). |
| |
| VisionPilot's ``scene_seg_model`` (``MR/ROS2/models/src/run_model_node.cpp:53-70``) converts the incoming image to |
| ``bgr8`` with ``cv_bridge`` and its backend (``MR/common/backends/onnx_runtime_backend.cpp:41-60`` == |
| ``tensorrt_backend.cpp:160-177``) squashes it to 640x320 with ``cv::resize`` (INTER_LINEAR, aspect ratio not kept), |
| scales by 1/255, normalises with the ImageNet constants in B, G, R order and splits it into B, G, R planes. The vendored |
| ``ttaw.image_linear`` (C14) reproduces every step bit for bit (OpenCV 4.8.1 and 4.11.0). |
| |
| The TT port uploads the 640x320 uint8 B, G, R squash (614 KB; :attr:`PreparedImage.u8_bgr`) and folds the |
| normalisation and, for the ``rgb`` variant, the channel permutation into the device graph (PLAN 2.2: mean subtracted |
| before the stem's zero padding, 1 / (255 std) and the permutation folded into the stem weights). The CPU reference |
| consumes :meth:`PreparedImage.network_input`, the ONNX input exactly as the node builds it. |
| |
| Channel order (LOAD-time variant): ``"bgr"`` is the deployed path (B, G, R planes into a network trained on R, G, B: |
| an upstream pre-processing bug the port keeps, user decision 1); ``"rgb"`` feeds R, G, B planes with the RGB |
| constants (training order). Each colour keeps its mean / std, so the two network inputs are channel permutations of |
| each other and both come from the same uint8 squash. |
| |
| Two bit-identical implementations of the squash (knob ``SCENESEG_HOST_SQUASH``, OPT_PLAN.md H1): ``numba`` |
| (default) is the compiled twin of ``host.fast`` with the RGB -> BGR swap folded into its reads (:func:`prepare_rgb` |
| takes the decoded RGB image directly); ``numpy`` is the vendored oracle on an explicit BGR copy. |
| """ |
| from __future__ import annotations |
|
|
| from dataclasses import dataclass |
| from typing import Any, Optional, Tuple |
|
|
| import numpy as np |
|
|
| from ..reference import config |
| from ..ttaw.image_linear import sceneseg_normalize, sceneseg_resize |
|
|
| __all__ = ["PreparedImage", "prepare_bgr", "prepare_rgb", "prepare_image", "rgb_to_bgr", "check_channel_order", |
| "squash_method"] |
|
|
| SQUASH_METHODS = ("numba", "numpy") |
| _METHOD: Optional[str] = None |
|
|
|
|
| def squash_method() -> str: |
| """The squash implementation of this process: knob ``SCENESEG_HOST_SQUASH`` (read once), ``numpy`` when numba is |
| missing.""" |
| global _METHOD |
| if _METHOD is None: |
| from ..tt.knobs import KNOBS |
| from . import fast |
|
|
| m = KNOBS.read().HOST_SQUASH |
| _METHOD = m if (m != "numba" or fast.HAVE_NUMBA) else "numpy" |
| return _METHOD |
|
|
|
|
| def check_channel_order(channel_order: str) -> str: |
| if channel_order not in config.CHANNEL_ORDERS: |
| raise ValueError(f"channel_order must be one of {config.CHANNEL_ORDERS}, got {channel_order!r}") |
| return channel_order |
|
|
|
|
| @dataclass(frozen=True) |
| class PreparedImage: |
| """One camera frame after the node's pre-processing. |
| |
| ``src_hw``: the source image size (the published mask has it); ``u8_bgr``: the 640x320 squash, uint8 |
| ``(320, 640, 3)`` B, G, R (the device input, for both channel orders); ``channel_order``: the variant the |
| network input is built for.""" |
|
|
| src_hw: Tuple[int, int] |
| u8_bgr: np.ndarray |
| channel_order: str = config.DEFAULT_CHANNEL_ORDER |
|
|
| def network_input(self) -> np.ndarray: |
| """The ONNX ``input`` float32 [1, 3, 320, 640] of this frame (``ttaw.image_linear.sceneseg_normalize``).""" |
| return sceneseg_normalize(self.u8_bgr, channel_order=self.channel_order) |
|
|
| def device_input(self) -> np.ndarray: |
| """The uint8 canvas the TT trace uploads, ``[1, 1, 320 * 640, 3]`` (NHWC rows, B, G, R), contiguous.""" |
| h, w = config.INPUT_HW |
| return np.ascontiguousarray(self.u8_bgr.reshape(1, 1, h * w, 3)) |
|
|
|
|
| def rgb_to_bgr(rgb: Any) -> np.ndarray: |
| """RGB uint8 HxWx3 (``ttaw.io`` decoding) -> B, G, R (``cv_bridge`` ``bgr8``), contiguous.""" |
| a = np.asarray(rgb) |
| if a.ndim != 3 or a.shape[2] != 3 or a.dtype != np.uint8: |
| raise ValueError(f"expected an RGB uint8 HxWx3 image, got {a.dtype} {a.shape}") |
| return np.ascontiguousarray(a[:, :, ::-1]) |
|
|
|
|
| def prepare_bgr(bgr: Any, *, channel_order: str = config.DEFAULT_CHANNEL_ORDER, |
| out: Optional[np.ndarray] = None, method: Optional[str] = None) -> PreparedImage: |
| """A B, G, R uint8 HxWx3 frame of any size (the node's ``bgr8`` input) -> :class:`PreparedImage`. |
| ``method``: ``numba`` / ``numpy`` (default: :func:`squash_method`), bit-identical.""" |
| check_channel_order(channel_order) |
| a = np.asarray(bgr) |
| if a.ndim != 3 or a.shape[2] != 3 or a.dtype != np.uint8: |
| raise ValueError(f"expected a BGR uint8 HxWx3 image, got {a.dtype} {a.shape}") |
| if (method or squash_method()) == "numba": |
| from . import fast |
|
|
| sq = fast.squash_into(a, config.INPUT_HW, swap=False, out=out) |
| else: |
| sq = sceneseg_resize(a, out=out) |
| return PreparedImage((int(a.shape[0]), int(a.shape[1])), sq, channel_order) |
|
|
|
|
| def prepare_rgb(rgb: Any, *, channel_order: str = config.DEFAULT_CHANNEL_ORDER, |
| out: Optional[np.ndarray] = None, method: Optional[str] = None) -> PreparedImage: |
| """A decoded R, G, B uint8 HxWx3 frame -> :class:`PreparedImage` = ``prepare_bgr(rgb_to_bgr(rgb))`` bit for bit. |
| The ``numba`` squash reads the RGB image with the channel swap folded in (no BGR copy); ``numpy`` makes the |
| copy and runs the oracle.""" |
| check_channel_order(channel_order) |
| a = np.asarray(rgb) |
| if a.ndim != 3 or a.shape[2] != 3 or a.dtype != np.uint8: |
| raise ValueError(f"expected an RGB uint8 HxWx3 image, got {a.dtype} {a.shape}") |
| if (method or squash_method()) == "numba": |
| from . import fast |
|
|
| sq = fast.squash_into(a, config.INPUT_HW, swap=True, out=out) |
| return PreparedImage((int(a.shape[0]), int(a.shape[1])), sq, channel_order) |
| return prepare_bgr(rgb_to_bgr(a), channel_order=channel_order, out=out, method="numpy") |
|
|
|
|
| def prepare_image(image: Any, *, channel_order: str = config.DEFAULT_CHANNEL_ORDER) -> PreparedImage: |
| """Any image the API takes (path, bytes, PIL image, RGB uint8 array, ``ttaw.io.CameraImage``) -> decoded to RGB by |
| ``ttaw.io`` (alpha dropped, as ``cv::imdecode`` / ``cv_bridge`` do), then :func:`prepare_bgr`.""" |
| from ..ttaw import io as _io |
|
|
| if isinstance(image, _io.CameraImage): |
| rgb = image.image |
| else: |
| rgb = _io.load_image(image) |
| return prepare_rgb(rgb, channel_order=channel_order) |
|
|
