Spatial Transformer Point Convolution

The standard convolution on images/videos is conditioned to grid shape region, and performs anisotropic filtering in order to encode subtle texture appearances. The point clouds data are often embedded in an unstructured and non-regular 3D space, and thus the standard convolution kernels are difficult to be defined on point clouds data due to its irregular property. Most existing convolution methods on point clouds  [22, 24] take the sum-/max-aggregation on local neighbor points without considering the local space topology of points therein. As a contrast, we attempt to construct an anisotropic filtering on 3D point clouds to better encode these finer characteristics of points in each local 3D spatial region. To this end, we introduce spatial direction dictionaries to coordinate these nearest points into latent normal spaces, where each atom of the dictionary implicitly defines a spatial direction. The previous filter operated on all nearest points can be untied to adapt to different atoms, which naturally results to anisotropic convolution filtering analogous to the standard convolution. This anisotropic filtering can be named as spatial transformer point convolution (STPC).

Given one reference point $p_{i}$, we can sample the $K$ nearest neighbor points as a set $\mathcal{N}(p_{i})=\{p_{k}|k=1,\cdots,K\}$. Obviously, each point has 3D coordinate value, denoted as $\mathbf{p}_{k}=(x_{k},y_{k},z_{k})^{\intercal}$¹¹1Here the lowercase $p$ denotes one general point, while the bold-type $\mathbf{p}$ denotes its spatial x-y-z coordinates. . Point clouds often carry some additional attributes, e.g., RGB color or intermediate learnt features. We denote the attributes of point $p_{i}$ as a vector $\mathbf{x}_{i}$. Next, we perform an anisotropic filtering on the point set $\mathcal{N}(p_{i})$ to produce a robust response of the reference point $p_{i}$. Below we define one-time spatial transformer point convolution process, which can be encapsulated into an STPC module as one layer. Formally,

	$\displaystyle\mathbf{x}_{i}^{l+1}$	$\displaystyle\longleftarrow\psi_{\text{conv}}([\widetilde{\mathbf{x}}_{1}^{l},\widetilde{\mathbf{x}}_{2}^{l},\cdots,\widetilde{\mathbf{x}}_{M}^{l}]),$		(1)
	$\displaystyle\widetilde{\mathbf{x}}_{m}^{l}$	$\displaystyle\longleftarrow\mathcal{T}_{\text{f}}(\{(\alpha_{km}^{l},\mathcal{F}_{\text{f}}^{l}(\mathbf{p}_{k},\mathbf{x}_{k}^{l}))|p_{k}\in\mathcal{N}(p_{i})\}),\quad m=1,\cdots M,$		(2)
	$\displaystyle\bm{\alpha}_{k}^{l}:$	$\displaystyle[\alpha^{l}_{k1},\alpha^{l}_{k2},\cdots,\alpha^{l}_{kM}]\longleftarrow\mathcal{T}_{\text{s}}(\mathbf{A}^{l},\mathcal{F}_{\text{s}}(\mathbf{p}_{k})),\quad p_{k}\in\mathcal{N}(p_{i}),$		(3)
	$\displaystyle\mathbf{A}^{l}$	$\displaystyle:[\mathbf{a}_{1}^{l},\mathbf{a}_{2}^{l},\cdots,\mathbf{a}_{M}^{l}]\longleftarrow\mathcal{F}_{\text{a}}(\mathbf{a}_{1}^{l-1},\mathbf{a}_{2}^{l-1},\cdots,\mathbf{a}_{M}^{l-1}),$		(4)

where $l$ is the layer number of neural network, $M$ is the number of atoms $\{\mathbf{a}_{1},\mathbf{a}_{2},\cdots,\mathbf{a}_{M}\}$ in the spatial direction dictionary $\mathbf{A}$, $\bm{\alpha}_{k}$ is the encoding coefficient w.r.t the point $p_{k}$, $\widetilde{\mathbf{x}}_{m}$ is the accumulated component at the $m$-th atom for all neighbor points of $p_{i}$, $\mathbf{x}_{i}^{l+1}$ is the anisotropic convolution response of point $p_{i}$. In the above formulas,

• •

  Direction dictionary leaning $\mathcal{F}_{\text{a}}$. Due to spatial encoding, this dictionary should emphasize spatial information (source from points positions), rather than attribute information. Meantime, the dictionary should be dynamic and depended on prior states by accompanying with stacked modules.

• •

  Spatial direction encoding $\mathcal{T}_{\text{s}}$. The point $p_{k}$ is represented with the atoms of the direction dictionary by using sparse coding or nearest neighbor reconstruction. The representation coefficients describe the latent coordinate information of this point in the dictionary space. The spatial transformation $\mathcal{F}_{\text{s}}$ is used to extract spatial features from 3D coordinates.

