Spaces:
Paused
Paused
| """BEV rasterization, per-cell modality fusion, and box decoding. | |
| This is the numpy analogue of BEVFusion's shared-BEV-space fusion: | |
| 1. Each modality's detections are splatted into a metric BEV feature map | |
| (:class:`ModalityBEV`) as Gaussian evidence. The splat covariance is the | |
| detection's position covariance (floored at the cell size for existence so | |
| tight LiDAR kernels are not lost between cell centers). Three kinds of | |
| evidence are accumulated per cell: | |
| - existence: noisy-OR of (confidence x kernel) — "is something here?"; | |
| - geometry: inverse-variance weighted votes for the detection's continuous | |
| mean/dims/yaw (weight = confidence x kernel / worst-axis variance), so | |
| LiDAR dominates position wherever it has evidence and there is no grid | |
| quantization in the decoded centers; | |
| - class: reliability-weighted class histograms. Camera class evidence is | |
| splatted with an inflated along-ray kernel: monocular semantics is | |
| positionally vague along the ray, which is precisely what lets a camera | |
| label flow to the LiDAR-anchored cell of the same object even when the | |
| camera's depth estimate is metres off. | |
| 2. :func:`fuse_maps` fuses per cell: noisy-OR for existence, summed | |
| inverse-variance geometry accumulators, summed weighted class histograms | |
| (camera weight 1.0 vs LiDAR 0.35 — camera is the semantics expert). | |
| 3. :func:`decode_bev` extracts local existence maxima above a threshold, | |
| applies euclidean NMS plus ray-aware NMS (two peaks on the same bearing | |
| within the monocular depth-ambiguity band are one object — the fused, | |
| LiDAR-anchored peak wins over the camera's displaced ghost), then decodes | |
| each surviving peak's neighbourhood into a 3D box with class + confidence. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| from dataclasses import dataclass, field | |
| from typing import Dict, List, Optional | |
| import numpy as np | |
| from scipy.ndimage import maximum_filter | |
| from sensorflow.bevfusion.geometry import BEVGrid | |
| from sensorflow.bevfusion.scenes import CLASSES | |
| from sensorflow.bevfusion.sensors import Detection | |
| CLASS_INDEX = {c: i for i, c in enumerate(CLASSES)} | |
| # Modality reliability weights for class evidence (camera >> lidar). | |
| CLASS_WEIGHT = {"camera": 1.0, "lidar": 0.35} | |
| # Camera class evidence is smeared further along the ray (see module docstring). | |
| CAMERA_CLASS_SIGMA_SCALE = 2.0 | |
| class ModalityBEV: | |
| """Per-modality BEV feature map built by splatting detections.""" | |
| grid: BEVGrid | |
| modality: str = "camera" | |
| miss: np.ndarray = field(default=None) # prod(1 - c*g): existence complement | |
| sum_w: np.ndarray = field(default=None) | |
| sum_wx: np.ndarray = field(default=None) | |
| sum_wy: np.ndarray = field(default=None) | |
| sum_wz: np.ndarray = field(default=None) | |
| sum_wl: np.ndarray = field(default=None) | |
| sum_ww: np.ndarray = field(default=None) | |
| sum_wh: np.ndarray = field(default=None) | |
| sum_wsin: np.ndarray = field(default=None) | |
| sum_wcos: np.ndarray = field(default=None) | |
| class_hist: np.ndarray = field(default=None) # (nx, ny, n_classes) | |
| def __post_init__(self): | |
| nx, ny = self.grid.nx, self.grid.ny | |
| self.miss = np.ones((nx, ny)) | |
| for name in ("sum_w", "sum_wx", "sum_wy", "sum_wz", "sum_wl", "sum_ww", | |
| "sum_wh", "sum_wsin", "sum_wcos"): | |
| setattr(self, name, np.zeros((nx, ny))) | |
| self.class_hist = np.zeros((nx, ny, len(CLASSES))) | |
| def existence(self) -> np.ndarray: | |
| return 1.0 - self.miss | |
| # -- splatting | |
| def _window(self, mx: float, my: float, ex: float, ey: float): | |
| g = self.grid | |
| ix0 = max(0, int((mx - ex - g.x_min) / g.cell)) | |
| ix1 = min(g.nx, int((mx + ex - g.x_min) / g.cell) + 1) | |
| iy0 = max(0, int((my - ey - g.y_min) / g.cell)) | |
| iy1 = min(g.ny, int((my + ey - g.y_min) / g.cell) + 1) | |
| if ix0 >= ix1 or iy0 >= iy1: | |
| return None | |
| xs = g.x_min + (np.arange(ix0, ix1) + 0.5) * g.cell | |
| ys = g.y_min + (np.arange(iy0, iy1) + 0.5) * g.cell | |
| X, Y = np.meshgrid(xs, ys, indexing="ij") | |
| return (slice(ix0, ix1), slice(iy0, iy1)), X - mx, Y - my | |
| def _kernel(dx: np.ndarray, dy: np.ndarray, cov: np.ndarray) -> np.ndarray: | |
| inv = np.linalg.inv(cov) | |
| q = inv[0, 0] * dx * dx + (inv[0, 1] + inv[1, 0]) * dx * dy + inv[1, 1] * dy * dy | |
| return np.exp(-0.5 * q) | |
| def splat(self, det: Detection) -> None: | |
| cov = np.array(det.cov, dtype=float) | |
| # Floor the existence/geometry kernel at ~the cell size so tight LiDAR | |
| # Gaussians are not lost between cell centers (grid smoothing only: | |
| # geometry votes still use the true covariance for their weights). | |
| floor = (0.8 * self.grid.cell) ** 2 | |
| cov_splat = cov + np.eye(2) * floor | |
| ex = 3.0 * math.sqrt(cov_splat[0, 0]) | |
| ey = 3.0 * math.sqrt(cov_splat[1, 1]) | |
| win = self._window(det.x, det.y, ex, ey) | |
| if win is None: | |
| return | |
| sl, dx, dy = win | |
| g = self._kernel(dx, dy, cov_splat) | |
| # Existence: noisy-OR accumulation. | |
| self.miss[sl] *= (1.0 - det.confidence * g) | |
| # Geometry: inverse-variance weighted votes for the *continuous* mean. | |
| # The kernel is normalized to unit mass so each detection contributes | |
| # exactly one measurement of total weight conf/sigma_eff^2 no matter | |
| # how widely its kernel is spread — proper inverse-variance fusion. | |
| eigvals = np.linalg.eigvalsh(cov) | |
| sigma_eff_sq = float(max(eigvals[-1], 1e-4)) # worst-axis variance | |
| g_mass = float(g.sum()) | |
| if g_mass <= 1e-12: | |
| return | |
| w = det.confidence * (g / g_mass) / sigma_eff_sq | |
| self.sum_w[sl] += w | |
| self.sum_wx[sl] += w * det.x | |
| self.sum_wy[sl] += w * det.y | |
| self.sum_wz[sl] += w * det.z | |
| self.sum_wl[sl] += w * det.dims[0] | |
| self.sum_ww[sl] += w * det.dims[1] | |
| self.sum_wh[sl] += w * det.dims[2] | |
| self.sum_wsin[sl] += w * math.sin(det.yaw) | |
| self.sum_wcos[sl] += w * math.cos(det.yaw) | |
| # Class: reliability-weighted histogram; camera semantics smeared | |
| # further along the ray so labels reach the LiDAR-anchored cell. | |
| cls_idx = CLASS_INDEX.get(det.class_name) | |
| if cls_idx is None: | |
| return | |
| if self.modality == "camera": | |
| cov_cls = cov * (CAMERA_CLASS_SIGMA_SCALE ** 2) + np.eye(2) * floor | |
| else: | |
| cov_cls = cov * 4.0 + np.eye(2) * floor | |
| exc = 3.0 * math.sqrt(cov_cls[0, 0]) | |
| eyc = 3.0 * math.sqrt(cov_cls[1, 1]) | |
| winc = self._window(det.x, det.y, exc, eyc) | |
| if winc is None: | |
| return | |
| slc, dxc, dyc = winc | |
| gc = self._kernel(dxc, dyc, cov_cls) | |
| wc = det.confidence * det.class_confidence * CLASS_WEIGHT.get(self.modality, 0.5) | |
| self.class_hist[slc[0], slc[1], cls_idx] += wc * gc | |
| def build_modality_map(dets: List[Detection], grid: BEVGrid, modality: str) -> ModalityBEV: | |
| bev = ModalityBEV(grid=grid, modality=modality) | |
| for det in dets: | |
| if det.modality == modality: | |
| bev.splat(det) | |
| return bev | |
| class FusedBEV: | |
| grid: BEVGrid | |
| existence: np.ndarray | |
| sum_w: np.ndarray | |
| sum_wx: np.ndarray | |
| sum_wy: np.ndarray | |
| sum_wz: np.ndarray | |
| sum_wl: np.ndarray | |
| sum_ww: np.ndarray | |
| sum_wh: np.ndarray | |
| sum_wsin: np.ndarray | |
| sum_wcos: np.ndarray | |
| class_hist: np.ndarray | |
| def fuse_maps(cam: ModalityBEV, lidar: ModalityBEV) -> FusedBEV: | |
| """Per-cell confidence-weighted fusion of the two modality feature maps.""" | |
| return FusedBEV( | |
| grid=cam.grid, | |
| existence=1.0 - cam.miss * lidar.miss, # noisy-OR across modalities | |
| sum_w=cam.sum_w + lidar.sum_w, | |
| sum_wx=cam.sum_wx + lidar.sum_wx, | |
| sum_wy=cam.sum_wy + lidar.sum_wy, | |
| sum_wz=cam.sum_wz + lidar.sum_wz, | |
| sum_wl=cam.sum_wl + lidar.sum_wl, | |
| sum_ww=cam.sum_ww + lidar.sum_ww, | |
| sum_wh=cam.sum_wh + lidar.sum_wh, | |
| sum_wsin=cam.sum_wsin + lidar.sum_wsin, | |
| sum_wcos=cam.sum_wcos + lidar.sum_wcos, | |
| class_hist=cam.class_hist + lidar.class_hist, | |
| ) | |
| def decode_bev(fused: FusedBEV, threshold: float = 0.30, | |
| nms_radius_m: float = 2.0, gather_radius_m: float = 2.0, | |
| ray_nms_bearing_deg: float = 1.5, | |
| ray_nms_range_m: float = 6.0) -> List[Dict]: | |
| """Decode fused BEV cells into 3D boxes with class + confidence.""" | |
| grid = fused.grid | |
| e = fused.existence | |
| local_max = (e >= maximum_filter(e, size=5)) & (e >= threshold) | |
| ixs, iys = np.nonzero(local_max) | |
| if ixs.size == 0: | |
| return [] | |
| order = np.argsort(-e[ixs, iys]) | |
| peaks = [(int(ixs[i]), int(iys[i])) for i in order] | |
| accepted: List[Dict] = [] | |
| gather_cells = max(1, int(round(gather_radius_m / grid.cell))) | |
| for ix, iy in peaks: | |
| cx, cy = grid.center(ix, iy) | |
| suppress = False | |
| for box in accepted: | |
| bx, by = box["_peak_xy"] | |
| if math.hypot(cx - bx, cy - by) < nms_radius_m: | |
| suppress = True | |
| break | |
| # Ray-aware NMS: same bearing, within the monocular depth | |
| # ambiguity band -> the weaker peak is the camera's range ghost. | |
| db = abs(math.atan2(cy, cx) - math.atan2(by, bx)) | |
| db = min(db, 2 * math.pi - db) | |
| if math.degrees(db) < ray_nms_bearing_deg and \ | |
| abs(math.hypot(cx, cy) - math.hypot(bx, by)) < ray_nms_range_m: | |
| suppress = True | |
| break | |
| if suppress: | |
| continue | |
| x0, x1 = max(0, ix - gather_cells), min(grid.nx, ix + gather_cells + 1) | |
| y0, y1 = max(0, iy - gather_cells), min(grid.ny, iy + gather_cells + 1) | |
| sl = (slice(x0, x1), slice(y0, y1)) | |
| w = float(fused.sum_w[sl].sum()) | |
| if w <= 1e-9: | |
| continue | |
| x = float(fused.sum_wx[sl].sum() / w) | |
| y = float(fused.sum_wy[sl].sum() / w) | |
| z = float(fused.sum_wz[sl].sum() / w) | |
| l = float(fused.sum_wl[sl].sum() / w) | |
| wd = float(fused.sum_ww[sl].sum() / w) | |
| h = float(fused.sum_wh[sl].sum() / w) | |
| yaw = float(math.atan2(fused.sum_wsin[sl].sum(), fused.sum_wcos[sl].sum())) | |
| hist = fused.class_hist[sl].sum(axis=(0, 1)) | |
| hist_total = float(hist.sum()) | |
| cls_idx = int(np.argmax(hist)) if hist_total > 0 else 0 | |
| accepted.append({ | |
| "bbox_3d": [x, y, z, l, wd, h, yaw], | |
| "class_name": CLASSES[cls_idx], | |
| "class_confidence": round(float(hist[cls_idx] / hist_total), 4) if hist_total > 0 else 0.0, | |
| "confidence": round(float(e[ix, iy]), 4), | |
| "_peak_xy": (cx, cy), | |
| }) | |
| for box in accepted: | |
| box.pop("_peak_xy", None) | |
| return accepted | |
| def fuse_frame(camera_dets: List[Detection], lidar_dets: List[Detection], | |
| grid: Optional[BEVGrid] = None, threshold: float = 0.30) -> List[Dict]: | |
| """Convenience: full rasterize -> fuse -> decode for one frame.""" | |
| grid = grid or BEVGrid() | |
| cam_map = build_modality_map(camera_dets, grid, "camera") | |
| lidar_map = build_modality_map(lidar_dets, grid, "lidar") | |
| return decode_bev(fuse_maps(cam_map, lidar_map), threshold=threshold) | |