Mix3D: Out-of-Context Data Augmentation for 3D Scenes

In plain words

3D scans of rooms and streets are used to train networks that label every point (wall, chair, car). Such networks pick up habits from their training data: a chair often sits near a table, so a model may guess "chair" partly from the surroundings. Mix3D is a training trick against this. It takes two scenes, moves and turns each one at random, and merges them into a single training example. Objects from one scene then appear in surroundings they never had, so the network cannot lean as much on the setting and has to look at the local shape.

Only the training input changes. The network and the loss stay the same, and every point keeps its own label. Because it works at the data level, it can be added to existing 3D segmentation code. The authors report gains for two network families on indoor and outdoor datasets.

TL;DR

• Mix3D makes a training sample by taking the union of two augmented 3D scenes and concatenating their labels §3.
• The network and cross-entropy loss are unchanged; the paper applies it to MinkowskiNet and KPConv §3 §4.1.
• Author-reported: MinkowskiNet gains on ScanNet, S3DIS and SemanticKITTI, with up to +4.9 mIoU on SemanticKITTI test §4.1 tab. 1.
• The 78.1 mIoU ScanNet test score uses a 4-model ensemble and post-processing, not a single model app. A.
• The paper's own training settings are incomplete (no learning rate). The released code differs from the paper in several augmentation details; see Paper vs code.

Problem

Large 3D scene datasets are small and have little variety. High-capacity segmentation networks with large receptive fields can memorise the global context of the training scenes and overfit to it. The authors argue this fails for rare or unusual arrangements, for example a pedestrian in the middle of a road §1. They also show that training on smaller crops, which removes context, does not help beyond a point §4.2, fig. 4.

Method

Mixing. A sample is built from two scenes, S1 and S2. Each is augmented independently. The two augmented point sets are then combined by union. Points are not interpolated or moved after augmentation, so the labels of the mixed sample are the concatenation of the two label sets §3 fig. 3.

Why keep labels intact. The authors contrast this with MixUp-style methods that interpolate inputs and labels. They argue that a label mix would make the prediction for one scene depend on the other scene, which is the opposite of the goal. They also say PointMixUp distorts point clouds and RSMix keeps only local chunks §3.

Augmentation of each scene (applied before mixing) §3 fig. 3:

1. Subtract the centroid, so that the two scenes overlap after mixing §3.
2. Random horizontal flips in both horizontal directions §3.
3. Random rotation about the up axis, and about the other axis within a small range §3.
4. Random sub-sampling, elastic distortion, and random scaling in U[0.9, 1.1] §3.
5. If colour is an input: random brightness, contrast and colour jitter §3.

Pipeline. A second scene is augmented in parallel and mixed with the first; the mixed scene is the model input. For segmentation the loss is cross-entropy over the predictions and the concatenated labels of both scenes fig. 3 §3.

Models and datasets §4.1 app. A:

• MinkowskiNet (voxel-based sparse convolution) and KPConv (point-based; rigid and deformable variants) §3 §4.1.
• ScanNet V2: 20 classes; public split of 1201 train, 312 validation and 100 test scans §4.1.
• S3DIS: 13 classes; train on all areas except Area 5 §4.1.
• SemanticKITTI: 28 annotated classes, 19 evaluated; 19130 training and 4071 validation scans; single-scan setting §4.1.

Pairing and batching. The augmented scenes are concatenated as consecutive entries of a batch, so the model sees each pair as one sample §3 app. A. To keep the number of points per iteration equal to the baseline, the authors halve the batch size when mixing, for example from 6 to 3 for the 2 cm MinkowskiNet baseline app. A.

Pseudocode (written from the paper text; the released code is cited where it fixes a detail the paper leaves open):

# Mix3D training sample  [§3] [fig. 3]
# Input:  two labelled scenes, each as points (N, 3), colors (N, 3),
#         labels (N,)
# Output: one mixed sample: points, colors, labels