• •

  Spatial transformer $\mathcal{T}_{\text{f}}$. The neighbor points set $\mathcal{N}(p_{i})$ is projected onto atom points, and forms a regular coordinate system defined by $M$ atoms. The feature transformation $\mathcal{F}_{\text{f}}$ is used to learn point features/attributes.

• •

  Anisotropic convolution (i.e., STPC) $\psi_{\text{conv}}$. As the dictionary implicitly defines a regular coordinate system, the spatial transformed features $\{\widetilde{\mathbf{x}}_{m}|m=1,\cdots,M\}$ just corresponds to the atoms therein. Hence, we could perform different filtering for each of them, which really works like the standard convolution on different relative spatial positions.

The above convolution process can be framed into a STPC network module as shown in Fig. 1, whose implementation details are introduced in section 3.2,  3.3 and  3.4.

In order to produce consistent responses under different permutation layouts, most existing methods mix these geometric information together to explore the relationship between a point and its neighbor points. Intuitively, these geometric/direction information of point clouds lie in a dense space with several latent components/directions. Inspired by the idea of bag-of-words, here we learn a spatial direction dictionary in the $c$-dimensional space, which can be then used to sparsely quantize local neighborhood structure of point clouds. Specifically, we build a spatial direction dictionary $\mathbf{A}=[\mathbf{a}_{1},\mathbf{a}_{2},\dots,\mathbf{a}_{M}]$ with $M$ atoms through random initialization, where each atom $\mathbf{a}_{m}\in\mathbb{R}^{c}$ refer to a spatial direction information. Here the initialized dictionary can be noted as $\mathbf{A}^{0}=[\mathbf{a}_{1}^{0},\mathbf{a}_{2}^{0},\dots,\mathbf{a}_{M}^{0}]$. Due to the sparsity of point clouds, a certain number of atoms are sufficient to encode each neighborhood region, and the redundancy and expensive computing costs can be also avoided in practice. With the prior states of direction dictionary in previous layer, we can dynamically optimize/update it while satisfying the constrain of atom decorrelation.

$\displaystyle\min\sum_{p=1}^{M}\sum_{q\neq{p}}^{M}\mathcal{S}_{(\mathbf{a}_{p}^{l},\mathbf{a}_{q}^{l})},\qquad\text{s.t.}\ \ \mathbf{A}^{l}=\mathcal{F}_{\text{a}}(\mathbf{A}^{l-1}),$

(5)

where $\mathbf{A}^{l}:[\mathbf{a}_{1}^{l},\mathbf{a}_{2}^{l},\cdots,\mathbf{a}_{M}^{l}]$ and $\mathbf{A}^{l-1}:[\mathbf{a}_{1}^{l-1},\mathbf{a}_{2}^{l-1},\cdots,\mathbf{a}_{M}^{l-1}]$ denote the states of direction dictionary in the $l$-th and $l-1$-th layers, the dictionary learning function $\mathcal{F}_{\text{a}}$ is implemented by a shared multi-layer perceptron. In the direction dictionary learning process, we use the cosine similarity to measure the similarity between each of them.

$$\mathcal{S}_{(\mathbf{a}_{p}^{l},\mathbf{a}_{q}^{l})}=\frac{\mathbf{a}_{p}^{l}\cdot\mathbf{a}_{q}^{l}}{\|\mathbf{a}_{p}^{l}\|\|\mathbf{a}_{q}^{l}\|},\qquad p,q=1,2,\dots,M,$$

(6)

where the greater the cosine similarity $\mathcal{S}_{(\mathbf{a}_{p}^{l},\mathbf{a}_{q}^{l})}$ is, the more similar the spatial directions between $\mathbf{a}_{p}^{l}$ and $\mathbf{a}_{q}^{l}$ are. The constraint is that the atoms of the spatial direction dictionary are discrete and not similar between them. We can iteratively learn the spatial direction dictionary by minimizing the above constraint, and then the atoms can discretely represent different spatial directions in the high-dimensional dictionary space.

For a single point, its x-y-z coordinates can not fully explore the local geometric information with its neighbor points, which is important for the semantic segmentation task. Several recent works[45, 19, 8] have explicitly embedded the low-level elements of neighbor points to encode local features. These above methods focus on expressing local structure relationships by employing multiple low-level elements, but the distinguishable spatial direction information of each point would be drowned out. Therefore, we introduce a spatial direction encoding method to capture and represent implicit geometric structures and distinguishable direction information. For the $i$-th point, its spatial direction information can be represented as $[\mathbf{d}_{1},\mathbf{d}_{2},\dots,\mathbf{d}_{K}]\in\mathbb{R}^{K\times c}$, and each of them can be encoded as:

