Context-Aware Transformer for 3D Point Cloud Automatic Annotation

Unlike 2D images, where pixels are arranged on regular grids and can be processed by CNNs, 3D point clouds are inherently sparse, unstructured, and scattered in 3D space: they are essentially set. These properties make typically convolutional architectures hardly applicable. Two ad hoc solutions to consume point clouds are transforming irregular points into voxels or grids and directly ingesting point clouds. Voxel-based networks (Maturana and Scherer 2015; Wu et al. 2015; Xie et al. 2018; Yang et al. 2018) divide point clouds into volumetric grids and employ 2D/3D CNNs for predictions. Point-based methods (Qi et al. 2017a, b; Lang et al. 2019; Shi, Wang, and Li 2019; Qi et al. 2018) directly consume raw point clouds to preserve more original geometric information for 3D tasks.

As an alternative, attention-based models have been recently used to represent 3D point clouds for 3D vision tasks, e.g., retrieval (Zhang and Xiao 2019), indoor and outdoor segmentation (Zhao et al. 2021a), object classification (Liu et al. 2020), and 3D object detection (Misra, Girdhar, and Joulin 2021; Pan et al. 2021), since Transformer and attention models have revolutionized natural language processing (Devlin et al. 2018; Wolf et al. 2020) and computer vision (Vaswani et al. 2017; Dosovitskiy et al. 2020; Liu et al. 2021; Carion et al. 2020). In Point Transformer (Zhao et al. 2021a), they design a point Transformer layer with vector self-attention layers and use it for constructing an Unet-like architecture for 3D point cloud segmentation and classification. In PCT (Guo et al. 2021), they make several adjustments on Transformers for 3D and introduce offset-attention and neighbor embedding to make a point Transformer framework suitable for point cloud learning. Adopting a Transformer encoder-decoder architecture, (Yu et al. 2021) proposes PoinTr to reformulate point cloud completion as a set-to-set translation problem. By taking advantage of the solid capability to process long-range token-wise interactions and enhance information communications, Transformers begin their journey in 3D tasks.

However, existing Transformer-based models for 3D point clouds bring in handcrafted inductive biases (e.g., local feature aggregation, neighbor grouping, and embeddings) and usually overlook the inter-object relation for effective feature learning on 3D tasks. Although the above methods apply specific 3D properties to modify the Transformer, we push the limits of the standard Transformer for 3D intra-object and inter-object feature extractions.

Human point cloud annotation is laborious (Meng et al. 2021; Wei et al. 2021; Liu et al. 2022a, b), hindering the training of high-quality 3D object detectors on massive datasets (Geiger, Lenz, and Urtasun 2012; Caesar et al. 2020; Sun et al. 2020). 2D weak annotations are more easily-acquired and cheaper compared with 3D labels. 3D object annotation from weak labels raises ever-growing attention. To our knowledge, very few works attempt to research 3D automatic annotation models. Using 2D boxes, FGR (Wei et al. 2021) proposes a heuristic algorithm to first locate key points and edges in the sub-point-cloud frustum after a coarse 3D segmentation and then estimate the 3D box. SDF (Zakharov et al. 2020) adopts the predefined models to process 3D segmented point clouds to estimate 3D annotations from 2D bounding boxes. WS3D (Meng et al. 2021, 2020) explores center clicks as weak annotations to generate cylindrical object proposals on bird’s view (BEV) maps and refine them to get 3D labels at the second stage. MAP-Gen (Liu et al. 2022b) designs a 3-stage automated 3D-box annotation workflow with 2D weak boxes: first foreground segmentation, point generation, and final 3D box generations. MTrans (Liu et al. 2022a) leverages both LiDAR scans and camera images to generate precise 3D labels from weak 2D bounding boxes with a multi-task design of segmentation, point generation, and box regression.

While they have demonstrated state-of-the-art performance on the KITTI benchmark, such complicated annotation designs involve pipelined training models, intermediately processing points, designing special loss functions, and multi-modality information combinations. Compared to these methods, our model is an end-to-end Transformer-based 3D annotator ( no convolutional backbone) that can be trained from scratch. CAT has removed many cumbersome designs, i.e., 3D segmentation, point completion, and multi-modality information combination. CAT requires minimal modifications to the vanilla Transformer layer.

