heatmap
Gaussian center-heatmap targets for center-based 3D detection heads.
Center-based detectors (VoxelNeXt, VoxelMamba, LION) supervise a per-class BEV heatmap whose peaks mark object centers, alongside per-object regression targets read back at those peak cells. A ground truth box is splatted as a 2D Gaussian whose radius is chosen so that any box overlapping the true box by at least min_overlap still lands inside the positive region.
References
Objects as Points (the min-overlap Gaussian radius), Center-based 3D Object Detection and Tracking (Yin et al., 2021) (the BEV center-heatmap formulation used by the 3D heads).
Functions:
-
gaussian_radiusβPer-box Gaussian splat radius from the standard three min-overlap cases.
-
draw_gaussian_to_heatmapβSplat a 2D Gaussian at an integer center into a heatmap, max-combining in place.
-
draw_heatmap_targetsβAssign center-heatmap and per-object regression targets for one scene.
-
transpose_gatherβGather per-object channel vectors from a dense map at flat cell indices.
gaussian_radius [source] ΒΆ
Per-box Gaussian splat radius from the standard three min-overlap cases.
Approximates the largest radius \(r\) such that a box overlapping the ground truth by at least min_overlap (IoU) still has its center inside the positive Gaussian region, as the minimum over the inscribed, enclosing, and shifted-box cases. The quadratic roots keep the un-normalized \(r_2, r_3\) of the original Objects as Points formulation (no \(1 / (2a)\) factor), so the radii match the published detectors' training targets rather than the exact closed-form solutions. The formula is symmetric in height and width.
Parameters:
-
height(Tensor) βBox heights (\(y\) extent) in feature-map cells, shape \((N,)\).
-
width(Tensor) βBox widths (\(x\) extent) in feature-map cells, shape \((N,)\).
-
min_overlap(float, default:0.5) βMinimum IoU a candidate box must keep with the ground truth.
Returns:
-
TensorβPer-box radius in feature-map cells, shape \((N,)\) (float; caller rounds and clamps).
Shape
- height: \((N,)\)
- width: \((N,)\)
- output: \((N,)\)
draw_gaussian_to_heatmap [source] ΒΆ
draw_gaussian_to_heatmap(
heatmap: Tensor,
center: Tensor,
radius: Union[int, Tensor],
k: float = 1.0,
) -> Tensor
Splat a 2D Gaussian at an integer center into a heatmap, max-combining in place.
The Gaussian (peak \(k\) at center) is clipped to the heatmap bounds and combined with an element-wise maximum, so overlapping objects keep the stronger response. heatmap is modified in place and also returned. Pass a single channel slice (heatmap[class_id]) to target one class.
Parameters:
-
heatmap(Tensor) βTarget map to draw into, shape \((H, W)\). Modified in place.
-
center(Tensor) βCenter cell \((x, y)\) in feature-map coordinates, shape \((2,)\); truncated to int.
-
radius(Union[int, Tensor]) βGaussian radius in cells (scalar); the splat spans \(2 \cdot \text{radius} + 1\) cells.
-
k(float, default:1.0) βPeak value at the center.
Returns:
-
TensorβThe same
heatmaptensor, modified in place.
Shape
- heatmap: \((H, W)\)
- center: \((2,)\)
- output: \((H, W)\)
draw_heatmap_targets [source] ΒΆ
draw_heatmap_targets(
boxes: Tensor,
labels: Tensor,
num_classes: int,
feature_map_size: Tuple[int, int],
voxel_size: Sequence[float],
point_cloud_range: Sequence[float],
feature_map_stride: int,
*,
num_max_objs: int = 500,
gaussian_overlap: float = 0.1,
min_radius: int = 2,
) -> Tuple[Tensor, Tensor, Tensor, Tensor]
Assign center-heatmap and per-object regression targets for one scene.
Projects each ground truth box center to the BEV feature map, splats a per-class Gaussian, and records the regression target at that peak cell. The regression code is the sub-cell center offset, absolute \(z\), log extents, and \((\cos\theta, \sin\theta)\), followed by any extra box columns (e.g. velocity): \(8 + (D - 7)\) channels for a \((M, D)\) box tensor. Extents are clamped to \(10^{-5}\) before the log so a degenerate box does not produce a non-finite target.
Parameters:
-
boxes(Tensor) βGround truth boxes \((c_x, c_y, c_z, d_x, d_y, d_z, \theta, \ldots)\), shape \((M, D)\), \(D \ge 7\).
-
labels(Tensor) βZero-based class ids, shape \((M,)\).
-
num_classes(int) βNumber of heatmap channels.
-
feature_map_size(Tuple[int, int]) βBEV feature-map size as \((W, H)\) (x then y).
-
voxel_size(Sequence[float]) βVoxel size \((v_x, v_y, v_z)\) in metric units.
-
point_cloud_range(Sequence[float]) βRange \((x_\min, y_\min, z_\min, x_\max, y_\max, z_\max)\).
-
feature_map_stride(int) βStride from voxel grid to feature map.
-
num_max_objs(int, default:500) βCapacity of the per-object target buffers.
-
gaussian_overlap(float, default:0.1) βMin-overlap passed to
gaussian_radius. -
min_radius(int, default:2) βLower clamp on the integer splat radius.
Returns:
-
TensorβTuple
(heatmap, reg_targets, inds, mask): -
Tensorβheatmap: per-class Gaussian map, shape \((\text{num\_classes}, H, W)\).
-
Tensorβreg_targets: regression targets, shape \((\text{num\_max\_objs}, 8 + (D - 7))\).
-
Tensorβinds: flat peak-cell index \(y \cdot W + x\) per object, shape \((\text{num\_max\_objs},)\) (long).
-
Tuple[Tensor, Tensor, Tensor, Tensor]βmask: \(1\) for assigned objects, shape \((\text{num\_max\_objs},)\) (long).
Shape
- boxes: \((M, D)\)
- labels: \((M,)\)
- heatmap: \((\text{num\_classes}, H, W)\)
Example
>>> import torch
>>> boxes = torch.tensor([[0.0, 0.0, -1.0, 4.0, 2.0, 1.5, 0.3]])
>>> labels = torch.tensor([0])
>>> hm, reg, inds, mask = draw_heatmap_targets(
... boxes, labels, num_classes=1, feature_map_size=(16, 16),
... voxel_size=[0.5, 0.5, 0.5], point_cloud_range=[-4.0, -4.0, -2.0, 4.0, 4.0, 2.0],
... feature_map_stride=1,
... )
>>> float(hm.max()), int(mask.sum())
(1.0, 1)
transpose_gather [source] ΒΆ
Gather per-object channel vectors from a dense map at flat cell indices.
Reads the \(C\)-dim vector at each flat cell index \(y \cdot W + x\) (the inds produced by draw_heatmap_targets) out of a dense \((B, C, H, W)\) prediction map, e.g. to compare head outputs against per-object regression targets at the Gaussian peak cells.
Parameters:
-
feat(Tensor) βDense prediction map, shape \((B, C, H, W)\).
-
ind(Tensor) βFlat per-object cell indices, shape \((B, M)\) (long).
Returns:
-
TensorβGathered per-object vectors, shape \((B, M, C)\).
Shape
- feat: \((B, C, H, W)\)
- ind: \((B, M)\)
- output: \((B, M, C)\)