Boundary-Aware Geometric Encoding for Semantic Segmentation of Point Clouds

First, we automatically annotate each point in training samples as its indicator of the boundary $g$, which is defined in accordance with the label of every point as below. In the target boundary, $g_{i}$ is $0$ if the $i^{th}$ point is on the boundary, otherwise equal to $1$. For every point $p$, whether it is located on the boundary is determined by its local neighborhood. That is, given fixed number of neighboring points for $p$, if there are more than a predefined ratio (detailed description is in experiments) of points that do not belong to the same category as $p$, then $p$ is assumed to be the point on the boundary, otherwise it is not.

The boundary prediction task is a slightly different from semantic segmentation, as boundary prediction should be aware of the difference of semantic information in a local region. To this end, as shown in Fig. 2. (b), we collect features of $k$ nearest neighbors in the local region for each point and take the variance of collected features as the input of the following part of BPM. Then, like PointNet (Qi et al. 2017a), we utilize several shared MLPs to predict the boundary annotation $\hat{g}$ for the whole input point cloud. Compared with a carefully designed network, our BPM is compact and easy to train. Specifically, its training loss is following:

$$\mathcal{L}_{BPM}=-\sum_{i=1}^{n}(w_{1}\cdot g_{i}\log\hat{g_{i}}+w_{2}\cdot(1-g_{i})\log(1-\hat{g_{i}})),$$

(1)

where $w_{1}$ and $w_{2}$ are used to balance the huge difference between the numbers of two categories. We also utilize cross-entropy loss to regularize the final semantic segmentation output, and the total loss is a simple addition of boundary prediction loss and semantic segmentation loss.

As mentioned above, in the proposed boundary-aware GEM (Fig. 2 (c)), we attempt to block the propagation of local features from points on the boundary in the early stage of encoding process. Therefore, according to the predicted boundary, we utilize the boundary information as a mask/filter to assign different weights to different points during feature aggregation. Before that, we also utilize the GCO (detailed description will be given later) to provide extra geometric features. The main difference between boundary-aware GEM decoder and encoder is that, we do not use the GCO in decoder as the output features of corresponding encoder which have already contained the geometry information will be concatenated to the input of the decoder.

Given the predicted boundary points (the red points in Fig. 2. (c)), during feature aggregation for a grey point, it will collect features in a neighborhood but ignore those points on the boundary. Therefore, the local feature aggregation for point $p_{i}$ can be expressed as follows:

$$f_{en\_\,l}=\sigma(\mathcal{A}(\{\phi(r_{ij})\cdot\mathcal{M}(\hat{g_{j}}\cdot f_{p_{j}})\})),\ \forall p_{j}\in\mathcal{N}(p_{i}),$$

(2)

where $f_{p_{j}}$ represents the feature of neighboring $p_{j}$ containing both original features and geometric features, and $\hat{g_{j}}$ works as a mask to assign weight to $f_{p_{j}}$. In this formula, $\mathcal{N}(p_{i})$ is the neighborhood of $p_{i}$, and $\mathcal{M}$ means shared MLPs to combine the original features and extracted geometric features at this scale. Referring to Fig. 2 (c), we can know $\phi$ learns weight from the relative position $r_{ij}$ for neighbor $p_{j}$ through another few MLPs. Additionally, $\mathcal{A}$ is the aggregation function that is done through matrix product and $\sigma$ represents the activation function. It is noteworthy that $\hat{g_{j}}$ is 0 if $p_{j}$ is on the boundary and this boundary point would not contribute to the aggregated feature.

In this way, we prevent the features of points on the boundary to be fused into the extracted local features, thus information is less likely to cross over the boundary to pollute features belonging to other categories (shown in Fig. 2. (c)). We only need to predict the boundary of point clouds for the input layer, while in the later encoding stages, points and predicted boundary labels are down-sampled at the same time. Unlike local features in the first few layers, the global features can propagate among different objects through boundary points. Therefore, in the latter stage, we extract global features as follows:

$$f_{en\_\,h}=\sigma(\mathcal{A}(\{\phi(r_{ij})\cdot\mathcal{M}(f_{p_{j}})\})),\ \forall p_{j}\in\mathcal{N}(p_{i}).$$

(3)

In the decoding stage, the feature extraction procedure is symmetrical. Specifically, when the number of points remains small, global features propagate without impeding to better recognize the global context. While in the later stage of the decoder, we prevent the propagation of features across the boundary again to obtain distinguishing local features.

