# SPDX-License-Identifier: Apache-2.0 """Host pre-processing of TransFusion-L with Autoware's semantics (research/transfusion/SPEC.md section 3; PLAN.md 2.5), shared by the device path and the CPU reference so both see **identical** voxel and feature tensors. Input contract (the cloud is in the model frame, Autoware's ``base_link``; point features ``x, y, z, intensity, time_lag``, ``num_point_feature_size: 5``): 1. ``points`` with a ``time_lag`` column: an already densified cloud (what Autoware's ``PointCloudDensification`` hands to the voxelizer), used as is; 2. ``points`` (+ optional ``sweeps=[{"points", "time_lag_s", "T_current_from_sweep"}]``): client-side densification (the stateless API of PLAN.md D2); each past sweep goes through the float32 sweep kernel (``ttaw.pointcloud.densify_sweeps``); 3. ``points`` + ``stream={"id", "timestamp_s", "T_world_from_ego", "reset"}``: the optional server-side cache, Autoware's state machine per stream id (``densification_num_past_frames: 1``, cold start, no age limit, capacity drops whole sweeps; ``ttaw.pointcloud.SweepDensifier``). A missing pose is Autoware's TF failure: the frame is refused (HTTP 400) and not cached; 4. ``points`` alone: one sweep, ``time_lag = 0``. ``intensity`` is required: Autoware reads it as the uint8 at byte offset 12 and feeds ``float(uint8)`` (0..255, no normalisation; ``preprocess_kernel.cu:66``). Float inputs get that semantics here: non-finite -> 0, clipped to [0, 255], truncated like a C cast to ``uint8_t`` (the number of changed values is reported in ``meta``). Hygiene (documented deviations): rows with a non-finite x / y / z never enter (``ttaw.io.PointCloud`` drops them at decode; Autoware's range test lets NaN through, S:transfusion:117), and all-zero x, y, z rows (organized-cloud fillers such as ``test.pcd``'s) are dropped, as in the CenterPoint port. Sampling (PLAN.md D2, S:transfusion:104-111): - ``"clean"`` (default): every point; the range test of ``generateVoxels_random_kernel`` (strict interior: a point with ``x <= min_x`` or ``x >= max_x`` ... is rejected, ``preprocess_kernel.cu:171-174``); the cell ``floorf((x - min_x) / 0.24)`` in float32, clamped; the **first 20 points of a cell in input order** (Autoware: the GPU atomics' order); pillar ids in raster (y-major) order, at most ``N_MAX = 65,536`` (overflow keeps the first pillars and is flagged). Autoware would skip frames with < 5,000 or (UNVERIFIED) > 60,000 pillars: reported in ``meta.autoware_would_skip``, the frame is still run; - ``"autoware_compat"``: the deployed node's ``shufflePoints_kernel`` emulated: the sweep points are read through a permutation of the WHOLE 2,000,000-point capacity at an offset, so a buffer slot receives a real point only with probability ``count / capacity`` and a zero point (0, 0, 0, 0, 0) otherwise; the tail of the last 256-thread block is +INF; the whole buffer is voxelized, so one pillar at the origin always holds 20 zero points (S:transfusion:106-110). The permutation is ``numpy.random.default_rng(0).permutation(capacity)`` (Autoware's ``std::shuffle(..., default_random_engine(0))`` is libstdc++-specific) and the offset a per-request integer (Autoware's is ``std::rand() % capacity``, time-seeded), so the emulation is deterministic and statistically equivalent. Frames below 5,000 or (UNVERIFIED) above 60,000 pillars are skipped, as Autoware does. Decoration (:func:`decorate_pillars`, the in-graph ``PillarFeatureNet(legacy=True)`` of ``transfusion.onnx``, in float32, PLAN.md 2.5: bf16 here would cost 0.6 m): ``mean = (sum of the 20 slots) / num_points``, ``center_zyx = coords * (voxel_z, voxel_y, voxel_x) + (voxel / 2 + min)`` (the ONNX constants), and the 11 features ``[x - x_c, y - y_c, z - z_c, intensity, time_lag, x - mean_x, y - mean_y, z - mean_z, x - x_c, y - y_c, z - z_c]`` (the pillar-centre offsets twice, no absolute coordinates), padded slots 0 (S:transfusion:157). Device inputs built here: the ``(N_MAX, 20, 11)`` float32 decorated features (rows past the pillar count zero) and the gather index of the canvas (C19: every cell ``iy * 768 + ix`` -> its pillar row, or row ``N_MAX``, an explicit zero row); the ONNX inputs ``voxels`` / ``num_points`` / ``coors`` of the reference (:meth:`Prepared.onnx_inputs`). numpy only. """ from __future__ import annotations import math from dataclasses import dataclass, field from typing import Any, Dict, Mapping, Optional, Sequence, Tuple import numpy as np from .. import io as tio from ..reference.config import (COMPAT_SHUFFLE_SEED, DEFAULT_SAMPLING, DROP_ZERO_POINTS, N_MAX, NUM_POINT_FEATURES, SAMPLING_MODES, TransFusionConfig) from ..ttaw.geometry import as_transform from ..ttaw.pointcloud import StreamDensifiers, densify_sweeps, nonzero_rows from ..ttaw.voxelize import PillarGridSpec, PillarSet, assign_pillars, canvas_gather_index __all__ = ["Prepared", "TransFusionPreprocessor", "TIME_LAG_FIELD", "INTENSITY_FIELD", "intensity_u8", "strict_range_mask", "autoware_shuffle", "decorate_pillars", "parse_sampling", "voxel_inputs"] TIME_LAG_FIELD = "time_lag" # a column with this name marks a pre-densified cloud (seconds) INTENSITY_FIELD = "intensity" THREADS_PER_BLOCK = 256 # transfusion_config.hpp threads_for_voxel_ (the shuffle's INF tail granularity) # ------------------------------------------------------------------------------------------------ helpers def intensity_u8(values: Any) -> Tuple[np.ndarray, int]: """Autoware's ``float(uint8 intensity)`` from any numeric column: non-finite -> 0, clip to [0, 255], truncate toward zero (a C cast). Returns (float32 values, number of values that changed).""" v = np.asarray(values, dtype=np.float64) out = np.trunc(np.clip(np.nan_to_num(v, nan=0.0, posinf=255.0, neginf=0.0), 0.0, 255.0)).astype(np.float32) changed = int(np.count_nonzero(out.astype(np.float64) != v)) return out, changed def strict_range_mask(points: np.ndarray, spec: PillarGridSpec) -> np.ndarray: """``generateVoxels_random_kernel``'s range test (``preprocess_kernel.cu:171-174``): a point is rejected when ``x <= min_x || x >= max_x || y <= min_y || y >= max_y || z <= min_z || z >= max_z`` (float32 compares; the open interval, unlike CenterPoint's half-open one). A NaN coordinate fails here (Autoware would let it pass).""" p = np.asarray(points, dtype=np.float32) keep = np.ones(len(p), dtype=bool) for a in range(3): keep &= (p[:, a] > spec.min_f32[a]) & (p[:, a] < spec.max_f32[a]) return keep def autoware_shuffle(points: np.ndarray, capacity: int, perm: np.ndarray, offset: int) -> np.ndarray: """``shufflePoints_kernel`` and its buffers (``transfusion_trt.cpp:185-207``, ``preprocess_kernel.cu:99-139``): the (capacity, 5) float32 buffer the voxelizer reads. ``dst < count`` takes ``aux[perm[(dst + offset) % capacity]]`` from the zero-cleared aux buffer (a real point only when that source index is < count), ``count <= dst < ceil256(count)`` is +INF, the rest stays zero (the voxelizer is called with ``points_size = capacity``).""" pts = np.asarray(points, dtype=np.float32) count = len(pts) if count > capacity: raise ValueError(f"{count} points exceed the cloud capacity {capacity}") out = np.zeros((capacity, pts.shape[1]), dtype=np.float32) nthreads = min(int(math.ceil(count / THREADS_PER_BLOCK)) * THREADS_PER_BLOCK, capacity) dst = np.arange(nthreads, dtype=np.int64) src = np.asarray(perm, dtype=np.int64)[(dst + int(offset)) % capacity] real = dst < count take = real & (src < count) # aux[src] is a real point; otherwise the zero point stays out[dst[take]] = pts[src[take]] out[dst[~real]] = np.inf return out def decorate_pillars(pillars: PillarSet, scale_zyx: np.ndarray, offset_zyx: np.ndarray) -> np.ndarray: """The in-graph decoration of ``transfusion.onnx`` (module docstring) -> (max_pillars, K, 11) float32; rows past ``num_pillars`` and slots past each pillar's ``num_points`` are zero. Summation over the slots is sequential in slot order (a fixed order; ONNX Runtime's ``ReduceSum`` may differ in the last bit).""" raw = pillars.points P, K, F = raw.shape if F < NUM_POINT_FEATURES: raise ValueError(f"pillar rows have {F} columns, the decoration needs {NUM_POINT_FEATURES}") n = pillars.num_pillars out = np.zeros((P, K, 11), dtype=np.float32) if n == 0: return out r = raw[:n] xyz = r[:, :, :3] s = np.zeros((n, 3), dtype=np.float32) for k in range(K): # ReduceSum over all slots (padded slots are zero) s += xyz[:, k] mean = s / pillars.num_points[:n].astype(np.float32)[:, None] center_zyx = pillars.coords[:n].astype(np.float32) * np.asarray(scale_zyx, np.float32) + np.asarray( offset_zyx, np.float32) center_xyz = center_zyx[:, ::-1] f_center = xyz - center_xyz[:, None, :] f = out[:n] f[:, :, 0:3] = f_center f[:, :, 3:5] = r[:, :, 3:5] f[:, :, 5:8] = xyz - mean[:, None, :] f[:, :, 8:11] = f_center f *= (np.arange(K)[None, :] < pillars.num_points[:n, None]).astype(np.float32)[:, :, None] return out def voxel_inputs(pillars: PillarSet) -> Dict[str, np.ndarray]: """The ONNX inputs of the kept pillars: ``voxels`` (N, 20, 5) float32 (unused slots 0), ``num_points`` (N,) int32, ``coors`` (N, 4) int32 ``(0, 0, iy, ix)`` (``preprocess_kernel.cu:236``).""" n = pillars.num_pillars coors = np.zeros((n, 4), dtype=np.int32) coors[:, 2:] = pillars.coords[:n, 1:] return {"voxels": np.ascontiguousarray(pillars.points[:n]), "num_points": pillars.num_points[:n].astype(np.int32), "coors": coors} def parse_sampling(sampling: Any) -> Dict[str, Any]: """``"clean"`` / ``"autoware_compat"`` / ``{"mode": ..., "offset": int}`` -> ``{"mode", "offset"}`` (``offset``: the shuffle offset of ``autoware_compat``, 0 <= offset < cloud_capacity).""" if sampling is None: sampling = DEFAULT_SAMPLING if isinstance(sampling, str): spec: Dict[str, Any] = {"mode": sampling} elif isinstance(sampling, Mapping): spec = dict(sampling) else: raise tio.InputError("sampling must be 'clean', 'autoware_compat' or {mode, offset}") unknown = sorted(set(spec) - {"mode", "offset"}) if unknown: raise tio.InputError(f"sampling has unknown keys {unknown} (allowed: mode, offset)") mode = spec.get("mode", DEFAULT_SAMPLING) if mode not in SAMPLING_MODES: raise tio.InputError(f"sampling mode {mode!r} is not one of {list(SAMPLING_MODES)}") offset = spec.get("offset", 0) if isinstance(offset, bool) or not isinstance(offset, (int, np.integer)) or offset < 0: raise tio.InputError("sampling.offset must be a non-negative integer") if mode == "clean" and offset: raise tio.InputError("sampling.offset only applies to mode 'autoware_compat'") return {"mode": mode, "offset": int(offset)} # ------------------------------------------------------------------------------------------------ prepared @dataclass class Prepared: """One frame, ready for the device or the CPU reference.""" points: np.ndarray # (M, 5) float32 x, y, z, intensity, time_lag: the densified cloud pillars: PillarSet # capacity N_MAX, 20 slots, the 5 raw features per slot canvas_index: np.ndarray # (768 * 768,) uint32: cell -> pillar row; N_MAX = the zero sentinel row frame_id: str = "base_link" densify: Dict[str, Any] = field(default_factory=dict) sampling: Dict[str, Any] = field(default_factory=dict) skipped: Optional[str] = None # autoware_compat: Autoware's reason to skip the frame (no inference, no objects) autoware_would_skip: Optional[str] = None intensity_changed: int = 0 # the decoration's (scale_zyx, offset_zyx); ``features`` is computed from it on first access (OPT_PLAN H1: the # pair staging's numba twin writes the bf16 rows straight from ``pillars`` and never builds the float32 array) decoration: Optional[Tuple[np.ndarray, np.ndarray]] = None _features: Optional[np.ndarray] = field(default=None, repr=False) @property def features(self) -> np.ndarray: """(N_MAX, 20, 11) float32 decorated features (``decorate_pillars``; computed once, on first access).""" if self._features is None: self._features = decorate_pillars(self.pillars, *self.decoration) return self._features @property def features_ready(self) -> bool: return self._features is not None @property def num_pillars(self) -> int: return self.pillars.num_pillars @property def coords(self) -> np.ndarray: return self.pillars.coords def onnx_inputs(self) -> Dict[str, np.ndarray]: return voxel_inputs(self.pillars) def stats(self) -> Dict[str, Any]: """Pillar statistics, the sampling, Autoware's skip verdicts and the densification info (``meta``).""" s = self.pillars.stats() s.update({"sampling": dict(self.sampling), "skipped": self.skipped, "autoware_would_skip": self.autoware_would_skip, "intensity_changed": self.intensity_changed, "densify": self.densify}) return s def _cloud(source: Any, what: str) -> tio.PointCloud: if source is None: raise tio.InputError(f"{what} is missing") return tio.load_points(source) def _xyzi(cloud: tio.PointCloud, what: str) -> Tuple[np.ndarray, np.ndarray, int]: if not {"x", "y", "z"} <= set(cloud.fields): raise tio.InputError(f"{what} needs fields x, y, z (got {list(cloud.fields)})") if INTENSITY_FIELD not in cloud.fields: raise tio.InputError(f"{what} needs an 'intensity' field (0..255, Autoware's uint8 value; TransFusion uses " f"it as a network input), got {list(cloud.fields)}") xyz = np.ascontiguousarray(cloud.select(("x", "y", "z")), dtype=np.float32) inten, changed = intensity_u8(cloud.select((INTENSITY_FIELD,))[:, 0]) return xyz, inten, changed class TransFusionPreprocessor: """Stateful only through the per-stream densification caches (``stream`` inputs).""" def __init__(self, config: TransFusionConfig, *, max_streams: int = 16, drop_zero_points: bool = DROP_ZERO_POINTS, decoration: Optional[Tuple[np.ndarray, np.ndarray]] = None, voxelizer: str = "numpy", densify: str = "numpy"): if voxelizer not in ("numpy", "numba"): raise ValueError("voxelizer must be 'numpy' or 'numba'") if densify not in ("numpy", "numba"): raise ValueError("densify must be 'numpy' or 'numba'") # "numba" (OPT_PLAN H3, TRANSFUSION_HOST_DENSIFY): host/fast.densify_single, the one-pass twin of the # single-sweep path (x, y, z, uint8 intensity, zero-row drop); the numpy path stays the oracle self.densify = densify # "numba" (OPT_PLAN H2, TRANSFUSION_HOST_VOXELIZER): host/fast.voxelize_strict, the counting-sort twin of # strict_range_mask + assign_pillars (field-by-field equal PillarSet); the numpy path stays the oracle self.voxelizer = voxelizer self.config = config self.spec = config.grid_spec self.drop_zero_points = drop_zero_points self.capacity = int(config.node["cloud_capacity"]) self.scale_zyx, self.offset_zyx = decoration or (config.decoration_scale_zyx, config.decoration_offset_zyx) # lidar_transfusion drops a sweep when ``count + sweep >= cloud_capacity`` (voxel_generator.cpp:72), the # shared densifier when ``count + sweep > capacity`` (CenterPoint's rule): capacity - 1 makes them equal self.densify_capacity = self.capacity - 1 self.streams = StreamDensifiers(num_past_frames=int(config.node["densification_num_past_frames"]), cloud_capacity=self.densify_capacity, max_streams=max_streams) self._perm: Optional[np.ndarray] = None @property def compat_permutation(self) -> np.ndarray: """The fixed shuffle table of ``autoware_compat`` (computed once: 2,000,000 entries).""" if self._perm is None: self._perm = np.random.default_rng(COMPAT_SHUFFLE_SEED).permutation(self.capacity) return self._perm # ---------------------------------------------------------------------------------------- densification def densified_points(self, points: Any = None, *, sweeps: Optional[Sequence[Mapping[str, Any]]] = None, stream: Optional[Mapping[str, Any]] = None) -> Tuple[np.ndarray, Dict[str, Any], str, int]: """The (N, 5) float32 cloud (x, y, z, intensity, time_lag) of one request, its densification info, frame id and the number of intensity values changed by the uint8 semantics.""" cloud = _cloud(points, "this model needs 'points'") frame_id = cloud.frame_id if TIME_LAG_FIELD in cloud.fields: # 1. pre-densified if sweeps or stream is not None: raise tio.InputError("a cloud with a 'time_lag' column is already densified: do not send 'sweeps' " "or 'stream' with it") if self.densify == "numba" and {"x", "y", "z", INTENSITY_FIELD} <= set(cloud.fields): from .fast import densify_single idx = {n: i for i, n in enumerate(cloud.fields)} got = densify_single(cloud.points, [idx[n] for n in ("x", "y", "z", INTENSITY_FIELD)], self.drop_zero_points, idx[TIME_LAG_FIELD]) if got is not None: # the H3 twin pts, changed = got return pts, {"mode": "pre-densified", "num_points": int(len(pts))}, frame_id, changed xyz, inten, changed = _xyzi(cloud, "points") lag = np.asarray(cloud.select((TIME_LAG_FIELD,))[:, 0], dtype=np.float32) if not np.isfinite(lag).all(): raise tio.InputError("time_lag holds non-finite values") pts = np.concatenate([xyz, inten[:, None], lag[:, None]], axis=1) if self.drop_zero_points: pts = pts[nonzero_rows(pts)] return np.ascontiguousarray(pts), {"mode": "pre-densified", "num_points": int(len(pts))}, frame_id, changed if self.densify == "numba" and not sweeps and stream is None \ and {"x", "y", "z", INTENSITY_FIELD} <= set(cloud.fields): from .fast import densify_single idx = {n: i for i, n in enumerate(cloud.fields)} got = densify_single(cloud.points, [idx[n] for n in ("x", "y", "z", INTENSITY_FIELD)], self.drop_zero_points) if got is not None: # 4. single sweep (H3 twin) pts, changed = got return pts, {"mode": "single", "num_points": int(len(pts))}, frame_id, changed xyz, inten, changed = _xyzi(cloud, "points") if self.drop_zero_points: keep = nonzero_rows(xyz) xyz, inten = xyz[keep], inten[keep] if sweeps and stream is not None: raise tio.InputError("give 'sweeps' (client-side densification) or 'stream' (server-side), not both") if sweeps: # 2. client-side past, past_i = [], [] for i, s in enumerate(sweeps): if not isinstance(s, Mapping): raise tio.InputError(f"sweeps[{i}] must be an object {{points, time_lag_s, T_current_from_sweep}}") try: lag = float(s.get("time_lag_s", 0.0)) except (TypeError, ValueError): lag = float("nan") if not np.isfinite(lag): raise tio.InputError(f"sweeps[{i}].time_lag_s must be finite") T = s.get("T_current_from_sweep") T = None if T is None else as_transform(T, field=f"sweeps[{i}].T_current_from_sweep") sx, si, sc = _xyzi(_cloud(s.get("points"), f"sweeps[{i}].points"), f"sweeps[{i}].points") if self.drop_zero_points: k = nonzero_rows(sx) sx, si = sx[k], si[k] changed += sc past.append((sx, lag, T)) past_i.append(si) pts, info = densify_sweeps(xyz, past, point_feature_size=NUM_POINT_FEATURES, cloud_capacity=self.densify_capacity, current_intensity=inten, sweep_intensity=past_i) return pts, {"mode": "sweeps", **info.to_dict()}, frame_id, changed if stream is not None: # 3. server-side, per stream id if not isinstance(stream, Mapping): raise tio.InputError("stream must be an object {id, timestamp_s, T_world_from_ego, reset}") sid = str(stream.get("id", "default")) try: stamp = float(stream["timestamp_s"]) except (KeyError, TypeError, ValueError): stamp = float("nan") if not np.isfinite(stamp): raise tio.InputError("stream.timestamp_s (seconds) is required for server-side densification") T = stream.get("T_world_from_ego") T = None if T is None else as_transform(T, field="stream.T_world_from_ego") d = self.streams.get(sid, reset=bool(stream.get("reset", False))) if not d.enqueue(xyz, stamp, T, intensity=inten): raise tio.InputError( f"stream.T_world_from_ego is required with densification_num_past_frames={d.num_past_frames}: " "Autoware skips a frame without its map->base_link transform and does not cache it (nor did this " "request)") pts, info = d.sweep_points(point_feature_size=NUM_POINT_FEATURES) return pts, {"mode": "stream", "stream_id": sid, **info.to_dict()}, frame_id, changed pts = np.zeros((len(xyz), NUM_POINT_FEATURES), np.float32) # 4. single sweep pts[:, :3] = xyz pts[:, 3] = inten return pts, {"mode": "single", "num_points": int(len(pts))}, frame_id, changed # ------------------------------------------------------------------------------------------ voxelization def voxelize(self, points5: np.ndarray, sampling: Any = None) -> Tuple[PillarSet, Dict[str, Any]]: """(N, 5) cloud -> deterministic pillars (capacity N_MAX) under the sampling mode; returns the sampling record (mode, offset, the voxelizer's input size).""" smp = parse_sampling(sampling) pts = np.asarray(points5, dtype=np.float32) if smp["mode"] == "autoware_compat": if smp["offset"] >= self.capacity: raise tio.InputError(f"sampling.offset must be < cloud_capacity ({self.capacity})") pts = autoware_shuffle(pts, self.capacity, self.compat_permutation, smp["offset"]) smp["voxelizer_points"] = int(len(pts)) pillars = None if self.voxelizer == "numba": from .fast import voxelize_strict pillars = voxelize_strict(pts, self.spec) if pillars is None: pillars = assign_pillars(pts[strict_range_mask(pts, self.spec)], self.spec, shuffle_seed=None) pillars.num_points_input = int(len(points5)) return pillars, smp def __call__(self, points: Any = None, *, sweeps: Optional[Sequence[Mapping[str, Any]]] = None, stream: Optional[Mapping[str, Any]] = None, sampling: Any = None) -> Prepared: pts, info, frame_id, changed = self.densified_points(points, sweeps=sweeps, stream=stream) pillars, smp = self.voxelize(pts, sampling) n_total = pillars.num_pillars_total verdict = None if n_total < self.config.min_voxels: verdict = "too_few_voxels" # transfusion_trt.cpp:217-223 elif n_total > self.config.max_voxels: verdict = "too_many_voxels" # UNVERIFIED: TensorRT shape above the profile max (:225-250) skipped = verdict if smp["mode"] == "autoware_compat" else None index = canvas_gather_index(pillars.coords, pillars.num_pillars, self.spec, sentinel=self.spec.max_pillars) return Prepared(points=pts, pillars=pillars, canvas_index=index, frame_id=frame_id, densify=info, sampling=smp, skipped=skipped, autoware_would_skip=verdict, intensity_changed=changed, decoration=(self.scale_zyx, self.offset_zyx)) def describe(self) -> Dict[str, Any]: return {"n_max": N_MAX, "points_per_pillar": self.spec.max_points_per_pillar, "order": self.spec.order, "drop_zero_points": self.drop_zero_points, "cloud_capacity": self.capacity, "sampling_modes": list(SAMPLING_MODES), "default_sampling": DEFAULT_SAMPLING, "compat_shuffle_seed": COMPAT_SHUFFLE_SEED, "streams": self.streams.describe()}