Sparse Single Sweep LiDAR Point Cloud Segmentation via Learning Contextual Shape Priors from Scene Completion

In plain words

A self-driving car's laser scanner returns a thin, uneven set of dots for each sweep. Far-away objects and small things such as bicycles or people often get only a handful of dots, so labelling each dot correctly is hard.

The authors' idea is to use the sequence of sweeps during training. Many neighbouring sweeps are stitched together into one dense picture of the scene, which shows the full shape of each object. A second part of the network learns to fill in the complete 3D shape from the single sparse sweep. That shape knowledge is then used to label the dots of the single sweep more accurately.

The filling-in part is only a training aid. It is dropped when the model runs, so the segmentation part runs as fast as it would without it.

TL;DR

JS3C-Net is a LiDAR segmentation network for one sweep at a time. It is trained together with a scene-completion decoder whose targets come from many merged sweeps. Both tasks are linked by a point-voxel interaction module. At inference only the segmentation branch runs. The authors report 66.0 mIoU on the SemanticKITTI benchmark and 60.2 average IoU on SemanticPOSS, compared against the baselines in their tables.

Problem

Single-sweep LiDAR points are sparse, unevenly spaced and noisy. Distant objects lose most of their shape, which makes segmentation of them hard §1. Earlier methods that exploit sequences look only at past frames, cannot use information about newly appearing objects, and add a feature-aggregation cost that is a poor fit for real-time driving §1 §2.

Method

Overview. A sparse segmentation network gives each point a label and a shape embedding. A cascaded completion branch takes these outputs and predicts a dense voxel volume of the whole scene. A point-voxel interaction (PVI) module passes information from the incomplete points to the complete voxels and back. The completion branch is used only in training §3.1 §3.6.

Segmentation backbone. The input is a sparse voxel grid at 0.05 m, built with a hash table so that only occupied voxels are computed. Its features come from a sparse U-Net, which is mapped back to points. The segmentation head adds a shape embedding (SE) to the U-Net output and then predicts 19 class logits §3.2 SubSparseConv.py. The paper states that only point coordinates are used as input features §3.2. The released code uses the LiDAR remission value as the single input feature, with coordinates only for placement kitti_dataset.py. See Paper vs code.

Completion branch (cascaded). Segmentation output and SE are fed to the completion branch per point, then averaged into the voxels they fall in §3.3 train.py. The completion volume covers 51.2 m forward, ±25.6 m sideways and 6.4 m high, at 0.2 m, i.e. 256×256×32 voxels §4.1. The branch is a voxel encoder-decoder with a stride-2 convolution, several 32-channel residual stages, multi-scale concatenation and a dense upsampling step back to full resolution §3.3 SSCNet.py. It predicts 19 semantic classes plus empty, i.e. C+1 outputs with C semantic classes §3.3.

Point-voxel interaction. For each coarse voxel, the module finds its nearest input voxels, computes a learned weight for each neighbour and pools their features. It then fuses the pooled result with the coarse features to refine the completion output §3.4 app. A.1. The paper gives a DGCNN-style edge feature (eq. 1) for the neighbour relation. The code uses a different relation, described under Paper vs code.

e_{ij} = \phi\big([p^v_i, f^v_i],\ [p^v_i, f^v_i] - [p_j, f_j]\big) eq. 1