$\displaystyle\mathbf{d}_{k}=\mathcal{F}_{\text{s}}(\mathbf{p}_{i}-\mathbf{p}_{k}),\ \ \ \ p_{k}\in\mathcal{N}(p_{i}),$

(7)

where the point set $\{p_{k}|k=1,2,\dots,K\}$ refers to the $K$ nearest neighbor points of the point $p_{i}$, $c$ is the feature dimension of $\mathbf{d}_{k}$, and the spatial transformation $\mathcal{F}_{\text{s}}$ is performed by the shared multi-layer perceptron for extracting the spatial direction information.

In the $l$-th convolution layer, we can encode the local spatial information of point $p_{i}$ by employing the direction dictionary $\mathbf{A}^{l}$ and spatial direction information $[\mathbf{d}_{k}|k=1,2,\dots,K]$. For one nearest point $p_{k}$ of the point $p_{i}$, we can encode it with the atoms of the direction dictionary by using sparse coding. Analogous to the self-attention mechanism [34] in the NLP field, $\mathbf{d}_{k}$ and $\mathbf{a}_{m}^{l}$ refer to a query and one dictionary atom, respectively. For describing the latent coordinate information of the point $p_{k}$, we can get the representation coefficient by computing the correlation between spatial direction information $\mathbf{d}_{k}$ and the atom $\mathbf{a}_{m}^{l}$. Specifically, we use the cosine similarity as Eqn. (6) to measure the correlation between them and an exponential function is used to make the representation coefficients more sparse. The spatial encoding process $\mathcal{T}_{\text{s}}$ can be expressed as computing the direction-dependent relationship between $p_{k}$ and each atom $\mathbf{a}_{m}^{l}\in\mathbf{A}^{l}$:

$$\alpha_{km}^{l}=\frac{\exp\left(\mathcal{S}_{(\mathbf{d}_{k},\mathbf{a}_{m}^{l})}\right)}{\sum_{m=1}^{M}\exp\left(\mathcal{S}_{(\mathbf{d}_{k},\mathbf{a}_{m}^{l})}\right)},$$

(8)

where $\alpha_{km}^{l}$ may be viewed as the attention score/coefficient of the point $p_{k}$ w.r.t the atom $\mathbf{a}_{m}$. Thus the spatial information of the point $p_{k}$ in the dictionary space is encoded as $\bm{\alpha}_{k}^{l}=[\alpha^{l}_{k1},\alpha^{l}_{k2},\cdots,\alpha^{l}_{kM}]$. Besides, the study on sparse coding [44] demonstrates that the sparsity of representation coefficients is important to overcomplete dictionaries. Hence, we discard those trivial coefficients in $\bm{\alpha}$ by setting a small threshold value (i.e., $\tau$ = 0.01), which results into sparse representation in practice.

To encode the unordered local information of the point $p_{i}$, we introduce a spatial transformer $\mathcal{T}_{\text{f}}$ for projecting the neighbor points set $\mathcal{N}(p_{i})$ into a canonical order dictionary space. After obtaining the attention score, the feature of all neighbor points $p_{k}\in\mathcal{N}(p_{i})$ can be transformed into the canonical ordered dictionary space.

	$\displaystyle\widetilde{\mathbf{x}}_{m}^{l}$	$\displaystyle=\mathcal{T}_{\text{f}}(\{(\alpha_{jm}^{l},\mathcal{F}_{\text{f}}^{l}(\mathbf{p}_{k},\mathbf{x}_{k}^{l}))|p_{k}\in\mathcal{N}(p_{i})\})$		(9)
		$\displaystyle=\sum_{k=1}^{K}\alpha_{km}^{l}\mathcal{F}_{\text{f}}^{l}(\mathbf{p}_{k},\mathbf{x}_{k}^{l}),\quad m=1,\cdots M,$

