PoinTr: Diverse Point Cloud Completion with Geometry-Aware Transformers

Introduction

Recent developments in 3D sensors largely boost researches in 3D computer vision. One of the most commonly used 3D data format is the point cloud, which requires less memory to store but convey detailed 3D shape information. However, point cloud data from existing 3D sensors are not always complete and satisfactory because of inevitable self-occlusion, light reflection, limited sensor resolution, etc. Therefore, recovering complete point clouds from partial and sparse raw data becomes an indispensable task with ever-growing significance.

Over the years, researchers have tried many approaches to tackle this problem in the realm of deep learning. Early attempts on point cloud completion try to migrate mature methods from 2D completion tasks to 3D point clouds by voxelization and 3D convolutions. However, these methods suffer from a heavy computational cost that grows cubically as the spatial resolution increases. With the success of PointNet and PointNet++ , directly processing 3D coordinates becomes the mainstream of point cloud based 3D analysis. The technique is further applied to many pioneer works in point cloud completion task, in which an encoder-decoder based architecture is designed to generate complete point clouds. However, the bottleneck of such methods lies in the max-pooling operation in the encoding phase, where fine-grained information is lost and can hardly be recovered in the decoding phase.

Reconstructing complete point cloud is a challenging problem since the structural information required in the completion task runs counter to the unordered and unstructured nature of point cloud data. Therefore, learning structural features and long-range correlations among local parts of the point cloud becomes the key ingredient towards better point cloud completion. In this paper, we propose to adopt Transformers , one of the most successful architecture in Natural Language Processing (NLP), to learn the structural information of pairwise interactions and global correlations for point cloud completion. Our model, named PoinTr, is characterized by five key components: 1) Encoder-Decoder Architecture: We adopt the encoder-decoder architecture to convert point cloud completion as a set-to-set translation problem. The self-attention mechanism of transformers models all pairwise interactions between elements in the encoder, while the decoder reasons about the missing elements based on the learnable pairwise interactions among features of the input point cloud and queries; 2) Point Proxy: We represent the set of point clouds in a local region as a feature vector called Point Proxy. The input point cloud is convert to a sequence of Point Proxies, which are used as the inputs of our transformer model; 3) Geometry-aware Transformer Block: To facilitate transformers to better leverage the inductive bias about 3D geometric structures of point clouds, we design a geometry-aware block that models the geometric relations explicitly; 4) Query Generator: We use dynamic queries instead of fixed queirs in the decoder, which are generated by a query generation module that summarizes the features produced by the encoder and represents the initial sketch of the missing points; 5) Multi-Scale Point Cloud Generation: We devise a multi-scale point generation module to recover the missing point cloud in a coarse-to-fine manner.

As another contribution, we argue that existing benchmarks are not representative enough to cover real-world scenarios of incompleted point clouds. Therefore, we introduce two more challenging benchmarks that contain more diverse tasks (i.e., joint upsampling and completion of point cloud), more object categories (i.e., from 8 categories to 55 categories), more diverse views points (i.e., from 8 viewpoints to all possible viewpoints) and more diverse level of incompleteness (i.e., missing 25% to 75% points of the ground-truth point clouds). We evaluate our method on both the new benchmarks and the widely used PCN dataset and KITTI benchmark . Experiments demonstrate that PointTr outperforms previous state-of-the-art methods on all benchmarks by a large margin. The main contributions of this paper are summarized in Figure 1.

Related Work

3D Shape Completion. Traditional methods for 3D shape completion tasks often adopt voxel grids or distance fields to describe 3D objects . Based on such structured 3D representations, the powerful 3D convolutions are used and achieve a great success in the tasks of 3D reconstruction and shape completion . However, this group of methods suffers from heavy memory consumption and computational burden. Although these issues are further alleviated by methods based on sparse representations , the quantization operation in these methods still cause a significant loss in detailed information. Different from the above methods, researchers gradually start to use unstructured point clouds as the representation of 3D objects, given the small memory consumption and strong ability to represent fine-grained details. Nevertheless, the migration from structured 3D data understanding to point clouds analysis is non-trivial, since the commonly used convolution operator is no longer suitable for unordered points clouds. PointNet and its variants are the pioneer work to directly process 3D coordinates and inspire the researches in many downstream tasks. In the realm of point cloud completion, PCN is the first learning-based architecture, which proposes an Encoder-Decoder framework and adopts a FoldingNet to map the 2D points onto a 3D surface by mimicking the deformation of a 2D plane. After PCN, many other methods spring up, pursuing point clouds completion in higher resolution with better robustness.

