One Thing One Click: A Self-Training Approach for
Weakly Supervised 3D Semantic Segmentation

To start, we leverage the 3D geometrically homogeneous super-voxels to build a graph. Compared with building on points, our graph has significant fewer nodes to facilitate efficient label propagation.

To derive the prediction $P(y_{j,c}|v_{j},\Theta)$ of the $j$-th super-voxel, we apply a super-voxel pooling to aggregate the semantic prediction of the $n_{j}$ points in $v_{j}$ as below:

where $P({y_{i,c}|p_{i},c_{i},\Theta})$ is the probability of $p_{i}$ in class $c$.

To build the graph, we treat each super-voxel as a graph node and compute the similarity between each pair of super-voxels $v_{j},v_{j^{\prime}}$, which is represented as an edge. Further, to propagate labels to unknown regions through the graph, we formulate it as an optimization problem that considers both the network prediction and similarities among the super-voxels to achieve the global optimum with the energy function below similar to Conditional Random Field (CRF).

where $\mathcal{R}$ is the relation network to be detailed later. The unary term $\psi_{u}(y_{j}|V,\Theta)$ represents the super-voxel pooled prediction of the 3D U-Net $P(y_{j})$ on super-voxel $v_{j}$. Specifically, it denotes the minus $\log$ probability of predicting super-voxel $v_{j}$ to have label $y_{j}$. We define it as below.

The pairwise term $\psi_{p}(j_{k})$ in Equation 3 represents the similarity between super-voxels $v_{j}$ and $v_{j^{\prime}}$. We employ both the low-level features and learned features for measuring the similarity, as shown in Equation 5 below:

where $\mathds{1}(y_{j},y_{j^{\prime}})$ is 1, if $v_{j}$ and $v_{j^{\prime}}$ have different predicted labels, and 0 otherwise. The pairwise term means that the cost will be higher if super-voxels with similar features are predicted to be different classes. Here, $c_{j},c_{j^{\prime}}$, $p_{j},p_{j^{\prime}}$ and $u_{j},u_{j^{\prime}}$ are the normalized mean color, mean coordinates and mean 3D U-Net feature, respectively, of super-voxels $v_{j}$ and $v_{j^{\prime}}$. Unlike existing works [56, 5, 54], which build the graph on 2D image pixels, we build our graph on 3D super-voxels, which have irregular and complex geometrical structures, as shown in the supplementary material. Therefore, hand-crafted features $p_{j},p_{j^{\prime}}$ and $c_{j},c_{j^{\prime}}$ have inferior capability to measure the similarity between super-voxels. To address this issue, we propose the Relation Network to better leverage the 3D geometrical information and explicitly learn the similarity among super-voxels.

Existing works Co-Training [37] and Tri-net [10] showed that semi-supervised training benefits from having two complementary tasks or components. In our framework, we propose a relation net to complement the 3D U-Net. The relation network $\mathcal{R}$ shares the same backbone architecture as the 3D U-Net $\Theta$ except for removing the last category-wise prediction layer. It aims to predict a category-related embedding $f_{j}$ for each super-voxel $v_{j}$ as the similarity measurement. Similar to Equation 2, $f_{j}$ is the per super-voxel pooled feature in $\mathcal{R}$. In other words, the relation network groups the embeddings of same category together, while pushing those of different categories apart. To this end, we propose to learn a prototypical embedding for each category, inspired by the Prototypical Network [40].

However, the per-category prototypes in [40] are fully determined by the sampled mini-batch, and may deviate from the actual categorical center. Consequently, they may not be stable and could keep changing during the training, thereby hard to converge. To assist the training of the relation network with sparse and unbalanced training data, we present a memory bank $K=\{k\}$ to generate one categorical prototype for each category, instead of simply regarding the average embedding as the prototype as in [40].

The embedding $f_{j}$ generated by $\mathcal{R}$ serves as a “query,” and we compare it with the corresponding “key” $k_{c}$ in the memory bank with a dot product. The two modules are optimized simultaneously with contrastive learning [31] as below.

Following [15], we update the key representations via a moving average with momentum as shown below