Here p^v_i is a coarse voxel centre, p_j an input point, f the features, [·,·] concatenation, and φ a shared MLP (the paper's word) §3.4.

Joint loss. Segmentation and completion are both weighted cross-entropy losses. They are combined with learned uncertainty weights σ1 and σ2 §3.5:

\mathcal{L}(W,\sigma_1,\sigma_2) = \frac{1}{2\sigma_1^2}\mathcal{L}_{seg}(W_{seg}) + \frac{1}{2\sigma_2^2}\mathcal{L}_{complet}(W) + \log\sigma_1 + \log\sigma_2 eq. 2

The paper says segmentation gradients reach only the segmentation parameters, while both losses update the shared network §3.5.

Training step. Written from the paper and checked against the released training script. Each line below is a paraphrase, so it does not reproduce the code.

# Joint training step, segmentation plus completion  [§3.1]
# Input:  points (N, 3); remission (N, 1); seg_labels (N,) in 0..18
#         scene_labels (256, 256, 32) in 0..19; invalid (256, 256, 32)
# Output: scalar loss

def training_step(points, remission, seg_labels, scene_labels, invalid):
    # Sparse U-Net on 0.05 m voxels, remission as the only input feature
    # [§3.2] [code: models/SubSparseConv.py#L21-L34]
    unet_feat = sparse_unet(points, remission)               # (N, 16)

    # Shape embedding, added back to the U-Net output before the classifier
    # [§3.2] [code: models/SubSparseConv.py#L35-L46]
    shape_emb = shape_mlp(unet_feat)                         # (N, 16)
    seg_logits = seg_linear(unet_feat + shape_emb)           # (N, 19) logits

    # Completion input: raw logits and shape embedding for each point
    # [§3.3] [code: train.py#L116-L119]
    point_feat = concat(seg_logits, shape_emb)               # (N, 35)

    # Mean over up to 5 points stored in each non-empty 0.2 m voxel
    # [code: models/model_utils.py#L76-L83]
    # [code: opt/JS3C_default_kitti.yaml#L39]
    voxel_feat = mean_over_voxel_points(point_feat, fuse_k=5)  # (V, 35)

    # Random flip, applied to features and labels with the same state
    # [code: kitti_dataset.py#L131-L136] [code: kitti_dataset.py#L215-L232]
    voxel_feat, scene_labels, invalid = random_flip(
        voxel_feat, scene_labels, invalid)

    # Coarse completion: stride-2 conv to 0.4 m, then five 32-channel stages
    # [§3.3] [code: models/SSCNet.py#L16-L83]
    coarse_logits = ssc_decoder(dense(voxel_feat))  # (128, 128, 16, 20)

    # Point-voxel interaction with k = 8 nearest input voxels
    # [§3.4] [code: models/model_utils.py#L137-L208]
    fine_logits = pvi(coarse_logits, voxel_shape_emb, k=8)  # (128, 128, 16, 20)

    # Upsample to 0.2 m: 1x1x1 conv to 8x channels, then pixel shuffle
    # [code: models/SSCNet.py#L98-L99] [code: models/SSCNet.py#L119]
    completion = upsample_pixel_shuffle(fine_logits)  # (256, 256, 32, 20)

    # Weighted cross-entropy; unlabeled points carry label -100 and are ignored
    # [§3.5] [code: kitti_dataset.py#L89-L90] [code: train.py#L217]
    loss_seg = weighted_cross_entropy(seg_logits, seg_labels, w_seg)

    # Weighted cross-entropy over the volume; invalid voxels are ignored
    # [code: models/model_utils.py#L112-L121]
    scene_labels[invalid] = 255
    loss_comp = weighted_cross_entropy(
        completion, scene_labels, w_comp, ignore=255)

    # Uncertainty weighting. The code uses 2 x eq. 2 with the same minimum
    # [eq. 2] [code: models/model_utils.py#L124-L130]
    loss = (loss_seg / sigma_seg ** 2 + loss_comp / sigma_comp ** 2
            + 2 * log(sigma_seg) + 2 * log(sigma_comp))
    return loss

Inference (segmentation only). The completion branch and PVI are not needed. The paper describes test-time voting with random augmentations §3.6 §4.2. The released code runs the segmentation head alone and adds up the logits of each augmented copy test_kitti_segment.py test_kitti_segment.py.

# Segmentation-only inference with test-time voting  [§3.6] [§4.2]
# Input:  points (N, 3); remission (N, 1); votes (code default 10)
# Output: class index per point (N,), classes 0..18

def predict(points, remission, votes=10):
    # Accumulate logits per original point over all augmented copies
    # [code: test_kitti_segment.py#L120] [code: test_kitti_segment.py#L194]
    logit_sum = zeros(N, 19)
    for vote in range(votes):
        # Random mirror, rotation, scale jitter and shift for this copy
        # [code: test_kitti_segment.py#L137-L142]
        augmented = random_augment(points, remission)
        # Only the segmentation head runs at inference [§3.6]
        # [code: test_kitti_segment.py#L73-L76]
        logits = seg_head(augmented)                         # (N', 19)
        logit_sum = add_back_to_points(logit_sum, logits)
    return argmax(logit_sum, axis=1)

Hyperparameters

Values in the paper and the released code often differ. Both are listed.

Results

Author-reported. Setup: SemanticKITTI official benchmark for Table 1 and the SSC table (single-sweep test submission app. A.2). SemanticPOSS split is not stated. Ablations use the SemanticKITTI validation split §4.5.

Other tables in the paper: tab. 1 has per-class IoU for 19 classes; tab. 2 has the 11-class POSS table; tab. 5 lists PointNet++, TangentConv, RandLA-Net and KPConv latencies. Historical-frame comparisons are in app. B.3: a SparseConv baseline at 63.1 mIoU, a SpSequenceNet variant at 64.3, and JS3C-Net with past frames only at 66.6, with past and future frames at 67.5 (validation).

Paper vs code

Code: official repository yanx27/JS3C-Net at d505d0b, MIT, by the first author LICENSE. Compared file by file with the paper.

• Input features. Paper: only point coordinates are used §3.2. Code: remission is the only feature, and use_coords is off kitti_dataset.py JS3C_default_kitti.yaml.
• Completion input. Paper: the segmentation output is the "semantic probability" §3.3. Code: raw linear logits with no softmax, concatenated with SE SubSparseConv.py train.py.
• PVI features. Paper: PVI uses SE §3.4. Code: PVI sees only the 16 SE channels; the 19 logits are dropped there SSCNet.py.
• PVI edge function. Paper: a DGCNN-style MLP over concatenated features and differences (eq. 1). Code: a 1×1 conv on a 10-value relation of positions and distance, multiplied with the neighbours' features, then averaged. The optional feature-relation term is off model_utils.py model_utils.py JS3C_default_kitti.yaml.
• PVI residual. Paper: the refined feature is added back to the coarse output §3.4. Code: the pooled features and coarse features are concatenated and passed through a fully connected layer to 20 channels model_utils.py model_utils.py.
• Completion stem. Paper: a convolution, then a pooling layer §3.3. Code: a stride-2 7×7×7 convolution; the pooling layer is defined but never called SSCNet.py SSCNet.py.
• Stage concatenation. Paper: features from different scales are concatenated §3.3. Code: stages 3–5 are concatenated at the same resolution; stages 4–5 use dilation SSCNet.py SSCNet.py.
• Uncertainty loss. Paper eq. 2 has 1/(2σ²) and log σ. Code uses 1/σ² and 2 log σ, which is exactly twice eq. 2 model_utils.py. σ is initialised uniformly in 0.2 to 1 train.py. Paper: σ trained at 10 × LR §4.2. Code: base LR train.py.
• Training schedule. See Hyperparameters: epochs, batch size and the LR decay interval all differ.
• Segmentation augmentation. Paper: rotation about y and scale 0.9–1.1 §4.2. Code: rotation about z (vertical in the LiDAR frame), a random 3×3 matrix I + N(0, 0.1²) scaled by the voxel scale, and a ±1 mirror on x kitti_dataset.py.
• Completion augmentation. Paper: flips along x and z, and 90° rotation about y §4.2. Code: flips along the two horizontal voxel axes only, chosen from 4 states, and the 90° rotations are commented out kitti_dataset.py kitti_dataset.py.
• Voting. Paper: averaging of predictions §4.2. Code: logits are summed over 10 votes and then the argmax is taken test_kitti_segment.py test_kitti_segment.py.
• Vertical shift. Code adds +1 voxel to the vertical index for SemanticKITTI, as a workaround for the old data version; the paper does not mention it train.py [README].
• Model selection. Code keeps the checkpoint with the best validation segmentation mIoU train.py. The paper does not say how the reported checkpoint was chosen.
• Released checkpoint. The file name model_segiou_0.6327_compltiou_0.2350_epoch79.pth records 63.27 segmentation mIoU and 23.50 SSC mIoU on validation [log/JS3C-Net-kitti]. The paper's full model is 67.5 and 24.0 on the same split tab. 4. The README lists 66.0 and 56.6, which are benchmark numbers.
• Completion branch data. Paper: 70 consecutive frames are merged §4.1. Repository: no merging code. Labels come from an external voxelizer with "future scans: 70" [README], which suggests future frames only. Paper does not say.

Pitfalls

• Input feature. The paper says coordinates only. Code and the released run use remission. Implement the feature the code uses, and confirm it before comparing numbers.
• Completion input is logits. Paper says probability. Code feeds raw logits. Adding a softmax changes the PVI and SSC inputs.
• Neighbourhood and graph depth are not in the paper. k is not stated; the code uses k = 8 with one relation step JS3C_default_kitti.yaml app. A.1.
• Loss scale. The code's loss is twice eq. 2. With Adam this has a small effect, but a value copied from the paper's loss will not match the logged loss.
• Learned σ has no lower bound. Code does not clamp σ. As σ approaches 0, 1/σ² grows quickly model_utils.py.
• Class weights are steep. Completion weights are (max frequency / frequency)^(1/3), from the config counts. Empty voxels get 1; the rarest completion class gets about 40 JS3C_default_kitti.yaml kitti_dataset.py. Segmentation weights range from 1 to about 14.5.
• Dense completion cost. The decoder runs dense 3D convolutions over the whole 256×256×32 grid at training time SSCNet.py. Training memory grows with batch size. The "disposable" claim covers inference only §3.6.
• Shape embedding stays at inference. The SE branch lives in the segmentation head when interaction is on, so it runs at inference. Only the completion decoder and PVI are dropped SubSparseConv.py test_kitti_segment.py.
• Labels are external. The repository does not generate scene labels. Training needs SemanticKITTI voxel files, or the voxelizer with the README settings [README].
• Paper arithmetic. The text says the ablation gain from uncertainty weighting is 0.4% for segmentation §4.5. Table 4 gives 66.1 → 66.4, which is 0.3.
• Two versions of Table 3. The main-text JS3C row has precision 71.5 and recall 73.5. The appendix table has 70.2 and 74.5. Both have IoU 56.6 and mIoU 23.8 tab. 3 app. A.3.
• Abstract gains use the baseline. The "4% and 3%" in the abstract match the gains over the authors' SparseConv baseline (61.8 → 66.0 and 57.2 → 60.2). Over KPConv they are 7.2 and 5.0 points.
• Complexity text disagrees with the table. The text says one fifth of KPConv's size and one third of its time §4.5. Table 5 gives 2.7 M against 18.3 M (about one seventh) and 471 ms against 1117 ms (about 0.42).
• Moving objects. Dynamic objects reconstruct as streaks in the merged volume. Removing them changes segmentation by 0.2 points on the validation split app. B.2. The authors kept them in training to match the benchmark.
• Section numbers. The LaTeX sets secnumdepth to 0, so the compiled paper may have unnumbered sections. The § numbers above count sections and subsections in source order. Appendix tables restart at 1 in the source; those are cited by appendix section.

Lineage

• Builds on: submanifold sparse convolution as the segmentation backbone §3.2. SSCNet for semantic scene completion §2. Uncertainty weighting for multi-task losses (Kendall et al., 2018) §3.5. DGCNN edge features §3.4.
• Compared against: KPConv, RandLA-Net, PointASNL, TangentConv, RangeNet++, SqueezeSegV2 and V3, PolarNet tab. 1. TS3D, TS3D², TS3D³ and EsscNet for completion tab. 3. SpSequenceNet for historical frames app. B.3.
• Superseded by: not stated in this wiki.