Transformers. Transformers are first introduced as an attention-based framework in Natural Language Processing (NLP). Transformer models often utilize the encoder-decoder architecture and are characterized by both self-attention and cross-attention mechanisms. Transformer models have proven to be very helpful to the tasks that involve long sequences thanks to the self-attention mechanism. The cross-attention mechanism in the decoder exploit the encoder information to learn the attention map of query features, which making transformers powerful in generation tasks. By taking the advantages of both self-attention and cross-attention mechanisms, transformers have a strong capability to handle long sequence input and enhance information communications between the encoder and the decoder. In the past few years, transformers have dominated the tasks that take long sequences as input and gradually replaced RNNs in many domains. Now they begin their journey in computer vision .

Approach

The overall framework of PoinTr is illustrated in Figure 2. We will introduce our method in detail as follows.

The primary goal of our method is to leverage the impressive sequence-to-sequence generation ability of transformer architecture for point cloud completion tasks. We propose to first convert the point cloud to a set of feature vectors, point proxies, that represent the local regions in the point clouds (we will describe in Section 3.2). By analogy to the language translation pipeline, we model point cloud completion as a set-to-set translation task, where the transformers take the point proxies of the partial point clouds as the inputs and produce the point proxies of the missing parts. Specifically, given the set of point proxies $\mathcal{F}=\{F_{1},F_{2},...,F_{N}\}$ that represents the partial point cloud, we model the process of point cloud completion as a set-to-set translation problem:

where $\mathcal{M}_{E}$ and $\mathcal{M}_{D}$ are the encoder and decoder models, $\mathcal{V}=\{V_{1},V_{2},...,V_{N}\}$ are the output features of the encoder, $\mathcal{Q}=\{Q_{1},Q_{2},...,Q_{M}\}$ are the dynamic queries for the decoder, $\mathcal{H}=\{H_{1},H_{2},...,H_{M}\}$ are the predicted point proxies of the missing point cloud, and $M$ is the number of the predicted point proxies. The recent success in NLP tasks like text translation and question answering have clearly demonstrated the effectiveness of transformers to solve this kind of problem. Therefore, we propose to adopt a transformer-based encoder-decoder architecture to solve the point cloud completion problem.

The encoder-decoder architecture consists of $L_{E}$ and $L_{D}$ multi-head self-attention layers in the encoder and decoder, respectively. The self-attention layer in the encoder first updates proxy features with both long-range and short-range information. Then a feed forward network (FFN) further updates the proxy features with an MLP architecture. The decoder utilizes self-attention and cross-attention mechanisms to learn structural knowledge. The self-attention layer enhances the local features with global information, while the cross-attention layer explores the relationship between queries and outputs of the encoder. To predict the point proxies of the missing parts, we propose to use dynamic query embeddings, which makes our decoder more flexible and adjustable for different types of objects and their missing information. More details about the transformer architecture can be found in the supplementary material and .

Note that benefiting from the self-attention mechanism in transformers, the features learned by the transformer network are invariant to the order of point proxies, which is also the basis of using transformers to process point clouds. Considering the strong ability to capture data relationships, we expect the transformer architecture to be a promising alternative for deep learning on point clouds.

2 Point Proxy

