BADet: Boundary-Aware 3D Object Detection from Point Clouds

In plain words

Self-driving cars often find objects in a laser scan (a point cloud) in two steps. First a network proposes many candidate boxes around possible objects. Then a second step looks at each box alone and adjusts it. Many of the proposals sit a little off the true object, so a single box often cannot see the whole object.

BADet changes the second step. It treats the nearby proposals as nodes in a small graph and lets them pass messages to each other, so each box gathers context from its neighbours and can infer where the object's edges are. Each node also gets richer features by combining three views of the scene: 3D voxels, a flat bird's-eye-view map and the raw points.

The idea is that a box does not have to be right on its own: its neighbours together reveal the object's boundary. The authors evaluate it on the KITTI and nuScenes driving benchmarks.

TL;DR

BADet is a LiDAR 3D detector that refines proposals jointly instead of one at a time. Proposals within a cut-off distance form a graph, and graph message passing (three iterations in the released code) updates each proposal's features. A lightweight Region Feature Aggregation module supplies those features from voxel-wise, pixel-wise and point-wise sources. Results are author-reported. The paper says BADet ranked first on the KITTI BEV leaderboard for Car at Moderate difficulty as of April 2021.

Problem

Two-stage 3D detectors refine each region proposal independently, and proposals tend to deviate from the ground truth. Each refined box then lacks the full boundary of its object, and the authors argue that existing networks have no mechanism to compensate for this §1 fig. 1. Taking RoI features from a bird's-eye-view map alone also loses information: downsampling dilutes the features, and squeezing out the depth axis drops 3D context §3.

Method

BADet has three main parts: a backbone with a region proposal network (RPN), the RFA module, and a boundary-aware graph detector called BA-Refiner §3 fig. 2.

1. Voxelization. Points are mapped to a regular grid with quantization step [v_L, v_W, v_H]. Points beyond a per-voxel capacity are dropped at random. Each voxel feature is the mean of its points' (x, y, z, reflectance) §3.1.1.
2. 3D backbone and RPN. Sparse 3D convolutions with batch normalization and ReLU produce feature volumes. The volumes are reshaped to a bird's-eye-view (BEV) map, a 2D convolutional network processes it, and two sibling 1×1 convolutions output proposals per location §3.1.2 §4.1.2.
3. RFA module. Each proposal's feature comes from three sources §3.2:
   • Voxel-wise: 3D sparse features from three stages are interpolated back to the raw points. An auxiliary head predicts foreground and centre offsets from these features, and the head is then discarded. The interpolated features are combined and mapped to the voxel-wise component §3.2.1 eq. 1.
   • Pixel-wise: for each proposal, an m₁ × m₂ grid is laid over its BEV footprint, and BEV features are sampled at the grid points §3.2.2.
   • Point-wise: a small PointNet++-style network is trained from scratch on the raw points and summarised into keypoints. Their features are interpolated to the proposal centres §3.2.3. The three components are concatenated into the node state x_i eq. 1.
4. Graph construction. Two proposals are connected when their centres are closer than a cut-off r. Each node is connected to itself §3.3.1 eq. 2.
5. Graph update. Each step updates every node from its neighbours. For neighbour j of node i, the message is a function of the relative position x_i − x_j, minus a predicted alignment offset Δx_i, concatenated with the neighbour's state. Messages are max-aggregated over the neighbourhood, and the node state gets a residual update §3.3.2 alg. 1 eq. 3. The paper runs this step T times; it reports T = 3 for the test-server submission and T = 4 as best on the validation split §4.4.4.
6. Heads. Two branches with stacked MLPs output a class score and an oriented 3D box per node §3.3.3.
7. Losses. The overall loss is L = L_rpn + L_gnn + L_offset + L_seg eq. 6. Stage one and the graph head use focal loss for classification and smooth-L1 for regression §3.4. The offset loss eq. 4 and foreground-segmentation loss eq. 5 supervise the voxel-wise auxiliary head.