where $m$ is a momentum coefficient to control the evolving speed of the memory bank. On the one hand, the representations in the memory bank are initialized with random vectors, and are updated during training to generate the prototype for each category. On the other hand, the embeddings generated from the relation network are grouped towards the prototype of its category. In this way, the relation network generates similar embeddings for the same category and distinct ones for different categories. The memory bank updates the prototypes in a category-balanced manner by randomly sampling the same number of training samples $s$ per category in every forward pass.

Our relation net complements with 3D U-Net. It measures the relations between super-voxels using different training strategies and losses, while 3D U-Net aims to project the inputs into the latent feature space for category assignment. The prediction of relation network is further combined with the prediction of 3D U-Net by multiplying the predicted possibilities of each category to boost the performance. In addition, the relation net offers a stronger measurement of the pairwise term in CRF vs. handcrafted features like colors and also complements with the 3D U-Net features.

Table 1 reports the benchmark result on ScanNet-v2 test set. The baselines can be roughly divided into two branches. (i) Fully supervised approaches with 100% supervision, including several representative works in 3D semantic segmentation. These methods are the upper bounds of weakly supervised ones. (ii) Weakly supervised approaches, including a recent work [48].

With less than 0.02% annotated points, our result (69.1% mIoU) outperforms many existing works with full supervision. As for weakly supervised approaches, MPRM [48] is trained with scene-level or subcloud-level labels. The scene-level annotation leads to an inferior performance of 24.4%, and the subcloud-level annotation takes around 3 minutes per scene as reported in [48], which is longer than our “One Thing One Click” scheme (2 minutes). More importantly, our result outperforms [48] by more than 26% mIoU.

In this section, we first present three important baselines as shown in Table 5 on ScanNet-v2 validation set.

• •

  Table 5 “Our fully sup baseline” is trained with the official 100% annotation provided by ScanNet-v2. It serves as the upper bound of our method.

• •

  The model directly trained with the raw annotated points as Figure 3 (c) cannot converge well due to the extreme sparsity of the training data.

• •

  Table 5 “One Thing One Click^∗”. The model trained with the initial pseudo labels as Figure 3 (d) achieves 62.18% mIoU. It serves as the starting point of our self-training approach and is denoted as “our baseline” in the following.

Table 5 “One Thing One Click” manifests that our self-training approach surpasses the baseline by nearly $10\%$ mIoU, attaining a $16\%$ relative improvement. Compared with the fully supervised baseline with the same network architecture, our performance is only 2% lower.

Table 5 “One Thing One Click^†” refers to disabling the graph propagation and relation network in inference. Note that they are still being used in training for generating the pseudo labels. This brings no extra computational burden during the inference, but helps to improve nearly 7% mIoU, comparing with the baseline (68.96% vs 62.18%).

The quantitative results in Figure 5 indicate our result (c) is very similar to the fully supervised baseline (e) [12] in ScanNet-v2. Check error maps (d) (f) for better comparison.

In this section, we show results on ScanNet-v2 “3D Semantic label with Limited Annotations” benchmark [16]. We report the results on the most challenging setting with only 20 points annotated each scene in Table 1 “Ours on Data Efficient” and Table 5 “Data Efficient”. In this experiment, we use the officially provided 20 points instead of “One-Thing-One-Click”, and then employ our self-training approach for semantic segmentation. The results show that our approach is not limited to “One-Thing-One-Click” and is applicable to other annotation schemes. However, the performance is inferior to “One-Thing-One-Click”, since the annotations are more uneven among the categories.

To further study the effectiveness of self-training, graph propagation and relation network, we conduct ablation studies on these three modules on ScanNet-v2 validation set as shown in Table 3 with single view evaluation.

“3D U-Net” indicates that the labels are propagated only based on the confidence score of the 3D U-Net itself, \ie, the unary term in Equation 3. This ablation is designed to manifest the effectiveness of self-training. The “3D U-Net” column in Table 3 manifests that the performance is consistently improved with self-training strategy even without pairwise energy term in Equation 3 and super-voxel partition.