The Transformers in NLP take as input a 1D sequence of word embeddings . To make 3D point clouds suitable for transformers, the first step is to convert the point cloud to a sequence of vectors. A trivial solution is directly feeding the sequence of $xyz$ coordinates to the transformers. However, since the computational complexity of the transformers is quadratic to the sequence length, this solution will lead to an unacceptable cost. Therefore, we propose to represent the original point cloud as a set of point proxies. A point proxy represents a local region of the point clouds. Inspired by the set abstraction operation in , we first conduct furthest point sample (FPS) to locate a fixed number $N$ of point centers $\{q_{1},q_{2},...,q_{N}\}$ in the partial point cloud. Then, we use a light-weight DGCNN with hierarchical downsampling to extract the feature of the point centers from the input point cloud. The point proxy $F_{i}$ is a feature vector that captures the local structure around $q_{i}$, which can be computed as:

where $F^{\prime}_{i}$ is the feature of point $q_{i}$ that is extracted using the DGCNN model, and $\varphi$ is another MLP to capture the location information of the point proxy. The first term represents the semantic patterns of the local region, and the second term is inspired by the position embedding operation in transformers, which explicitly encodes the global location of the point proxy. The detailed architecture of the feature extraction model can be found in Supplementary Material.

3 Geometry-aware Transformer Block

One of the key challenges of applying transformers for vision tasks is the self-attention mechanism in transformers lacks some inductive biases inherent to conventional vision models like CNNs and point cloud networks which explicitly model the structures of vision data. To facilitate transformers to better leverage the inductive bias about 3D geometric structures of point clouds, we design a geometry-aware block that models the geometric relations, which can be a plug-and-play module to incorporate with the attention blocks in any transformer architectures. The details of the proposed block are shown in Figure 3. Different from the self-attention module that uses the feature similarity to capture the semantic relation, we propose to use kNN model to capture the geometric relation in the point cloud. Given the query coordinates $p_{Q}$, we query the features of the nearest keys according to the key coordinates $p_{k}$. Then we follow the practice of DGCNN to learn the local geometric structures by feature aggregation with a linear layer followed by the max pooing operation. The geometric feature and semantic feature are then concatenated and mapped to the original dimensions to form the output.

4 Query Generator

The queries $\mathcal{Q}$ serve as the initial state of the predicted proxies. To make sure the queries correctly reflect the sketch of the completed point cloud, we propose a query generator module to generate the query embeddings dynamically conditioned on the encoder outputs. Specifically, we first summarize $\mathcal{V}$ with a linear projection to higher dimensions followed by the max pooing operation. Similar to , we use a linear projection layer to directly generate $M\times 3$ dimension features that can be reshaped as the $M$ coordinates $\{c_{1},c_{2},...,c_{M}\}$. Lastly, we concatenate the global feature of the encoder and the coordinates, and use an MLP to produce the query embeddings.

5 Multi-Scale Point Cloud Generation

The goal of our encoder-decoder network is to predict the missing parts of incomplete point clouds. However, we can only get predictions for missing proxies from the transformer decoder. Therefore, we propose a multi-scale point cloud generation framework to recover missing point clouds at full resolution. To reduce redundant computations, we reuse the $M$ coordinates produced by the query generator as the local centers of the missing point cloud. Then, we utilize a FoldingNet $f$ to recover detailed local shapes centered at the predicted proxies:

where ${\mathcal{P}}_{i}$ is the set of neighboring points centered at $c_{i}$. Following previous work , we only predict the missing parts of the point cloud and concatenate them with the input point cloud to obtain the complete objects. Both predicted proxies and recovered point clouds are supervised during the training process, and the detailed loss function will be introduced in the following section.

6 Optimization

Note that we also concatenate the predicted local centers and the centers of the input point cloud to form the local centers of the whole object $\mathcal{C}$. We directly use the high-resolution point cloud $\mathcal{G}$ to supervise the sparse point cloud $\mathcal{C}$ to encourage them to have similar distributions. Our final objective function is the sum of these two objectives $J=J_{0}+J_{1}$.

Experiments

In this section, we first introduce the new benchmarks for diverse point cloud completion and the evaluation metric. Then, we show the results of both our method and several baseline methods on our new benchmarks. Lastly, we demonstrate the effectiveness of our model on the widely used PCN dataset and KITTI benchmark. We also provide ablation study and visual analysis of our method.