In our method, we propose a geometric kernel $K_{geo}$ with three directional vectors $\{v_{1},v_{2},v_{3}|v_{i}\in\mathbb{R}^{3}\}$. Each vector represents a direction in the 3D space, thus the kernel itself can describe a distribution of points over directions, so as to tell where the points is located (e.g., on a plane or curved surface). Unlike (Shen et al. 2018; Thomas et al. 2019) which employ a large amount of kernel points, only three 3-D directional vectors are adopted in our method. Even though the proposed operation has a much simpler structure, the performance is comparable with some sophisticated operators, that is proved in ablation study. Because, tetrahedron is the simplest polyhedron and these three directional vectors along with the origin can represent a tetrahedron. Furthermore, more complex geometry pattern can be recognized through hierarchical geometric feature extraction. Fig. 3 illustrates learnt kernels and heat maps for different objects to show effectiveness.

For a point in the point cloud, the local pattern is represented by the relative positions from this point to its neighbors. Similar to 2D convolution, if the geometric pattern of the neighborhood is very similar to the learnt GCO kernel, the response will be large, thus geometric pattern is recognized.

Our geometric convolution focuses more on the angular distributions of neighbors rather than their relative displacement like KCNet (Shen et al. 2018). For every point $p_{i}$, relative positions of three neighbors, that are used to represent the local pattern, are represented by $\{\vec{d}_{ij}|\ j\in[1,2,3]\}$. Convolving with $K_{geo}$, the output could be expressed as

$$O_{i}=\mathop{\max}_{\mathcal{P}_{i}}\sigma(b+\mathop{\sum}_{j}\vec{d}_{ij}\cdot\vec{v}_{{\mathcal{P}_{i}}(j)}),$$

(4)

where $b$ is the bias and $\sigma$ is the activation function. $\mathcal{P}_{i}(\cdot):\{1,2,3\}\mapsto\{1,2,3\}$ represents a mapping function which finds the matching vector in the kernel for $\vec{d}_{ij}$. It is noteworthy that because the point cloud is unordered, it is hard to use a fixed mapping. Additionally, if $K_{geo}$ describes the same pattern as the neighborhood, each pair of $\vec{d}_{ij}$ and the matching $\vec{v}_{{\mathcal{P}_{i}}(j)}$ would be in the same direction making the dot production maximum. Therefore, in our proposed convolution procedure, we dynamically choose the mapping function that makes the output maximum. Obviously, our geometric convolution is more sensitive to the angle between two vectors $cos\langle\vec{d}_{ij},\vec{v}_{{\mathcal{P}_{i}}(j)}\rangle$ rather than the displacement between vectors in the neighborhood and kernel $|\vec{d}_{ij}-\vec{v}_{{\mathcal{P}_{i}}(j)}|$, which is more easily be influenced by scales and density of point clouds.

After extracting geometric features, they are concatenated to the original features of points for further boundary-aware geometric encoding (Fig. 2 (c)), making points with different geometry more distinguishable. In the encoder, geometric patterns can be learnt from different scales, thus complex geometric pattern can be represented by the combination of geometrical features of different scales.

In scene semantic segmentation task, we evaluate our method on ScanNet v2 (Dai et al. 2017) and S3DIS (Armeni et al. 2016). In ScanNet v2, there are totally $1,201$ scanned scenes for training and $312$ scenes for validation. Additionally, another 100 scenes are provided as the testing samples, and there are 20 different categories. Following (Wu, Qi, and Fuxin 2019), we randomly sample $3m\times 1.5m\times 1.5m$ cubes from rooms with 8,192 points as the training samples, and test over the entire scan. In S3DIS, there are six indoor areas including 271 rooms from three different buildings. Each point is annotated with a corresponding label from 13 categories. We split points by room and sample all rooms into $0.5m\times 0.5m$ blocks with $0.25m$ padding. Like experiment setting used in previous works (Qi et al. 2017a; Li et al. 2018), we split Area 5 as the test set and use others for training. In the training areas, 4,096 points are sampled for each block and all points in the testing areas are used for testing block-wisely.

In our method, we take an efficient way to implement the weight computation and feature aggregation using matrix multiplication like PointConv (Wu, Qi, and Fuxin 2019). Therefore, we take PointConv as our baseline, but we do not use density information during feature extraction because it have limited improvement in performance.

In the BPM, to automatically annotate the target boundary points for each input point cloud, points with more than $40\%$ of $32$ neighbors not belonging to the same category are assumed to be boundary points. Then, because boundary points are predicted based on neighborhood information, and color information is highly related to boundary prediction, we take the variance of color features of $32$ neighbors as the aggregated feature for each point and further predict the boundary points. After predicting boundary points, we build an encoder-decoder network based on the boundary-aware GEM and take both the color and coordinate information as its input. Our model is trained by Adam optimizer with batch size 8 for ScanNet and batch size 12 for S3DIS on a GTX 1080Ti GPU. Also, we analyze the number of the ground truth of boundary and non-boundary points in different scenes. Accordingly, for ScanNet, $w_{1}$ and $w_{2}$ used in $\mathcal{L}_{BPM}$ are 1 and 10, and for S3DIS, $w_{1}$ and $w_{2}$ are 1 and 2.