where the feature transformation $\mathcal{F}_{\text{f}}^{l}(\mathbf{p}_{k},\mathbf{x}_{k}^{l})$ refers to the feature extracting process by employing the shared multi-layer perceptron, the x-y-z coordinates $\mathbf{p}_{k}$ and the geometry-aware features $\mathbf{x}_{k}^{l}$ of the point $p_{k}$ are used as inputs. Thus, any disordered neighborhood can be transformed into the canonical ordered dictionary space, and then we can perform an anisotropic convolution filtering $\psi_{\text{conv}}$ to produce a robust response of the point $p_{i}$.

	$\displaystyle\mathbf{x}_{i}^{l+1}$	$\displaystyle=\psi_{\text{conv}}([\widetilde{\mathbf{x}}_{1}^{l},\widetilde{\mathbf{x}}_{2}^{l},\cdots,\widetilde{\mathbf{x}}_{M}^{l}])$		(10)
		$\displaystyle=\mathbf{w}_{1}*\widetilde{\mathbf{x}}_{1}^{l}+\mathbf{w}_{2}*\widetilde{\mathbf{x}}_{2}^{l}+\dots+\mathbf{w}_{M}*\widetilde{\mathbf{x}}_{M}^{l},$

where $[\widetilde{\mathbf{x}}_{1}^{l},\widetilde{\mathbf{x}}_{2}^{l},\cdots,\widetilde{\mathbf{x}}_{M}^{l}]$ denotes the spatial transformed features with the ordered structure in the dictionary space, $[\mathbf{w}_{1},\mathbf{w}_{2},\dots,\mathbf{w}_{M}]$ refers to the learnt weights for performing anisotropic convolution.

To demonstrate the efficiency of our proposed STPC method, we conduct semantic segmentation experiments on three large-scale point clouds datasets, including S3DIS [1], Semantic3D [6] and SemanticKITTI [3]. Given an point clouds as input, the widely-used encoder-decoder architecture is used to predict the point-wise label in an end-to-end way. In the encoding stage, we stack five STPC modules to extract per-point features, each of which is followed by a random sampling operation [8] to reduce the size of the point clouds. In the decoding stage, five upsampling layers are stacked to propagate features by employing the nearest point interpolation technique. Meanwhile, skip connection is applied to concatenate the upsampled features with the intermediate features produced by encoding layers. To predict the categories of input point clouds, we stack several fully-connected layers on this responses, and use the cross-entropy loss during training. Follow the same protocols as in [8], we split piont cloud data into training set and testing set, and adopt the same original features as inputs (3D coordinates and color information for S3DIS and Semantic3D, and only 3D coordinates for SemanticKITTI). We adopt the mean IoU (mIoU), mean class Accuracy (mAcc) and Overall Accuracy (OA) over the total classes as the standard metrics. The learning rate is initialized at 0.01 and decreases by 5% after each epoch. The number of nearest points K is set as 16. As in RandLA-Net [8], we use the Adam optimizer with default parameters and the same data processing methods. Our learnt spatial direction dictionary contains 25 atoms (i.e., $M=25$), the dimension of each atom is 16 (i.e., $c=16$), each of which represents a different spatial direction. For smooth updates, the learning rate of direction dictionary learning is set to 0.01 times the global network. All models in our experiment are trained and tested based on an NVIDIA RTX2080Ti GPU.

S3DIS dataset[1]: Table 1 shows the performance of our STPC method and comparisons with several state-of-the-arts on OA, mAcc and mIoU metrics. our STPC method can significantly outperform four baselines: 2.1% over SPGraph [14], 2.6% over PointCNN [17], 3.1% over PointConv [41] and 1.5% over PointWeb [47] in term of OA score. When compared with these existing point clouds semantic segmentation approaches, our method in term of mAcc score also achieves the best result, representing improvement of 26.3%, 13.1%, 11.4%, 4.5%, and 2.5% over PointNet [22], TangentConv [30], PointCNN [17], PAT [45], KPConv deform [32]. The proposed STPC method can also obtain a better performance than these exiting methods on mIoU score, except for KPConv deform method [32]. Furthermore, we also report the mIoU scores for each class in Table 1. Generally, when compared with these baselines, our STPC method shows the comparable performance, e.g, 76.3% vs 75.4% [32] for sofa, 89.2% vs 91.0 [32] for chair, 66.5% vs 64.9 [47] for board, and 82.1% vs 79.4% [17] for wall. It demonstrates that our STPC method performs very well on predicting the point-wise labels by considering the anisotropic convolution filtering technology on 3D point clouds data.

Semantic3D dataset [6]: We report the results of comparisons between the baselines and our STPC method on the Semantic3D dataset (reduced-8). As shown in Table 2, comparisons with previous methods [31, 46, 32, 33] demonstrate that our proposed STPC obtains the best performance, achieving improvements of 0.3% over RGNet [33], 1.9% over KPConv rigid [32], 1.6% over ShellNet [46] and 6.7% over SegCloud [31] in term of OA score. The STPC achieves 76.2% in term of OA score, which can outperform most existing methods [31, 46, 32, 33] , except the RandLA-Net method [8]. For example, when compared with the classic SPGraph [14], we can boost the semantic segmentation performance by 0.8 and 3.0% in terms of OA and mIoU scores, respectively. Furthermore, when predicting the labels for each category, the mIoU scores for man-made and natural can achieve the best performance with 97.9% and 93.0% segmentation results. This indicates that our SPTC can better capture these finer local features of point clouds and then improve the discriminative capability of semantic segmentation network.