Our goal is to generate 3D pseudo labels that can be used to train any off-the-shelf 3D object detector (e.g., PointRCNN). Given LiDAR point clouds and weak 2D boxes, we first extract the frustum area tightly fitted with 2D bounding boxes, following (Qi et al. 2018; Wei et al. 2021; Liu et al. 2022a, b). With a known LiDAR-to-Image calibration matrix, a 3D point cloud in the form of $(x,y,z)$ can be projected onto its 2D image plane $(h,w)$ by the mapping function $f_{cal}$. Therefore, the 2D projected sub-cloud $\mathcal{P}_{2D}\in\mathbb{R}^{N\times 2}$ within a 2D box can be defined as:

$$\mathcal{P}_{2D}=\{(h,w)\;|\;(h,w)=f_{cal}(x,y,z),(h,w)\in\mathcal{B}\}.$$

(1)

Therefore, the frustum sub-cloud $\mathcal{P}_{F}\in\mathbb{R}^{N\times 3}$ within its 2D bounding box can be extracted as:

$$\mathcal{P}_{F}=\{(x,y,z)\;|\;f_{cal}(x,y,z)\in\mathcal{P}_{2D}\},$$

(2)

where $(x,y,z)$ is the point coordinate, and $\mathcal{B}$ is the 2D region in the 2D box. We use $N$ to represent input frustum sub-cloud size (i.e., the number of points).

CAT takes as input $B$ frustum sub-clouds $\mathcal{P}_{F}\in\mathbb{R}^{N\times 3}$ within 2D boxes, and generates 3D labels in the form of 3D bounding boxes ${Box}_{3D}\in\mathbb{R}^{B\times 7}$, where batch size is denoted as $B$. An input point sub-cloud is an unordered set of points associated with 3-dimensional coordinates $(x,y,z)$. The number of points varies from individual to individual. Some are very large when LiDAR scans are complete, but some are small due to occlusion, truncation, limited sensor resolution, light reflection, etc. We use the random sampling from (Qi et al. 2018; Wei et al. 2021; Liu et al. 2022a, b) to keep the same size (i.e., $N=1024$ points) of each sample and project them to $d=512$ dimensional point embeddings through a multilayer perceptron (MLP). The resulting point embeddings are passed through the CAT encoder to obtain a set of features. The additional box tokens are taken as box representations along with point embeddings in the encoder, inspired by the class token in ViT (Dosovitskiy et al. 2020). A decoder takes as input box tokens as queries, points features as keys, and values from encoder output and predicts 3D bounding boxes.

Both encoder and decoder employ standard Transformer layers with the pre-norm, self-attention, and MLP. We refer the readers to (Vaswani et al. 2017; Carion et al. 2020) for more details on the Transformer encoder and decoder. Figure 2 illustrates our CAT model for 3D automatic annotation.

CAT encoder comprises an intra-object encoder (local) and an inter-object encoder (global) to perform self-attention along the sequence and batch dimensions. The local encoder explicitly models local contexts among all-pairs points at an object level. The resulting features are then fed to the global encoder to capture the contextual interactions at a scene level for further point feature learning.

To generate high-quality 3D bounding boxes, CAT should focus on essential parts of point features for 3D box regression. We introduce a decoder for the box to one-way attends point features to refine object-level features for the final box regression. We take box tokens as queries and point features as contexts from global encoder output $(box\_tokens,point\_features)$ to query point features.

Following (Carion et al. 2020), our decoder operates $B$ objects in parallel using multiheaded self- and cross-attention mechanisms in one batch $B$. Box queries are learned $box\_tokens$ from global encoder output. Each query has 7 box tokens in the form of $(x,y,z,width,length,height,yaw)$. The decoder takes the $point\_features$ as keys and values $(K,V)$ and $box\_tokens$ as queries $(Q)$. The resulting decoder outputs are then used to generate boxes, ${Box}_{3D}\in\mathbb{R}^{B\times 7}$. In our framework, box queries represent 3D boxes in 3D space around which our final 3D boxes in $B\times 7$ are generated by three MLP heads to predict locations in $(x,y,z)$, dimensions in $(width,length,height)$ and rotation along the z-axis $yaw$, respectively. Each MLP head consists of one linear layer and leaky ReLU nonlinearity. Thus, the Transformer decoder in (Carion et al. 2020) has not many specific modifications for 3D.

