Active Self-Training for Weakly Supervised 3D Scene Semantic Segmentation

We adopt the 3D U-Net architecture [3] as the backbone of our segmentation network. The input to U-Net is a point cloud $X$ of $N$ points, with each point $x_{i}$ containing both 3D coordinates and color information, where $i\in\{1,...,N\}$. The network predicts the probability of each semantic category for each point $x_{i}$, denoted as $p_{i,c}$, and the probability corresponding to the ground truth category $\bar{c}$, provided either by $T$ or $P$, is denoted as $p_{i,\bar{c}}$. The network is then trained with the softmax cross-entropy loss:

where $\lambda$ is a combination weight.

In the first iteration of the self-training, the network is trained with the set $T$ derived from a set of randomly sampled points $\hat{T}$ annotated by the user, and the set $P$ is empty. In the subsequent iterations, the set $\hat{T}$ is expanded with new samples annotated by users, where the samples are selected via active learning as explained in Section 3.2, which are then propagated through the super-voxels to form a new set $T$. The set $P$ is updated with the prediction results of the current network, as explained in Section 3.3.

Note that, during inference, the average of predictions of all the points inside the same super-voxel is used as the prediction result of the super-voxel, so that all the points in the same super-voxel have the same prediction. This “voting method” ensures that the prediction for the points inside each super-voxel is consistent, which improves the final prediction accuracy, as we show in our ablation studies.

During the active learning process, the user is asked to annotate a sparse set of points with labels, illustrated in Fig. 2 (d). To select the most effective set of points to annotate for improving the accuracy of the segmentation, we measure the uncertainty of the labeling of each point based on the current prediction results. The uncertainty of each point is measured by calculating the standard deviation of several stochastic forward passes, and using the one corresponding to the category with the highest mean prediction confidence. More specifically, for each input point cloud, we first get $K$ different versions via standard data augmentation operations [3], and then for each point, we compute the mean and standard deviation of these $K$ probability distributions predicted from those $K$ different input versions. Finally, the standard deviation of the category with the highest mean probability is used as the uncertainty $u_{i}$ of the point $x_{i}$:

where $p_{i,c}^{k}$ is the predicted probability of point $x_{i}$ in the $k$-th point cloud version for category $c$ and $\overline{p}_{i,c}=\frac{\sum_{k}p_{i,c}^{k}}{K}$ is the mean probability for the $K$ versions for category $c$, with $\hat{c}$ being the category with highest mean probability.

For each iteration, we select $m$ points according to the uncertainty distribution over points, since intuitively points with high uncertainty require more reliable user input.

Similar to [17], we iteratively update the set of pseudo labels $P$. Starting with the label predictions of the segmentation network trained at a given iteration, we take the predictions with high confidence (larger than a given threshold $\tau$) and use them as updated pseudo labels $P$. The pseudo labels $P$ are then used together with the labels $T$, propagated from the true labels $\hat{T}$, to train the network in the next iteration. Note that we also use the mean prediction probability $\overline{p}_{i,\hat{c}}$ of $K$ different point cloud versions, as in Eq.(2), to compute the confidence in this step.

To limit error propagation in our iterative training and pseudo-labeling process, we generate new labels for all unlabeled samples and reinitialize the neural network after each pseudo-labeling step, similarly to Rizve et al. [21].

Note that, for a fair comparison, we use less or the same amount of user annotations as the existing methods. The actual amount of user annotations used is reported in all the tables. As our method uses active learning to select samples to annotate, the number of samples is evenly distributed. For example, if $n$ is the total amount of user annotations and $k$ is the number of iterations, then $m=n/k$ points are sampled in each iteration for users to annotate.

Our method produces very competitive results with only 20 labeled points per scene. Firstly, our method outperforms the best weakly supervised approaches on the “Data Efficient” setting with 20 points by nearly 11% mIoU (70.3% vs. 59.4%). Our method also surpasses 1T1C by nearly 1% mIoU (70.3% vs. 69.1%) when 1T1C uses twice the amount of user annotations, i.e., 0.02% vs. 0.01% (20 points/scene), and requires object instance information. Secondly, the gap between our method and full supervision is less than 3.3% mIoU (70.3% vs. 73.6%), which shows the effectiveness of our method.

Fig. 3 shows visual examples of results obtained with our method and the fully-supervised baseline, compared to the ground truth. We can see that our method obtains comparable, sometimes even better, results than the fully-supervised baseline. For example, for the scene shown in the first row, our method is able to correctly label all the chairs in different arrangements. Our method is also able to recognize the cabinet for the scene shown in the second row, although the fully-supervised baseline misclassifies it as a counter.

We can see that with only 20 labeled points per scene (0.01% supervision), our method outperforms 1T1C (0.02% supervision) by nearly 6% mIoU. It also surpasses the 1T3C instance that annotates 3 random points per object and results in 0.06% supervision. In addition, our approach even outperforms several fully-supervised methods, which again demonstrates the effectiveness of our method.