# BADet inference for one point cloud  [§3] [fig. 2]
# Input:  points (N, 4): x, y, z, reflectance
# Output: oriented boxes (K, 7) with scores (K,)

def badet_inference(points):
    # Voxelize, then run the 3D sparse-conv backbone  [§3.1.1]
    voxel_feats, voxel_coords = voxelize(points)
    volume = sparse_conv_backbone(voxel_feats, voxel_coords)

    # Flatten depth into a BEV map and predict proposals  [§3.1.2]
    bev, bev_conv6 = bev_network(volume)
    proposals, proposal_scores = rpn_head(bev)

    # Build one node state per proposal from three sources  [§3.2] [eq. 1]
    centers = proposals[:, :3]
    pixel_feats = sample_bev_grid(bev_conv6, proposals)    # (K, 28)
    voxel_feats = interpolate(centers, volume_points)      # (K, 64)
    point_feats = interpolate(centers, keypoint_feats)     # (K, 28)
    states = concat(pixel_feats, voxel_feats, point_feats) # (K, 120)

    # Connect proposals whose centres are within the cut-off  [§3.3.1] [eq. 2]
    edges = radius_graph(centers, cutoff)                  # cutoff: not stated

    # Refine node states by message passing  [§3.3.2] [alg. 1]
    for step in range(T):                                  # T = 3 in code
        states = graph_update(states, centers, edges)

    # Score and regress every node  [§3.3.3]
    class_scores, boxes = graph_heads(states, proposals)
    return nms(boxes, class_scores)                        # IoU 0.1 on KITTI

# One extended graph update step  [alg. 1] [eq. 3]
# states (K, 120); positions (K, 3); edges: (i, j) pairs, self-loops included

def graph_update(states, positions, edges):
    # Alignment offset predicted for each centre node  [eq. 3]
    offsets = offset_mlp(states)                           # (K, 3)
    for centre i in nodes:
        # One message per neighbour j: relative position and its state  [eq. 3]
        messages = [
            message_linear(
                concat(positions[i] - positions[j] - offsets[i], states[j])
            )
            for j in neighbours(i)
        ]
        # Max over the neighbourhood, then a residual update  [alg. 1]
        pooled = channelwise_max(messages)
        states[i] = states[i] + activation(project(pooled))
    return states

Key equations, in the paper's notation:

\mathcal{X} = \{ x_i = (x_i, [f^{(voxel)}_i, f^{(pixel)}_i, f^{(point)}_i]) : i = 1..n \} eq. 1

\mathcal{E} = \{ (i, j) \mid \lVert x_i - x_j \rVert_2 < r \} eq. 2

\mathbf{h}^{k}_{\mathcal{N}(i)} = \max_{j \in \mathcal{N}(i)} \Big( \mathbf{W}^{k}_{g} \big[\, x_i - x_j - \Delta x^{k}_i ;\ \mathbf{h}^{k-1}_j \big] \Big) eq. 3

The paper typesets the aggregation subscript as k. Algorithm 1 defines the aggregator as a channel-wise max over neighbours, with k indexing the iteration, so the max is over j.

\mathcal{L}_{offset} = \frac{1}{N_{pos}} \sum_{i}^{N} \mathcal{L}_{smooth\text{-}L1}(\Delta\hat{d} - \Delta d^{*}) \cdot \mathbb{1}[b^{*}_i \ge 1] eq. 4

\mathcal{L}_{seg} = \frac{1}{N_{pos}} \sum_{i}^{N} -\alpha (1 - \hat{p}_i)^{\gamma} \log(\hat{p}_i) eq. 5

\mathcal{L} = \mathcal{L}_{rpn} + \mathcal{L}_{gnn} + \mathcal{L}_{offset} + \mathcal{L}_{seg} eq. 6

Hyperparameters

