anchors

Anchor-based dense detection heads for the voxel detectors (PointPillars, SECOND).

A packed-format port of the anchor head from open-mmlab/OpenPCDet.

generate_anchors: axis-aligned anchor generation; residuals are decoded with decode_box_residuals.
AnchorHeadSingle: the single-stage anchor head (per-anchor class logits, box residuals and a direction bin).
AnchorHeadMulti: the multi-group separate-head variant (sincos + velocity box code) used by the nuScenes detectors.
separate_branch: the per-attribute SeparateHead branch builder, also used by the Voxel Mamba center head.

Classes:

AnchorHeadOutput –

Raw and decoded predictions of AnchorHeadSingle.
AnchorTargets –

Per-anchor training targets from assign_anchor_targets.
AnchorHeadSingle –

Single-stage anchor head (AnchorHeadSingle).
AnchorHeadMultiOutput –

Predictions of AnchorHeadMulti.
MultiGroupSingleHead –

One RPN head of AnchorHeadMulti.
AnchorHeadMulti –

Multi-group anchor head (AnchorHeadMulti, separate-multihead).

Functions:

generate_anchors –

Generate axis-aligned anchors for a single class over a BEV feature map.
assign_anchor_targets –

Assign classification and box-regression targets to a single class group of axis-aligned anchors.
separate_branch –

Build a SeparateHead-style prediction branch: middle conv blocks, then a plain output conv.

AnchorHeadOutput ¶

Bases: TypedDict

Raw and decoded predictions of AnchorHeadSingle.

AnchorTargets ¶

Bases: TypedDict

Per-anchor training targets from assign_anchor_targets.

AnchorHeadSingle ¶

AnchorHeadSingle(
    input_channels: int,
    num_classes: int,
    grid_size: Tuple[int, int],
    point_cloud_range: Sequence[float],
    *,
    anchor_sizes: Sequence[Sequence[float]],
    anchor_bottom_heights: Sequence[float],
    feature_map_stride: int,
    anchor_rotations: Sequence[float] = (0.0, 1.57),
    num_dir_bins: int = 2,
    dir_offset: float = 0.78539,
    dir_limit_offset: float = 0.0,
)

Bases: Module

Single-stage anchor head (AnchorHeadSingle).

Three \(1\times1\) convs predict, per anchor, class logits, 7-DoF box residuals and a direction bin. At inference the residuals are decoded against the precomputed anchors and the predicted heading is snapped to the predicted direction bin (dir_offset / num_dir_bins).

Parameters:

input_channels (int) –

Channels of the BEV feature map fed to the head.
num_classes (int) –

Number of foreground classes.
grid_size (Tuple[int, int]) –

Full voxel grid size \((n_x, n_y)\) (before the head feature-map stride).
point_cloud_range (Sequence[float]) –

Range \((x_\min, y_\min, z_\min, x_\max, y_\max, z_\max)\).
anchor_sizes (Sequence[Sequence[float]]) –

Per-class box size \((d_x, d_y, d_z)\), shape \((\text{num\_classes}, 3)\).
anchor_bottom_heights (Sequence[float]) –

Per-class anchor bottom \(z\), shape \((\text{num\_classes},)\).
anchor_rotations (Sequence[float], default: (0.0, 1.57) ) –

Yaw angles (radians) shared by all classes.
feature_map_stride (int) –

BEV feature-map stride of the head.
num_dir_bins (int, default: 2 ) –

Number of direction bins.
dir_offset (float, default: 0.78539 ) –

Direction-classifier angle offset.
dir_limit_offset (float, default: 0.0 ) –

Offset used when wrapping the decoded heading before snapping.

Methods:

generate_predicted_boxes –

Decode raw head outputs into per-anchor class logits and absolute boxes.
decode –

Decode a forward output into raw per-anchor detections (no score threshold or NMS).

generate_predicted_boxes ¶

generate_predicted_boxes(
    batch_size: int,
    cls_preds: Tensor,
    box_preds: Tensor,
    dir_cls_preds: Tensor,
) -> Tuple[Tensor, Tensor]

Decode raw head outputs into per-anchor class logits and absolute boxes.

