Bayesian Self-Training for Semi-Supervised 3D Segmentation

TL;DR

This paper introduces a Bayesian self-training framework for semi-supervised 3D segmentation tasks, leveraging uncertainty estimation to improve pseudo-label quality and achieve state-of-the-art results across multiple datasets.

Contribution

It proposes a novel Bayesian self-training approach that effectively utilizes unlabeled data for 3D segmentation, extending to instance segmentation and 3D visual grounding.

Findings

Achieves state-of-the-art results on SemanticKITTI and ScribbleKITTI for 3D semantic segmentation.

Significantly improves semi-supervised 3D instance segmentation on ScanNet and S3DIS.

Enhances dense 3D visual grounding performance over supervised baselines.

Abstract

3D segmentation is a core problem in computer vision and, similarly to many other dense prediction tasks, it requires large amounts of annotated data for adequate training. However, densely labeling 3D point clouds to employ fully-supervised training remains too labor intensive and expensive. Semi-supervised training provides a more practical alternative, where only a small set of labeled data is given, accompanied by a larger unlabeled set. This area thus studies the effective use of unlabeled data to reduce the performance gap that arises due to the lack of annotations. In this work, inspired by Bayesian deep learning, we first propose a Bayesian self-training framework for semi-supervised 3D semantic segmentation. Employing stochastic inference, we generate an initial set of pseudo-labels and then filter these based on estimated point-wise uncertainty. By constructing a heuristic $n$-partite matching algorithm, we extend the method to semi-supervised 3D instance segmentation, and finally, with the same building blocks, to dense 3D visual grounding. We demonstrate state-of-the-art results for our semi-supervised method on SemanticKITTI and ScribbleKITTI for 3D semantic segmentation and on ScanNet and S3DIS for 3D instance segmentation. We further achieve substantial improvements in dense 3D visual grounding over supervised-only baselines on ScanRefer. Our project page is available at ouenal.github.io/bst/.