Values from the paper are cited to the paper. Values marked "code" appear only in the official repository at the commit above and are not stated in the paper.

Results

Author-reported. Setup: KITTI test (online leaderboard, Car, 40 recall positions, IoU 0.7, submitted 2021) and nuScenes val (trained on the train split) §4.2 §4.3.

KITTI test, Car (AP in %, 40 recall positions, IoU 0.7). Source: tab. 1.

The paper reports BADet on par with Voxel R-CNN for 3D (−0.01 points at Moderate) and first on the BEV Moderate leaderboard at the date of submission §4.2.2.

KITTI val, Car (AP3D in %, 11 recall positions, IoU 0.7). Source: tab. 2.

nuScenes val (trained on the train split). Source: tab. 3 (AP per class and mAP) and tab. 4 (NDS and error terms).

Per-class AP for BADet on nuScenes val tab. 3: Car 80.30, Pedestrian 80.20, Bus 57.10, Barrier 57.50, Traffic cone 58.20, Truck 43.70, Trailer 25.00, Motorcycle 36.00, Construction vehicle 11.40, Bicycle 27.10. The paper notes BADet is not best on Car, Bus, Truck, Trailer and Construction vehicle §4.3.2.

Ablations (KITTI val, Car Moderate, AP3D, 11 points, IoU 0.7) tab. 5 tab. 6 tab. 7:

Runtime (batch 1, Intel i7-6900K CPU and GTX 1080 Ti GPU, KITTI val, 3,769 samples) §4.5 tab. 8: average 140.36 ms overall, of which RPN 43.65 ms, RFA 26.73 ms, BA-Refiner 61.98 ms and NMS 1.49 ms. The paper reports 7.1 FPS against Point-GNN's 1.5 FPS tab. 9.

Paper vs code

Code: https://github.com/rui-qian/BADet at a2c260449ca355b526a9642a7bfe9073b22891e6. The repository is linked from the paper's abstract and the authors are its owners, so it belongs to the authors. Its only KITTI config is configs/car_cfg.py. There is no nuScenes config and no nuScenes dataset code, so the nuScenes settings above cannot be checked against code.

• Score threshold. Paper: 0.3 §4.1.4. Code: 0.32 car_cfg.py.
• Graph iterations. Paper: T = 3 for the test submission, T = 4 best on val §4.4.4. Code: fixed at 3 for both single_stage.py.
• Relative offset sign. Paper Eq. 3: x_i − x_j − Δx_i, with i the centre node eq. 3. Code: x_j − x_i − Δx, with Δx predicted from the centre node's state gnn.py. Paper describes the alignment offset as predicted for each neighbour §3.3.2, while code predicts it once per centre node.
• Message and update shape. Paper Algorithm 1: a single linear map W_g, then max, then a residual update with a non-linearity σ(W_f ·) alg. 1. Code: a three-layer MLP on each message before the max, and a three-layer MLP after the max, added residually with no outer non-linearity gnn.py gnn.py.
• Graph radius and neighbour cap. Paper: r unspecified §3.3.1. Code: 0.1 m, with at most 256 neighbours per node gnn.py.
• Pixel-wise sampling. Paper: "nearest pixel" interpolation with a spatial transformer sampler §3.2.2. Code: bilinear grid_sample ssd_rotate_head.py.
• 2D BEV backbone. Paper: six 3×3 convolutions §4.1.2. Code: seven 3×3 convolutions (conv0 to conv6) and a 1×1 conv7, with the conv6 output used for pixel-wise sampling cmn.py cmn.py.
• Auxiliary loss normalisation. Paper: each of L_offset and L_seg is normalised by N_pos eq. 4 eq. 5. Code: normalised by the positive count and then divided again by the batch size cmn.py.
• Focal and smooth-L1 constants. Paper: α, γ and the smooth-L1 β are not stated. Code: α = 0.25, γ = 2.0 losses.py, β = 1/9 ssd_rotate_head.py.
• Training additions. Paper: SGD, cosine annealing §4.1.3. Code adds momentum 0.9, a 500-iteration linear warmup, gradient clipping at norm 10, and ground-truth rotation and centre-noise augmentation car_cfg.py car_cfg.py.
• Second-stage matching. Paper does not state the refiner's assignment thresholds. Code: 3D IoU 0.7 for both positive and negative car_cfg.py.
• Direction classifier. Paper does not mention one. Code enables it car_cfg.py.
• Test-time refiner gate. In the test path the BA-Refiner runs only when len(bbox_score) == 0 single_stage.py. The paper does not describe this condition, and I did not trace the empty-proposal case in detail.
• Table 5 text. The text in §4.4.3 gives +5.81 (3D Moderate), +0.13 (3D Hard), and +0.11, +0.12, +0.08 (BEV Easy, Moderate, Hard) for adding the BA-Refiner §4.4.3. Table 5 rows 3 to 4 give +5.83, +0.15, and −0.03, +0.10, +0.05. The table values are shown above.