Fig. 3 shows visual results of our method on the S3DIS dataset. We see that even with more challenging and crowded scenes, our method can still obtain reasonably good prediction results similar to those obtained via the fully-supervised baseline.

Our ablation studies are conducted on the most challenging setting with only 20 points annotated in each scene, for both the ScanNet-v2 and S3DIS datasets. For ScanNet-v2, the evaluation is conducted on the validation set. For S3DIS, the evaluation is conducted on Area 5.

We test the two key components of our method: the self-training (denoted “Self-train.”) and the active learning (“Active learn.”). We also test the “voting” process used during the inference. The results are shown in Table 3. We see that each component of our method contributes to the quality of the final results.

More specifically, in the first row, we show the results of our baseline method, which is the segmentation network [3] trained with all the user annotations, where the samples are either selected with a pre-trained model provided by [9] for ScanNet-v2 or randomly selected for S3DIS. By adding “voting” of predictions within each super-voxels, we can see that the results improve, which confirms that maintaining label consistency inside of each super-voxel contributes to an improvement of the results. When adding one of our two key components individually in the third and forth row, we see that the performance is improved on both datasets in either case, while the performance is the best when using the full method with both key components.

Fig. 4 shows statistics on the classes of objects where the points were selected with the active learning in our method, compared to the method of Hou et al. [9]. We see that our method selects less samples on floors while selecting more samples on categories like wall, door, window, desk, counter, and sink, which leads to significant improvement of prediction accuracy for these categories. Our method even improves the prediction accuracy for floors as this class often gets misclassified as “door” in early-iteration results, while in the final results, there is less confusion with doors as our method selects more samples for doors.

We also show a visual comparison of selected samples in Fig. 4. For our method, we show the uncertainty map over the scene, which is used to guide the sample selection, together with the selected samples for each iteration. Comparing to the samples selected at once with the pre-trained model in [9], we see that the samples that we select are located on more complicated objects with more visual variability like chairs rather than the floor.

Note that, if we have $s$ scenes in the training set, there are two ways for active learning to sample the points. One way is to sample the same number of points in each scene, i.e., $n/(k*s)$ per scene, and the other way is to sample the points among all the scenes based on uncertainty only, which results in different numbers of sampled points in different scenes. We tested these two different options and found that the results are similar.

This is because both datasets have large scene variations, which leads to samples evenly distributed over the scenes even if they are sampled over the entire dataset based on pointwise uncertainty. If another dataset with scenes similar to each other were given, we believe that the results with sampling over all the scenes would be better.

In this section, we test our method under the “one thing one click” setting. The only change in our method is that, for the sample selection during the active learning, we avoid objects that have been sampled before, selecting the sample with highest uncertainty among the remaining objects. In other words, our active learning chooses which object to sample as well as which point inside each object should be sampled in the each iteration.

In more detail, we first compute the average number $n_{o}$ of objects per scene for both datasets, which results in $n_{o}=32$ for ScanNet-v2 and $n_{o}=36$ for S3DIS. Here we set the number of iterations to be $k=6$. For each of the first five iterations, we select $6$ samples from $6$ different objects in the scene and set them to be invalid for the selection in the next iteration. In the final iteration, we select one point with highest uncertainty from each remaining object to complete the “one thing one click” sample selection.

Table 4 shows how the performance changes across the iterations. We see that, for the S3DIS dataset, our method already obtains better results than 1T1C [17] when only $24$ points were annotated in iteration 4. Fig. 5 shows a visual comparison of a sample point selected by our method and the point randomly sampled as in [17]. We can see that our method learns to select points that are located in more important regions inside the objects, which is more informative and leads to more accurate prediction results. By focusing on labeling more challenging regions, our method can correctly predict other unlabeled regions, while with random samples as in [17], the method may not be able to extract enough information to correctly predict the labels for the challenging regions.

One interesting result that we observed is that, for both the ScanNet-v2 and S3DIS datasets, after adding $24$ points while constraining only one point per object, the results are worse than for our method under the 20 points/scene setting, where less annotations are given. To find the reason behind this behavior of the method, in Fig. 6, we show the statistics of number of points sampled on each object for our method under the 20 points/scene setting. We see in the results for ScanNet-v2 that, compared to the “one thing one click” setting where $100\%$ of the objects get exactly one click, most of the objects ($61\%$) in our setting do not get any samples, $14\%$ of the objects get more than one click, and only $25\%$ of the objects get exactly one click. We observe similar distributions on S3DIS. This indicates that the “one thing one click” procedure may not be a good way to sample points for semantic segmentation. Complicated objects may need more sample points than others, while for objects coming from simpler categories, we do not require annotations for each scene.

In this supplementary document, we first present more details of our ActiveST framework in Section A. Then, we report detailed benchmark results on the ScanNet-v2 test set with per-category performance of our method comparing to other weakly supervised methods on the most challenging setting with only 20 points annotated in each scene in Section B. Finally, we show more results on both ScanNet-v2 and S3DIS datasets with various numbers of annotated points in Section C.