We choose to generate the samples in our benchmarks based on the synthetic dataset, ShapeNet , because it contains the complete object models that cannot be obtained from real-world datasets like ScanNet and S3DIS . What makes our benchmarks distinct is that our benchmarks contain more object categories, more incomplete patterns and more viewpoints. Besides, we pay more attention to the ability of networks to deal with the objects from novel categories that do not appear in the training set.

ShapeNet-55 Benchmark: In this benchmark, we use all the objects in ShapeNet from 55 categories. Most existing datasets for point cloud completion like PCN only consider a relatively small number of categories (e.g., 8 categories in PCN). However, the incompleted point clouds from the real-world scenarios are much more diverse. Therefore, we propose to evaluate the point cloud completion models on all 55 categories in ShapeNet to more comprehensively test the ability of models with a more diverse dataset. We split the original ShapeNet using the 80-20 strategy: we randomly sample 80% objects from each category to form the training set and use the rest for evaluation. As a result, we get 41,952 models for training and 10,518 models for testing. For each object, we randomly sample 8,192 points from the surface to obtain the point cloud.

ShapeNet-34 Benchmark: In this benchmark, we want to explore another important issue in point cloud completion: the performance on novel categories. We believe it is necessary to build a benchmark for this task to better evaluate the generalization performance of models. We first split the origin ShapeNet into two parts: 21 unseen categories and 34 seen categories. In the seen categories, we randomly sample 100 objects from each category to construct a test set of the seen categories (3,400 objects in total) and leave the rest as the training set, resulting in 46,765 object models for training. We also construct another test set consisting of 2,305 objects from 21 novel categories. We evaluate the performance on both the seen and unseen categories to show the generalization ability of models.

Training and Evaluation: In both benchmarks, the partial point clouds for training are generated online. We sample 2048 points from the object as the input and 8192 points as the ground truth. In order to mimic the real-world situation, we first randomly select a viewpoint and then remove the $n$ furthest points from the viewpoint to obtain a training partial point cloud. Although the projection method proposed in is a better approximation to real scans, our strategy is more flexible and efficient. Our experiments on KITTI also show the model learned on our dataset works well when finetuning to real-world scans. Besides, our strategy also ensures the diversity of our training samples in the aspect of viewpoints. During training, $n$ is randomly chosen from 2048 to 6144 (25% to 75% of the complete point cloud), resulting in different level of incompleteness. We then down-sample the remaining point clouds to 2048 points as the input data for models.

During evaluation, we fix 8 view points and $n$ is set to 2048, 4096 or 6144 (25%, 50% or 75% of the whole point cloud) for convenience. According to the value of $n$, we divide the test samples into three difficulty degrees, simple, moderate and hard in our experiments. In the following experiments, we will report the performance for each method in simple, moderate and hard to show the ability of each network to deal with tasks at difficulty levels. In addition, we use the average performance under three difficulty degrees to report the overall performance (Avg).

2 Evaluation Metric

We follow the existing works to use the mean Chamfer Distance as the evaluation metric, which can measure distance between the prediction point cloud and ground-truth in set-level. For each prediction, the Chamfer Distance between the prediction point set $\mathcal{P}$ and the ground-truth point set $\mathcal{G}$ is calculated by:

3 Results on ShapeNet-55

4 Results on ShapeNet-34

On ShapeNet-34, we also conduct experiments for our method and other five state-of-the-art methods. The results are shown in Table 2. For the 34 seen categories, we can see our method outperforms all the other methods. For the 21 unseen categories, we using the networks that are trained on the 34 seen categories to evaluate the performance on the novel objects from the other 21 categories that do not appear in the training phase. We see our method also achieves the best performance in this more challenging setting. Comparing with the results of seen categories, we see in the simple setting (25% of point cloud will be removed), the performance drop of our method is less than 0.3. But when the difficulty level increases, the performance gap between seen categories and unseen categories significantly increases. We also visualize the results in Figure 4 to show the effectiveness of our method on the unseen categories.