Returns:

Tensor –

A tuple (batch_cls_preds, batch_box_preds) of shapes \((B, A, \text{num\_classes})\)
Tensor –

(raw logits, not sigmoided) and \((B, A, 7)\) where \(A\) is the number of anchors.

decode ¶

decode(out: AnchorHeadOutput) -> Detection3D

Decode a forward output into raw per-anchor detections (no score threshold or NMS).

Scores each anchor by its top sigmoid class probability and labels it by the argmax class. The full per-anchor set is returned; the evaluation pipeline applies score thresholding and per-class 3D NMS via the torch_pointcloud.utils.box3d utilities (see the benchmark examples).

Returns:

Detection3D –

Packed per-anchor detections `{"boxes": (B * A, 7), "scores": (B * A,), "labels": (B * A,),
Detection3D –

"batch": (B * A,)}` (PyG layout).

AnchorHeadMultiOutput ¶

Bases: TypedDict

Predictions of AnchorHeadMulti.

MultiGroupSingleHead ¶

MultiGroupSingleHead(
    input_channels: int,
    num_classes: int,
    num_anchors_per_location: int,
    code_size: int,
    reg_list: Sequence[str],
    head_label_indices: Tensor,
    *,
    num_middle_conv: int = 1,
    num_middle_filter: int = 64,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

One RPN head of AnchorHeadMulti.

A SeparateHead-style head over the shared feature: a classification branch plus one regression branch per box-code group (reg, height, size, angle, velo), whose outputs are concatenated into the per-anchor box code. Predictions are reshaped to the multihead layout \((B, A, \cdot)\) with \(A\) the anchors of this head's class group.

Parameters:

input_channels (int) –

Channels of the shared feature map.
num_classes (int) –

Number of classes handled by this head (separate-multihead).
num_anchors_per_location (int) –

Anchors per BEV cell for this head.
code_size (int) –

Box code size (e.g. 10 for sincos angle + velocity).
reg_list (Sequence[str]) –

Regression-branch spec, e.g. ["reg:2", "height:1", "size:3", "angle:2", "velo:2"].
num_middle_conv (int, default: 1 ) –

Number of middle convs per branch.
num_middle_filter (int, default: 64 ) –

Middle-conv channel width.
act (Union[str, Callable, None], default: 'relu' ) –

Activation type or callable for the middle convs.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra activation arguments.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization type or callable for the middle convs.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra normalization arguments.

AnchorHeadMulti ¶

AnchorHeadMulti(
    input_channels: int,
    num_classes: int,
    grid_size: Tuple[int, int],
    point_cloud_range: Sequence[float],
    *,
    anchor_sizes: Sequence[Sequence[float]],
    anchor_bottom_heights: Sequence[float],
    head_class_groups: Sequence[Sequence[int]],
    feature_map_stride: int,
    anchor_rotations: Sequence[float] = (0.0, 1.57),
    shared_conv_num_filter: int = 64,
    reg_list: Sequence[str] = (
        "reg:2",
        "height:1",
        "size:3",
        "angle:2",
        "velo:2",
    ),
    num_middle_conv: int = 1,
    num_middle_filter: int = 64,
    code_size: int = 9,
    encode_angle_by_sincos: bool = True,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
)

Bases: Module

Multi-group anchor head (AnchorHeadMulti, separate-multihead).

A shared conv feeds several MultiGroupSingleHeads, one per class group. Anchors (7-DoF, padded to the box-code size) are decoded with decode_box_residuals (sincos heading, velocity deltas); per-head class scores stay separate (with their global label mapping) for class-wise NMS downstream.

Parameters:

input_channels (int) –

Channels of the BEV feature map fed to the head.
num_classes (int) –

Number of foreground classes.
grid_size (Tuple[int, int]) –

Full voxel grid size \((n_x, n_y)\).
point_cloud_range (Sequence[float]) –

Range \((x_\min, y_\min, z_\min, x_\max, y_\max, z_\max)\).
anchor_sizes (Sequence[Sequence[float]]) –

Per-class box size \((d_x, d_y, d_z)\), shape \((\text{num\_classes}, 3)\).
anchor_bottom_heights (Sequence[float]) –

Per-class anchor bottom \(z\), shape \((\text{num\_classes},)\).
head_class_groups (Sequence[Sequence[int]]) –

Class-index groups, one per RPN head (e.g. [[0], [1, 2], ...]); the classes in each group share one SeparateHead.
anchor_rotations (Sequence[float], default: (0.0, 1.57) ) –

Yaw angles (radians) shared by all classes.
feature_map_stride (int) –

BEV feature-map stride of the head.
shared_conv_num_filter (int, default: 64 ) –

Channels of the shared conv.
reg_list (Sequence[str], default: ('reg:2', 'height:1', 'size:3', 'angle:2', 'velo:2') ) –

Regression-branch spec for each head.
num_middle_conv (int, default: 1 ) –

Middle convs per branch.
num_middle_filter (int, default: 64 ) –

Middle-conv channel width.
code_size (int, default: 9 ) –

Base box code size (9 for nuScenes; +1 internally for sincos).
encode_angle_by_sincos (bool, default: True ) –

Encode heading as \((\cos, \sin)\).
act (Union[str, Callable, None], default: 'relu' ) –

Activation type or callable for the shared conv and head middle convs.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra activation arguments.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization type or callable for the shared conv and head middle convs.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra normalization arguments.

Methods:

decode –

Decode a multihead forward output into raw per-anchor detections (no score threshold or NMS).

decode ¶

decode(out: AnchorHeadMultiOutput) -> Detection3D

Decode a multihead forward output into raw per-anchor detections (no score threshold or NMS).

Each head scores its anchors by their top sigmoid class probability and maps the argmax to the global label; the per-head results are concatenated in head order (matching batch_box's anchor order). When the box code carries velocity deltas the decoded \((v_x, v_y)\) columns are returned under velocity. The full per-anchor set is returned; the evaluation pipeline applies score thresholding and per-class 3D NMS via the torch_pointcloud.utils.box3d utilities (see the benchmark examples).

Returns:

Detection3D –

Packed per-anchor detections `{"boxes": (B * A, 7), "scores": (B * A,), "labels": (B * A,),
Detection3D –

"batch": (B * A,)}(PyG layout), plus"velocity"` \((B \cdot A, 2)\) when the head predicts it.

generate_anchors ¶

generate_anchors(
    point_cloud_range: Sequence[float],
    feature_map_size: Tuple[int, int],
    anchor_sizes: Sequence[Sequence[float]],
    anchor_rotations: Sequence[float],
    anchor_bottom_heights: Sequence[float],
    *,
    dtype: dtype = float32,
) -> Tensor

Generate axis-aligned anchors for a single class over a BEV feature map.

Mirrors the reference AnchorGenerator with align_center=False: anchor centers are placed on a grid spanning point_cloud_range (endpoints inclusive), then anchor_bottom_heights are lifted by half the box height to box centers.

Parameters:

point_cloud_range (Sequence[float]) –

Range \((x_\min, y_\min, z_\min, x_\max, y_\max, z_\max)\).
feature_map_size (Tuple[int, int]) –

BEV feature map size \((n_x, n_y)\).
anchor_sizes (Sequence[Sequence[float]]) –

Box sizes \((d_x, d_y, d_z)\), one row per size template.
anchor_rotations (Sequence[float]) –

Yaw angles (radians).
anchor_bottom_heights (Sequence[float]) –

Anchor bottom \(z\) per height template.
dtype (dtype, default: float32 ) –

Anchor dtype.

Returns:

Tensor –

Anchors \((1, n_y, n_x, n_\text{size}, n_\text{rot}, 7)\) as \((x, y, z, dx, dy, dz, \theta)\).

assign_anchor_targets ¶

assign_anchor_targets(
    anchors: Tensor,
    gt_boxes: Tensor,
    gt_labels: Tensor,
    *,
    matched_threshold: float,
    unmatched_threshold: float,
    match_height: bool = False,
) -> AnchorTargets

Assign classification and box-regression targets to a single class group of axis-aligned anchors.

Each anchor is matched to the ground-truth box of highest IoU: IoU \(\ge\) matched_threshold makes it a positive carrying that box's label, IoU \(<\) unmatched_threshold makes it background, and anything in between is ignored. Each ground-truth box additionally force-matches its single highest-IoU anchor, so a box with no anchor above threshold still receives one positive. Positive anchors' regression targets are the residual encoding of their matched box against the anchor (the inverse of decode_box_residuals).

Callers with several class groups (one anchor set per class) invoke this once per group with that class's anchors, ground truth, and thresholds, then concatenate the results.

Parameters:

anchors (Tensor) –

Anchors \((x, y, z, d_x, d_y, d_z, \theta)\) for one class group, shape \((A, 7)\).
gt_boxes (Tensor) –

Ground-truth boxes \((c_x, c_y, c_z, d_x, d_y, d_z, \theta)\), shape \((G, 7)\).
gt_labels (Tensor) –

Ground-truth class labels (\(1\)-based foreground indices), shape \((G,)\).
matched_threshold (float) –

IoU at or above which an anchor becomes a positive.
unmatched_threshold (float) –

IoU below which an anchor becomes background.
match_height (bool, default: False ) –

Match by 3D IoU when True, otherwise bird's-eye IoU.

Returns:

AnchorTargets –

A TypedDict with cls_labels \((A,)\) (\(-1\) ignore, \(0\) background, \(\ge 1\) foreground class) and
AnchorTargets –

box_reg_targets \((A, 7)\) (residual encodings, zero for non-positive anchors).

Shape

anchors: \((A, 7)\)
gt_boxes: \((G, 7)\)
gt_labels: \((G,)\)
cls_labels: \((A,)\)
box_reg_targets: \((A, 7)\)

Example

>>> anchors = torch.tensor([[0.0, 0.0, 0.0, 4.0, 2.0, 1.5, 0.0], [20.0, 0.0, 0.0, 4.0, 2.0, 1.5, 0.0]])
>>> gt_boxes = torch.tensor([[0.0, 0.0, 0.0, 4.0, 2.0, 1.5, 0.0]])
>>> gt_labels = torch.tensor([1])
>>> out = assign_anchor_targets(anchors, gt_boxes, gt_labels, matched_threshold=0.6, unmatched_threshold=0.45)
>>> out["cls_labels"].tolist()
[1, 0]

separate_branch ¶

separate_branch(
    in_channels: int,
    out_channels: int,
    num_middle_conv: int,
    num_middle_filter: int,
    *,
    act: Union[str, Callable, None] = "relu",
    act_kwargs: Optional[Dict[str, Any]] = None,
    norm: Union[str, Callable, None] = "batch_norm",
    norm_kwargs: Optional[Dict[str, Any]] = None,
    bias: bool = False,
) -> Sequential

Build a SeparateHead-style prediction branch: middle conv blocks, then a plain output conv.

The per-attribute branch shared by the separate detection heads (the anchor multi-head and the center head): num_middle_conv blocks of (\(3\times3\) conv, norm, act) followed by a \(3\times3\) output conv.

Parameters:

in_channels (int) –

Input channels.
out_channels (int) –

Output channels of the final conv.
num_middle_conv (int) –

Number of middle conv blocks.
num_middle_filter (int) –

Channel width of the middle convs.
act (Union[str, Callable, None], default: 'relu' ) –

Activation of the middle conv blocks.
act_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra activation arguments.
norm (Union[str, Callable, None], default: 'batch_norm' ) –

Normalization of the middle conv blocks.
norm_kwargs (Optional[Dict[str, Any]], default: None ) –

Extra normalization arguments.
bias (bool, default: False ) –

Whether the middle convs carry a bias (the output conv always does).

Returns:

Sequential –

The branch as an nn.Sequential.

Shape

Input: \((B, C_\text{in}, H, W)\)
Output: \((B, C_\text{out}, H, W)\)

Example

>>> branch = separate_branch(64, 2, num_middle_conv=1, num_middle_filter=64)
>>> branch(torch.rand(2, 64, 16, 16)).shape
torch.Size([2, 2, 16, 16])