SemanticKITTI dataset [3]: We finally compare the proposed STPC method with two families of recent point clouds semantic segmentation approaches, including point based methods [22, 14, 24, 30, 8] and projection based approaches [39, 40, 3, 21]. The quantitative comparison of segmentation results is reported in Table 3. When compared with all point based methods and projection based approaches, the STPC can achieve the best performance in term of mIoU score, and also obtain the better results in most categories (11 out of 19). Especially in the motorcyclist category, the mIoU results of all these existing methods are quite low (e.g., 0.0% with PointNet [22], SPLATNet [28] and PointNet++ [24], 0.9% with SqueezeSeg [39], 7.2% with RandLA-Net [8]), but we can achieve 15.3% segmentation result, which is much higher than the second-ranked TangentConv method [30] by improving 7.2%. Compared with the best projection based approach (i.e., RangeNet53++ [21]), we outperform it by 2.4% in term of mIoU score, and obtain the better results in 14 categories. Compared with the best point based method (i.e., RandLA-Net [8]), we outperform it by 0.7% in term of mIoU score, and achieve the better results in 16 categories. Some visualization results between our STPC method and the compared RandLA-Net [8] can be found in Fig. 4. These above experimental results indicate that our proposed STPC method can better predict the point-wise category information of point clouds data by improving the predictive capability of the network.

In this section, we conduct the following ablation studies for our spatial transformer point convolution module. All ablated networks are trained on area 1$\sim$4 and 6, and tested on area 5 of S3DIS dataset.

As illustrated in Figure.3, we compare the semantic segmentation performance with different number of atoms (i.e., $M=1,4,9,16,25,36$) in the spatial direction dictionary. When the number of atoms increases from 1 to 25, the performance of point clouds semantic segmentation is also improving from 64.87% to 66.56% in term of mIoU score. If the number of atoms continues to increase, the segmentation performance would be slightly decline. Our STPC method achieves the best result when the number of dictionary atoms is 25 ($M=25$). If the number of atoms is too small, the local geometric information cannot be captured finely in the anisotropic convolution filtering process. When the number of atoms equals to 1, our method would degenerate to the sum operation, while if the number of atoms is too large, the learnt weights tend to average on each atom, and these finer features would be lost. Therefore, a suitable number of dictionary atoms can make the proposed SPTC method to capture the feature varieties of point clouds.

For better understanding the direction dictionary learning process, we visualize the learned representation coefficients in Eqn. 8, which can reflect the direction-dependent relationship between one neighbor point $p_{k}$ and each atom $\mathbf{a}_{m}\in\mathbf{A}$. As shown in Fig.3, the vertical ordinates refer to the index of the K nearest neighbor points (i.e., $p_{k}\in\mathcal{N}(p_{i})$, $k=1,2,\dots,16$) and the horizontal ordinates denote the index of dictionary atoms (i.e., $\mathbf{a}_{m}\in\mathbf{A}$, $m=1,2,\dots,25$). The local neighbor points around the point $p_{i}$ can be projected onto different dictionary atoms. It indicates that the feature-dense local point clouds can be sparsely encoded by sparse discrete direction dictionary.

For exploring the effectiveness of the anisotropic filtering/convolution in our network, we evaluate the segmentation performance with several network variants on the S3DIS dataset. Specifically, we directly perform convolution on unordered point clouds (“Conv. on unordered points"), and conduct isotropic convolution by aggregating local spatial points after max-, mean-, and sum-pooling operations, named “Isotropic Conv. after max-pooling", “Isotropic Conv. after mean-pooling", and “Isotropic Conv. after sum-pooling", respectively. As can be seen in Table 4, when compared with “Conv. on unordered points", our proposed anisotropic filtering can improve the segmentation perform from 64.15% to 66.56%. When compared with three different isotropic convolutions, the performance of our proposed anisotropic convolution is better than three isotropic convolution methods, representing improvements of 1.25%, 2.39% and 1.69% over “Isotropic Conv. after max-pooling", “Isotropic Conv. after mean-pooling", and “Isotropic Conv. after sum-pooling". It proves that our anisotropic filtering method (i.e., STPC) can capture and represent implicit geometric structures, which is more robust to express those finer variances of local regions.