No discrepancies were found for the anchor sizes, the voxel grid, the point range, the matching IoU thresholds, the 3D augmentation ranges or the batch size and epochs.

Pitfalls

• Unstated graph constants. The cut-off r, the neighbour cap, the pixel grid size m₁ × m₂ and the proposal selection threshold are not in the paper. An implementation needs the code values: 0.1 m, 256 neighbours, 4 × 7, and a guided-anchor threshold of 0.1 single_stage.py.
• Iteration count. The reported val optimum is T = 4, but the released config runs T = 3 §4.4.4 single_stage.py. Results at T = 4 would need a code change.
• Sensitivity to T. Val AP3D Moderate falls from 85.77 at T = 3 to 79.87 at T = 5, and the authors attribute this to gradient vanishing without testing it §4.4.4.
• Sign and order in the graph update. Paper and code differ in the sign of the relative position and in which node's state predicts the offset (see Paper vs code). The two forms are not interchangeable. Follow the code to reproduce reported numbers, and the paper for its stated formula.
• Runtime sums. The four listed components sum to 133.85 ms against the stated overall average of 140.36 ms. About 6.5 ms are unaccounted for §4.5 tab. 8. The per-component percentages also do not quite match: NMS is written as 1.04 % but 1.49 / 140.36 is 1.06 %.
• Point and voxel caps. Code caps voxels at 5 points each and 20,000 voxels car_cfg.py. The paper states neither cap, only that points are dropped at random when over capacity §3.1.1.
• Metric variants. KITTI test uses 40 recall positions and val uses 11. Numbers are not directly comparable across the two splits §4.2.1. The paper attributes the val-test gap to distribution mismatch, not measured §4.2.2.
• Time-bound claims. The "first place" BEV claim refers to the leaderboard as of April 2021 §4.2.2.
• Test-set training. For the test-server submission the model was trained on train plus val, using the last epoch's weights §4.2.
• Figure and table numbers. Numbers in this card follow the LaTeX source order (Table 1 to Table 9, Figure 1 to Figure 4, Equations 1 to 9). They should be checked against the compiled PDF before publication.
• Publisher notice. The source carries an Elsevier © 2022 notice for the journal version. The license used is the arXiv CC BY 4.0 license recorded for version v8.

Lineage

• Builds on: Point-GNN for the graph update and max aggregation, with alignment offsets §3.3.2. SA-SSD for the backbone and the pixel-wise sampling §3.1.2 §3.2.2. PointNet++ for feature propagation and the point-wise network §3.2.3.
• Compared against: PointRCNN, PV-RCNN, Voxel R-CNN, STD, Part-A², 3DSSD and CIA-SSD on KITTI §4.2.2; SECOND, PointPillars and 3DSSD on nuScenes §4.3.2.
• Superseded by: not stated.