Using self- and cross-attention over the global encoder output, the CAT decoder globally refines $B$ object-level features for 3D box regression, while being able to use $B$ objects’ global point features as contexts.

To train the model, we use a loss directly on the dIoU as in (Zheng et al. 2020) for each generated box and its matched ground truth box denoted as $\mathcal{L}_{box}$. Moreover, we introduce a direction loss $\mathcal{L}_{dir}$ for a binary direction classification, since the IoU metric is direction-invariant (Liu et al. 2022a). The direction classification is performed by an MLP where orientation within $[-\pi/2,\pi/2)$ is predicted as the front, and $[\pi/2,\pi]\cup[-\pi,-\pi/2)$ as the back. This encourages our model to learn the right box direction, a property that helps our regression on yaws. We use Cross-Entropy loss for the direction loss $\mathcal{L}_{dir}$.

$$\mathcal{L}=\lambda_{box}\mathcal{L}_{box}+\mathcal{L}_{dir},$$

(3)

where the $\lambda_{box}$ is a weight, we empirically set it as 5.

Our loss function is a weighted sum of two terms above. CAT removes many handcrafted designs on the loss function while being way simpler to implement and understand.

Standard positional encoding for vision tasks is crafted manually, e.g., sinusoidal position embedding (Vaswani et al. 2017) based on sine and cosine functions and normalized range values (Vaswani et al. 2017). In 3D point cloud processing, 3D points contain information about $(x,y,z)$ coordinates that are natural positional information. For brevity, we go beyond this by introducing an MLP for trainable and parameterized positional embeddings. Trainable positional embedding on 3D coordinates is essential to adapt the 3D structure for point clouds. Notably, our model’s more straightforward, parameterized, trainable positional embedding outperforms other position encoding schemes. An MLP is used to encode the 3D position $pos$ as:

$$pos=MLP(x,y,z).$$

(4)

We concatenate box tokens and points in our sequence, $concat(box\_tokens,pos)$. For dimension consistency, we manipulate another MLP with one linear layer in $pos$ to describe the position of each entity in the sequence.

We now evaluate the performance of CAT as an auto labeler on the KITTI dataset. CAT is trained with 500 frames of labeled 3D data to generate 3D pseudo labels from weak 2D boxes of the KITTI dataset. CAT-generated 3D pseudo labels are used to train the 3D object detector of PointRCNN. We compare it to other baselines on the KITTI test and val sets, as shown in Table 1 and Table 2.

We conduct a series of experiments to confirm the effectiveness of each module in our proposed model, namely, the local encoder, global decoder, and decoder. Meanwhile, we also study the choice of position embedding strategies for CAT training. The 3D pseudo labels are compared with human annotations on the metrics of mIoU, recall with IoU threshold of 0.7, mAP, and mAP_R40 over the Car category.

In Table 3, we regard the standard Transformer encoder (Local Enc.) as a baseline. Model A (Local Enc.) shows the local encoder extracts intra-object interactions among points. Model B (Local+Global Enc.) brings a performance boost of 5.6%, 10.7%, 26.8%, 26.6% on mIou, Recall, mAP, and mAP_R40, demonstrating that global encoder models inter-object relations and facilitates samples’ interactions, particularly helpful on hard samples. Model C (Local+Global Enc.+Decoder) further improves mIoU by up to 2.31%, clarifying that the decoder helps to refine representations for more precise box regression. Model D (Local+Global Enc.+Decoder+Pos Embed.) achieves the final 72.78% Recall and 87.04% mAP in an MLP embedding way, verifying positional encoding on 3D coordinates is essential to adapt to the 3D structure. We can observe that the more straightforward parameterized and trainable positional embedding of MLP outperforms other hand-designed position encoding schemes in our model. Meanwhile, we evaluate the effectiveness of box tokens in CAT. Seven Box tokens represent seven box elements (XYZ localization, HWL size, and yaw orientation), as box representations. Removing box tokens in the CAT encoder results in a 2.69% IoU drop, suggesting queries from box tokens instead of position encoding in (Carion et al. 2020) can help to refine object-level features for the final 3D box generation in the decoder.