5 Results on the Existing Benchmarks

Apart from the experiments on the two newly proposed challenging benchmarks, we also conduct the experiments on the existing benchmarks including the PCN dataset and KITTI benchmark .

Results on the PCN Dataset. The PCN dataset is one of the most widely used benchmark datasets for the point cloud completion task. To verify the effectiveness of our method on existing benchmarks and compare it with more state-of-the-art methods, we conducted experiments on this dataset following the standard protocol and evaluation metric used in previous work . The results are shown in Table 3. We see our method largely improves the previous methods and establishes the new state-of-the-art on this dataset.

Results on KITTI Benchmark. To show the performance of our method in real-world scenarios, we follow to finetune our trained model on ShapeNetCars and evaluate the performance of our model on KITTI dataset, which contains the incomplete point clouds of cars in the real-world scenes from LiDAR scans. We report the Fidelity and MMD metrics in Table 5 and show some reconstruction results in Figure 5. Our method achieves better qualitative and quantitative performance.

6 Model Design Analysis

To examine the effectiveness of our designs, we conduct a detailed ablation study on the key components of PoinTr. The results are summarized in Table 4. The baseline model A is the vanilla transformer model for point cloud completion, which uses the encoder-decoder architecture with the standard transformer blocks. In this model, we form the point proxies directly from the point cloud using a single-layer DGCNN model. We then add the query generator between the encoder and decoder (model B). We see the query generator improve the baseline by 0.34 in Chamfer distance. When using DGCNN to extract features from the input point cloud (model C), we observe a significant improvement to 8.69. By adding the geometric block to all the transformer blocks (model D), we see the performance can be further improved, which clearly demonstrates the effectiveness of the geometric structures learned by the block. We find that only adding the geometric block to the first transformer block in both encoder and decoder can lead to a slightly better performance (model E), which indicates the role of geometric block is to introduce the inductive bias and a single layer is sufficient while adding more blocks may result in over-fitting. Besides, we see our method can achieve over 0.74 F-Score on the PCN dataset while obtaining only 0.46 F-Score on our ShapeNet-55, which also suggests our new datatset is much more challenging.

7 Qualitative Results

In Figure 6, we show some completion results for all methods and find our method perform better. For example, the input data in (a) nearly lose all the geometric information and can be hardly recognized as an airplane. In this case, other methods can only roughly complete the shape with unsatisfactory geometry details, while our method can still complete the point cloud with higher fidelity. These results show our method has a stronger ability to recover details and is more robust to various incomplete patterns. More results can be found in the supplementary material.

Conclusion

In this paper, we have proposed a new architecture, PoinTr, to convert the point cloud completion task into a set to set translation tasks. With several technical innovations, we successfully applied the transformer model to this task and achieved state-of-the-art performance. Moreover, we proposed two more challenging benchmarks for more diverse point cloud completion. Extending our transformer architecture to other 3D tasks can an interesting future direction.

This work was supported in part by the National Key Research and Development Program of China under Grant 2017YFA0700802, in part by the National Natural Science Foundation of China under Grant 61822603, Grant U1813218, and Grant U1713214, in part by a grant from the Beijing Academy of Artificial Intelligence (BAAI), and in part by a grant from the Institute for Guo Qiang, Tsinghua University.

Appendix A Implementation Details

Our proposed method PointTr is implemented with PyTorch . We utilize AdamW optimizer to train the network with initial learning rate as 0.0005 and weight decay as 0.0005. In all of our experiments, we set the depth of the encoder and decoder in our transformer to 6 and 8 and set $k$ of kNN operation to 16 and 8 for the DGCNN feature extractor and the geometry-aware block respectively. We use 6 head attention for all transformer blocks and set their hidden dimensions to 384. On the PCN dataset, the network takes 2048 points as inputs and is required to complete the other 14336 points. We set the batch size to 54 and train the model for 300 epochs with the continuous learning rate decay of 0.9 for every 20 epochs. We set $N$ to 128 and $M$ to 224. On ShapeNet-55/34, the model takes 2048 points as inputs and is required to complete the other 6144 points. We set the batch size to 128 and train the model for 200 epochs with the continuous learning rate decay of 0.76 for every 20 epochs. We set $N$ to 128 and $M$ to 96.

