| |
| """The three in-network TensorRT plugins of ``bevdet_one_lt_d.onnx`` in torch fp32 (the vendor fp32 path). |
| |
| Clean-room from research/bevdet/SPEC.md section 4.2 (no vendor code), ported from the verified research reference |
| (``research/bevdet/scripts/bevdet_common.py``): |
| |
| - :func:`bev_pool`: ``bevdet::BEVPool``: ``out[c, ranks_bev[k]] += feat[ranks_feat[k] -> (n, c, hw)] * |
| depth_flat[ranks_depth[k]]`` over every table entry, zero-initialised (the intended semantics; the TRT plugin |
| leaves never-hit cells uninitialised, S:162). One ``index_add_`` in fp32. |
| - :func:`align_bev`: ``bevdet::AlignBEV``: for history slot n and BEV pixel (h, w): ``ix = a w + b h + c``, |
| ``iy = d w + e h + f`` (pixel coordinates), bilinear over the 4 neighbours with zero outside |
| (= ``grid_sample(align_corners=True, padding_mode="zeros")``, as in training). The 4 taps are evaluated literally |
| in float32 like the CUDA kernel, so an identity transform is an exact copy. |
| - :func:`gather_bev`: ``bevdet::GatherBEV``: ``cat(curr, s_1 .. s_8)`` with ``s_j = aligned[j - 1]`` if ``flag`` |
| else ``curr``; channel order [t, t-1, ..., t-8]. |
| |
| Inputs and outputs are torch tensors or numpy arrays (outputs: torch). The Preprocess plugin is host work |
| (``host.preprocess``). |
| """ |
| from __future__ import annotations |
|
|
| from typing import Any |
|
|
| import numpy as np |
|
|
| from ..host.calib import BEVPoolTables |
| from .config import BEVDetConfig |
|
|
| __all__ = ["bev_pool", "align_bev", "gather_bev"] |
|
|
|
|
| def _t(x: Any): |
| import torch |
|
|
| return torch.as_tensor(np.asarray(x, np.float32)) if not isinstance(x, torch.Tensor) else x.to(torch.float32) |
|
|
|
|
| def bev_pool(depth: Any, feat: Any, tables: BEVPoolTables, cfg: BEVDetConfig): |
| """depth [N, D, H, W] (softmax), feat [N, C, H, W] -> curr [1, C, ny, nx] (fp32, cells never hit are 0).""" |
| import torch |
|
|
| depth_t, feat_t = _t(depth), _t(feat) |
| N, C, H, W = feat_t.shape |
| ny, nx = cfg.bev_hw |
| nz = cfg.bev_xyz[2] |
| dflat = depth_t.reshape(-1) |
| fflat = feat_t.reshape(N, C, H * W).permute(1, 0, 2).reshape(C, N * H * W) |
| rd = torch.from_numpy(tables.ranks_depth.astype(np.int64)) |
| rf = torch.from_numpy(tables.ranks_feat.astype(np.int64)) |
| rb = torch.from_numpy(tables.ranks_bev.astype(np.int64)) |
| contrib = fflat[:, rf] * dflat[rd][None, :] |
| out = torch.zeros(C, nx * ny * nz, dtype=torch.float32) |
| out.index_add_(1, rb, contrib) |
| return out.reshape(1, C, ny, nx) |
|
|
|
|
| def align_bev(adj_feats: Any, transforms: Any): |
| """adj_feats [1, A, C, H, W], transforms [1, A, 6] -> aligned [1, A, C, H, W] (bilinear, zero outside).""" |
| import torch |
|
|
| adj = _t(adj_feats) |
| trs = _t(transforms) |
| B, A, C, H, W = adj.shape |
| out = torch.zeros_like(adj) |
| hh, ww = torch.meshgrid(torch.arange(H, dtype=torch.float32), torch.arange(W, dtype=torch.float32), |
| indexing="ij") |
| identity = torch.tensor([1, 0, 0, 0, 1, 0], dtype=torch.float32) |
| for n in range(A): |
| if torch.equal(trs[0, n], identity): |
| out[0, n] = adj[0, n] |
| continue |
| a, b, c0, d, e, f = (trs[0, n, k] for k in range(6)) |
| ix = a * ww + b * hh + c0 |
| iy = d * ww + e * hh + f |
| x0 = torch.floor(ix).double() |
| y0 = torch.floor(iy).double() |
| x1, y1 = x0 + 1, y0 + 1 |
| ixd, iyd = ix.double(), iy.double() |
| |
| taps = ((y0, x0, (x1 - ixd) * (y1 - iyd)), (y0, x1, (ixd - x0) * (y1 - iyd)), |
| (y1, x0, (x1 - ixd) * (iyd - y0)), (y1, x1, (ixd - x0) * (iyd - y0))) |
| src = adj[0, n] |
| acc = torch.zeros((C, H, W), dtype=torch.float32) |
| for yy, xx, wt in taps: |
| ok = (yy >= 0) & (yy < H) & (xx >= 0) & (xx < W) |
| yi = yy.clamp(0, H - 1).long() |
| xi = xx.clamp(0, W - 1).long() |
| vals = src[:, yi, xi] |
| acc += torch.where(ok[None], vals * wt.float()[None], torch.zeros((), dtype=torch.float32)) |
| out[0, n] = acc |
| return out |
|
|
|
|
| def gather_bev(aligned: Any, curr: Any, flag: int): |
| """aligned [1, A, C, H, W], curr [1, C, H, W] -> [1, (1 + A) C, H, W].""" |
| import torch |
|
|
| al, cu = _t(aligned), _t(curr) |
| A = al.shape[1] |
| parts = [cu[0]] + [cu[0] if int(flag) == 0 else al[0, n] for n in range(A)] |
| return torch.cat(parts, 0)[None] |
|
|