def mix3d_sample(scene_1, scene_2):
    # Augment each scene independently before mixing  [§3]
    points_1, colors_1, labels_1 = augment_scene(*scene_1)
    points_2, colors_2, labels_2 = augment_scene(*scene_2)
    # Mixing is the union of the two point sets  [§3]
    points = concatenate(points_1, points_2)
    colors = concatenate(colors_1, colors_2)
    # Labels follow the same order as the points  [§3] [fig. 3]
    labels = concatenate(labels_1, labels_2)
    return points, colors, labels

# Per-scene augmentation  [§3] [code: mix3d/datasets/semseg.py#L152-L181]
# max_raw is the maximum of the uncentred input coordinates (see Pitfalls)

def augment_scene(points, colors, labels):
    # Centre the scene on its centroid  [§3]
    points = points - points.mean(axis=0)
    # Random shift per axis, drawn from [min, max] / 2  [code: mix3d/datasets/semseg.py#L154]
    points = points + uniform(points.min(0), points.max(0)) / 2
    # Mirror x and y, each with probability 0.5  [§3] [code: mix3d/datasets/semseg.py#L164-L167]
    for axis in (0, 1):
        if random() < 0.5:
            points[:, axis] = max_raw[axis] - points[:, axis]
    # Elastic distortion, probability 0.95, two scales  [code: mix3d/datasets/semseg.py#L168-L172]
    if random() < 0.95:
        for granularity, magnitude in ((0.2, 0.4), (0.8, 1.6)):
            points = elastic_distortion(points, granularity, magnitude)
    # Rotation and scale  [§3] [code: volumentations_aug.yaml]
    points, colors, labels = volume_augment(points, colors, labels)
    # Colour jitter on the colour channels  [§3] [code: albumentations_aug.yaml]
    colors = colour_augment(colors)
    return points, colors, labels

# Voxelisation after mixing  [code: mix3d/datasets/utils.py#L128, #L162-L165]
# voxel_size is 0.05 m for the analysis experiments  [§4.2] [app. A]

def voxelize(points, colors, labels, voxel_size):
    # Quantise each point to its voxel index  [code: mix3d/datasets/utils.py#L162]
    voxels = floor(points / voxel_size)
    # Keep one point per occupied voxel  [code: mix3d/datasets/utils.py#L164-L170]
    keep = one_point_per_voxel(voxels)
    return voxels[keep], colors[keep], labels[keep]

# Training loss on the mixed sample  [§3] [fig. 3]
# The network is unchanged: it sees one input and predicts one label
# per point.

def training_loss(model, points, colors, labels):
    # Predict per-point class scores for the mixed scene  [§3]
    logits = model(points, colors)  # (N1 + N2, num_classes)
    # Cross-entropy over both scenes' points  [§3]
    return cross_entropy(logits, labels)

# Variant with an unlabelled second scene  [§4.2]
# Points from the unlabelled scene stay in the input but get no loss

def mix_with_unlabelled(scene_1, scene_2_points, scene_2_colors):
    # The unlabelled scene gets the ignore label, so it has no loss  [§4.2]
    # [code: mix3d/datasets/utils.py#L105-L109]
    labels_2 = full(len(scene_2_points), IGNORE_LABEL)
    return mix3d_sample(scene_1, (scene_2_points, scene_2_colors, labels_2))

Hyperparameters

Values marked "paper" are stated in the paper; "code" means only the released repository states them. Where a value is missing from both, it is marked not stated.

Results

Author-reported. The baseline is the same network trained without Mix3D, with unchanged training settings §4.1.

Table 1: comparison across datasets (mIoU, mean over three runs with std where given) tab. 1:

Table 2: ScanNet validation, per class tab. 2. Largest gain on the rare "fridge" class: 51.6 to 59.3 for MinkowskiNet (+7.7) and 53.1 to 58.1 for deformable KPConv (+5.0). Overall mIoU values match Table 1.

Table 3: ScanNet test benchmark tab. 3. Mix3D (4-model MinkowskiNet ensemble, see Hyperparameters) scores 78.1 mIoU. The next entries in the table are OccuSeg (76.4) and O-Net (76.2). The table also lists methods that use 2D images or meshes; the Mix3D entry uses 3D input only.