We employ a lightweight DGCNN model to extract the point proxy features. To reduce the computational cost, we hierarchically downsample the original input point cloud to $N=128$ center points and use several DGCNN layers to capture local geometric relationships. The detailed network architecture is: $\texttt{Linear}(C_{in}=3,C_{out}=8)$ $\rightarrow$ $\texttt{DGCNN}(C_{in}=8,C_{out}=32,K=8,N_{out}=2048)$ $\rightarrow$ $\texttt{DGCNN}(C_{in}=32,C_{out}=64,K=8,N_{out}=512)$ $\rightarrow$ $\texttt{DGCNN}(C_{in}=64,C_{out}=64,K=8,N_{out}=512)$ $\rightarrow$ $\texttt{DGCNN}(C_{in}=64,C_{out}=128,K=8,N_{out}=128)$, where $C_{in}$ and $C_{out}$ are the numbers of channels of input and output features, $N_{out}$ is the number of points after FPS.

Appendix B Technical Details on Transformers

Encoder-Decoder Architecture. The overall architecture of the transformer encoder-decoder networks is illustrated in Figure 7. The point proxies are passed through the transformer encoder with $N$ multi-head self-attention layers and feed-forward network layers. Then, the decoder receives the generated query embeddings and encoder memory, and produces the final set of predicted point proxies that represents the missing part of the point cloud through $N$ multi-head self-attention layers, decoder-encoder attention layers and feed-forward network layers. We set $N$ to 6 in all our experiments following common practice .

Multi-head Attention. Multi-head attention mechanism allows the network to jointly attend to information from different representation subspaces at different positions . Speciacally, given the input values $V$, keys $K$ and queries $Q$, the multi-head attention is computed by:

where $W^{O}$ the weights of the output linear layer and each head feature can be obtained by:

where $W_{i}^{Q}$, $W_{i}^{K}$ and $W_{i}^{V}$ are the linear layers that project the inputs to different subspaces and $d_{k}$ is the dimension of the input features.

Feed-forward network (FFN). Following , we use two linear layers with ReLU activations and dropout as the feed-forward network.

Appendix C Detailed Experimental Results

In Table 6, we report the detailed results for FoldingNet , PCN , TopNet , PFNet , GRNet and the proposed method on ShapeNet-55. Each row in the table stands for a category of object. We test each method under three settings: simple, moderate and hard.

Detailed results on ShapeNet-34:

In Table 7, we report the detailed results for the novel objects from 21 categories in ShapeNet-34. Each row in the table stands for a category of object. We test each method under the three settings: simple, moderate and hard.

Appendix D Complexity Analysis

Our method achieves the best performance on both our newly proposed diverse benchmarks and the existing benchmarks. We provide the detailed complexity analysis of our method in Table 8. We report the number of parameters and theoretical computation cost (FLOPs) of our method and other five methods. We also provide the average Chamfer distances of all categories in ShapeNet-55 and unseen categories in ShapeNet34 as references. We can see our method achieves the best performance while using relatively low parameters and FLOPs among the methods in the table, which shows our method offers a decent trade-off between cost and performance.

Appendix E Visualization of the Predicted Centers

We visualize the local center prediction results on ShapeNet-55. We adopt a coarse-to-fine strategy to recover the point cloud. Our method starts with the prediction of local centers, then we can obtain the final results by adding the points around the centers. As shown in Figure 8, Line (a) shows the input partial point cloud and the predicted point centers. Line (b) is the predicted point clouds. We see the predicted point proxies can successfully represent the overall structure of the point cloud and the details then are added in the final predictions.

Appendix F Qualitative Results

In Figure 9, we provide more qualitative results on ShapeNet-55. We see our results are much better than baseline methods visually.

References