“3D U-Net+GP” refers to the label propagation with graph model, and the similarity among super-voxels are measured by the coordinates $p_{i}$ and colors $c_{i}$ without the learned feature $f_{i}$. This ablation study is to show the effectiveness of the graph model. The results in Table 3 indicate that the graph model benefits the label propagation, and finally boosts the overall performance by 2% over “3D U-Net” (67.92% vs. 65.91%).

“3D U-Net+Rel+GP” utilizes the relation network for similarity measurement based on “3D U-Net+GP”. In this setting, the similarity among super-voxels is measured with the averaged coordinates $p_{i}$, the colors $c_{i}$, the unary features $u_{i}$, and the relation network generated feature $f_{i}$, as shown in Equation 3. This experiment is to manifest that the relation network benefits the similarity measurement and pseudo label generation, compared with the hand-crafted feature, i.e., coordinates and color. It outperforms the hand-crafted features especially in the later iterations since the network benefits from the richer pseudo labels. It finally achieves 2.5% improvement compared with “3D U-Net+GP” (70.45% vs. 67.92%). As shown in Figure 4, the generated pseudo labels for each iteration expands to unknown regions step by step and finally gets close to the ground truth.

Further, we study whether the learned embeddings of the relation network outperform the hand-crafted features for similarity measurement. We randomly sample 200 super-voxels for each category in ScanNet-v2, and conduct a t-SNE visualization [27] on them. Figure 6 indicates that the relation network better groups the intra-class embeddings and distinguish the inter-class embeddings compared with hand-crafted features.

We also compare with fully supervised approaches and weakly supervised approaches on S3DIS. The latter includes existing works [21, 42] and recent works [52, 47].

As shown in Table 4, with the “One Thing One Click” scheme where less than 0.02% points are annotated, we achieve 50.1% mIoU. With “One Thing Three Clicks” scheme, our performance can be further improved to 55.3% mIoU. The above two results outperform [52] by 5.6% and 10.8% mIoU (0.2% annotations in [52]), and 2.1% and 7.3% mIoU (10% annotations in [52]) respectively.

Wang et al. [47] unprojects 2D semantic labels to 3D space for 3D semantic segmentation. To compare with [47], we first compare with the actual number of annotated points regardless of 2D or 3D. For S3DIS, the number of annotated 2D pixels (70,496 images with 1080$\times$1080 resolution) is 100$\times$ more than the officially annotated 3D points ($5.27\times 10^{8}$ in total), so both settings of [47] (100% 2D annotations and 16.7% 2D annotations) actually utilize a large quantity of annotations. Even with a large gap of annotation, the results in Table 4 show that our “One Thing Three Clicks” scheme with only 0.06% 3D annotation outperforms [47] with 100% 2D annotations by nearly 3% mIoU.

In addition, our approach achieves comparable results with several fully supervised methods as shown in Table 4.

We follow the similar settings in Section 4.1 to show several baselines for S3DIS.

• •

  Table 4 “Our fully-sup baseline”. The model trained with the full supervision of S3DIS achieves 63.7% mIoU. It serves as the upper bound of our approach.

• •

  The model directly trained with only the annotated points in Figure 3 (c) cannot converge well.

• •

  Table 4 “One Thing One Click^∗” and “One Thing Three Clicks^∗”. The model trained with the annotated super-voxels in Figure 3 (d) achieves 43.7% mIoU for “One Thing One Click” and 48.9% mIoU for “One Thing Three Clicks”. They are used as the baselines to calculate the “relative improvement” of our approach, and are denoted as “our baseline” in the following.

As shown in Table 4 “Rel. Imp.” column, we have 14.6% (“One Thing One Click”) and 13.1% (“One Thing Three Clicks”) relative improvement over our baseline, surpassing the relative improvement of [52], which is 1.1% (with 0.2% annotations) and 5% (with 10% annotations) over their own baselines, by a large margin. The significant improvement of “relative improvement over baseline” manifests the effectiveness of the proposed approach.

To evaluate without any extra computation burden, we further disable the label propagation and relation network in inference as shown in Table 4 “^†”. Note that they are still adopted in training. Our model still attains 13.0% (“One Thing One Click”) and 10.6% (“One Thing Three Clicks”) relative improvement over our baseline in this case.