Single-instance evaluation (MinkowskiNet, trained on full scenes, tested on cropped ground-truth instances) tab. 4 §4.2: 24.5 ± 0.88 without Mix3D, 36.0 ± 0.57 with Mix3D (+11.5 mIoU).

Alternative augmentations (ScanNet validation, MinkowskiNet 5 cm; baseline 66.6) tab. 5:

Scene overlap (MinkowskiNet 5 cm, ScanNet validation) tab. 6 §4.2:

Context during training (ScanNet validation; values read from the plot data in the TeX source) fig. 4 §4.2: MinkowskiNet on 1/16 scenes 65.4, on 1/4 scenes 66.9, on full scenes 66.6. KPConv deformable with crop radius 0.5 m 53.8, 1.0 m 67.2, 1.5 m 68.5, 2.0 m 69.3, 3.0 m 69.6.

Learning curves fig. 6 §4.2: with Mix3D the training loss is higher, the validation loss rises later, and mIoU ends higher (three runs).

Supplementary, selected app. B:

• Number of mixed scenes (ScanNet val): 1 → 66.6; 2 → 69.0; 3 → 68.6; 4 → 68.7; 7 → 68.3; 8 → 65.7.
• Training set size (MinkowskiNet, ScanNet): gains of +2.0 to +2.8 mIoU at 5%, 25%, 50%, 75% and 100% of the training set. At 5% the baseline is 43.6 and Mix3D reaches 45.6.
• Single-instance mixing (MinkowskiNet, ScanNet val): overlapping instances 67.7 (+1.0x) and 67.9 (+0.5x); free placement 68.4 (+1.0x) and 69.2 (+0.5x). The paper calls these comparable to Mix3D's three-run spread.

Per class, isolated instances app. C: a Mix3D model beats the baseline on 18 of 20 ScanNet classes when only isolated objects are shown.

Paper vs code

Repository: https://github.com/kumuji/mix3d, linked from the author project page at nekrasov.dev (which redirects to kumuji.github.io/mix3d). The owner is one of the authors. The repository has no LICENSE file; pyproject.toml declares MIT. Code paths below are relative to the repository root at code_commit. The repository also contains two patch files for the original MinkowskiNet and KPConv code, which the README says were used for the ScanNet benchmark and Table 1 README.md. The mix3d/ package is a separate re-implementation with different defaults. I compared both.