For ScanNet v2, we report the mean IoU (mIoU) over categories in Table 1, where we have achieved mIoU of $63.5\%$. It shows our method has outperformed lots of state-of-the-art competitors. Fig. 4 visualizes scene semantic segmentation result of PointConv and our method. Misclassification is easy to appear in the transition area of two adjacent objects. For example, in the second row third column, points of “wall” category are predicted as the “picture” that is adjacent to the wall, leading to the poor contour of the picture. By contrast, benefiting from the boundary awareness, our network perform well in this transition area.

For S3DIS, we report the mIoU over categories in Table 2. We achieve $61.43\%$ in mIoU on this benchmark which has better performance than many state-of-the-art competitors. Also, we visualize our segmentation results in Fig. 4. As can be seen in this figure, the segmentation results obtained by our method have better contour thanks to the Boundary-aware GEM for local feature extraction.

To show the effectiveness of boundary-aware GEM and the GCO, we conduct more ablative experiments. First, we only use MLPs to simulate the 3D convolution kernel like PointConv (Wu, Qi, and Fuxin 2019) and treat it as our baseline. Next, we simply introduce GCO into the baseline to validate the effectiveness of GCO. Then, we use boundary-aware GEM without GCO to build the network and prove the effectiveness of boundary-aware strategy. Finally, we report the result of our method on validation dataset. The results are shown in Table 3, in which we can see both of them can improve the performance on semantic segmentation task.

In our method, we attempt to prevent the propagation of features for points on the boundary for local feature extraction. Meanwhile, BSANet (Zhao et al. 2019) proposed to emphasize the features near boundary. So we try to enhance the influence of points on the boundary during local feature aggregation by:

$$f_{en\_\,l}=\sigma(\mathcal{A}(\{\phi(r_{ij})\cdot\mathcal{M}((2-\hat{g_{j}})\cdot f_{p_{j}})\})),$$

(5)

where $x_{j}\in\mathcal{N}(p_{i})$ and other settings are the same as our proposed method. Recall the Eq. (2), $\hat{g_{j}}$ is $0$ if $p_{j}$ is on the boundary. Given the Eq. (5), more emphasis is imposed on boundary points. The result is shown in Table 3 (“Boundary Augmented”), achieving $61.8\%$ in mIoU, which is better than not using boundary information. However, compared with our proposed method, it decreases the mIoU by $1.6\%$, that means preventing the propagation of features on the boundary is more effective than emphasizing the features near boundary.

Both KCNet and GCO utilize a vector-set kernel to represent localgeometric pattern. Compared with KCNet (Shen et al. 2018), our geometric features are more sensitive to direction rather than position, thus less sensitive to density of points. Additionally, we take a strategy to extract geometric features hierarchically with light-weight geometric convolution to learn complex pattern rather than a heavy-weight module to extract geometric features within one layer.

To show our advantages and give fair comparisons, we first replace GCO with Kernel Correlation (KC) proposed in KCNet with the settings same as KCNet (corresponding to the first row in Table 4). More specifically, the KC is only employed in the first layer with $16$ learnable kernel vectors. Moreover, following the settings of our proposed method, we use a light-weight version of KC, that reduces the learnable kernel vectors from $16$ to $3$, to extract geometric information hierarchically (denoted by KC (H) in Table 4). Also, the extra computational cost is illustrated for these geometry encoding methods in terms of FLOPs. Comparing the row $2$ vs. row $1$ in Table 4, it is shown that using light-weight version of KC to extract geometric features hierarchically obtains better performance than using heavy-weight KC in one layer like KCNet. In addition, using light-weight version of KC decreases the computation drastically. More importantly, keeping other settings the same and using GCO can further increase the mIoU by $2.4\%$. Compared with the light-weight version of KC, GCO has a much simpler form and require less computation resource.

In our implementation, we utilize a kernel unit with a set of only three $3$-D vectors. To check whether a kernel with more or less vectors can learn geometric pattern better, we separately take a kernel with six and two $3$-D vectors and other settings are the same. The results are shown in Table 4, where taking six $3$-D vectors as kernel achieves $61.7\%$ in mIoU which may be caused by overfitting. It also proves our claim that a kernel with three vectors are enough to learn 3D geometry in a hierarchical way and a kernel with less than three vectors is not able to learn 3D geometry, thus has worse performance.

Furthermore, we also conduct ablation study to show the robustness to prediction error introduced by BPM. First, we randomly flip 3% points on prediction results. We report the results on Table 5 and 62.4% mIoU is achieved on the validation set of ScanNet. We think that 3% is enough because the number of randomly flipped points is comparable to the number of target boundary points. Second, we select 5% of the predicted boundary points and exchange the label of each point with its nearest neighbor, making the boundary points shifted by one point and 61.8% mIoU was achieved. Both outperform the network without predicted boundary, showing the robustness to errors in boundary prediction.