• Rotation range. Paper: U[-π/64, π/64] for the non-up axes §3. Released config: U[-π/24, π/24] for those axes volumentations_aug.yaml; the MinkowskiNet patch uses the same π/24 SpatioTemporalSegmentation-ScanNet.patch. The paper's π/64 is not what the released code uses.
• Up-axis rotation. The paper gives no range §3. Code: full ±π about z volumentations_aug.yaml; the patch sets [-π, π] SpatioTemporalSegmentation-ScanNet.patch.
• Augmentation probability. The paper says the transforms are random, without probabilities §3. The released config applies each volume transform with p = 0.5 volumentations_aug.yaml; the MinkowskiNet patch applies the scale and rotations always SpatioTemporalSegmentation-ScanNet.patch.
• Which samples are mixed. The paper mixes batch entries in pairs, for every sample §3. In the released package the pairing is consecutive batch entries with proba = 1 voxelize_collate_merge.yaml utils.py. The MinkowskiNet patch mixes a sample only with probability 0.80, and the partner is drawn uniformly at random from the whole training set SpatioTemporalSegmentation-ScanNet.patch. So the patch leaves about 20% of samples unmixed.
• Second-scene translation. The paper transforms both scenes so that they overlap §3. In the patch, the second scene is centred but its random shift is commented out SpatioTemporalSegmentation-ScanNet.patch; the first scene gets the random shift SpatioTemporalSegmentation-ScanNet.patch. In the released package each scene gets a random shift semseg.py.
• KPConv batches. The supplement says KPConv gets "fully overlapping spheres" in each sample app. B. The TensorFlow patch keeps the last two spheres of each batch unmerged kpconv_tensorflow_mix3d.patch kpconv_tensorflow_mix3d.patch. The PyTorch SemanticKITTI patch leaves about one quarter of the batch (roughly 20%) unmerged kpconv_pytorch_mix3d.patch.
• Far placement. The alternative "far away" setup gives no distance §4.2. The released code offsets the second scene by a uniform vector in [-10, 10] per axis, times 200 utils.py.
• Single-instance mixing. The released code includes it. Instances are centred at random positions drawn from x, y in [-5, 5] and z in [-0.5, 2], with random rotations semseg.py. The paper does not give these ranges app. B.
• Optimizer and schedule. The paper does not state the learning rate or optimizer for the MinkowskiNet runs §4.1 app. A. The released package defaults to AdamW with lr 0.01 and a one-cycle schedule adamw.yaml onecyclelr.yaml for 600 epochs trainer.yaml. The paper's MinkowskiNet runs use 120k iterations §4.1. These are different settings and should not be taken as the paper's.
• Batch size. The paper uses 6 for the 2 cm baseline app. A. The released package defaults to 8 indoor.yaml; the MinkowskiNet patch default is 16 SpatioTemporalSegmentation-ScanNet.patch.
• Models for the released package. Its defaults are Res16UNet34C with first kernel 5 baseline.yaml. The paper's ScanNet benchmark uses Res16UNet34A and 34C app. A.
• Pipeline. The paper says it used the official code for all models §4.1, but for MinkowskiNet on S3DIS and SemanticKITTI it used its own pipeline §4.1. The released package is a re-implementation, not the code used for Table 1 README.md.

No discrepancies were checked for colour jitter magnitudes, since the paper gives none.

Pitfalls

• The 78.1 mIoU benchmark number is not a single-model result. It uses a 4-model ensemble, 40 test-time augmentation runs per model, and graph-based over-segmentation of the predictions app. A. A single-model implementation will not reach it.
• Voxel size changes the numbers. The ScanNet benchmark and Table 1 use 2 cm voxels; the analysis uses 5 cm app. A §4.2. A 2 cm baseline scores 72.4, a 5 cm baseline 66.6, so the two are not comparable.
• Keep points per iteration constant. Each mixed sample has about twice the points of a single scene. Without halving the batch size, the step is roughly twice as expensive app. A.
• The flip mirrors about the uncentred maximum. The released code takes the maximum of the original input coordinates and mirrors the centred coordinates about it semseg.py. The mirrored scene is not re-centred. The random shift that follows partly hides this, but an implementer who re-centres after flipping will get a different sample.
• Sub-sampling is voxel quantisation. The paper says "random sub-sampling" §3. The released code keeps one point per voxel after mixing utils.py, which is not the same as random subsampling.
• Unlabelled-partner variant. Points from the unlabelled scene are kept in the input and excluded from the loss by setting their label to the ignore value utils.py. The paper only says the loss is computed on one scene §4.2.
• Mixed scenes are not realistic. The authors note that the generated scenes are unrealistic and that realism may help §5.
• Noise drop wording. The text says uniform noise lowers mIoU by "more than 6.5%" §4.2. Table 5 gives 60.1 against a baseline of 66.6, a drop of 6.5 mIoU points, so "more than" is slightly off.
• Mixing more than two scenes hurts. Eight scenes give 65.7 against 69.0 for two app. B; keep the pair size at two unless re-tuned.
• Interaction with context. The paper argues that mixing creates overlap of context and that nearby but non-overlapping scenes give a smaller gain §4.2 tab. 6. A pipeline that only places scenes side by side, without overlap, gets the smaller gain.

Lineage

• Builds on: MinkowskiNet (sparse convolution) and KPConv (point-based), the two backbones evaluated §3 §4.1. Neither has a card in this wiki yet.
• Compared against: MixUp and its label-interpolating variants, PointMixUp and RSMix, which the paper contrasts with its design §3; CutOut and random noise §4.2; instance-level mixing app. B.
• Superseded by: not known from content in this wiki.