OctAttention: Octree-Based Large-Scale Contexts Model for Point Cloud Compression

Introduction

The point cloud is an essential data structure for 3D representation. It has been used in many fields such as virtual reality, smart city, robotics, and autonomous driving (Schwarz et al. 2018). Since massive point clouds are generated, efficient compression techniques are necessary for point cloud storage and transmission. However, point clouds are unordered and have various distributions; it is relatively difficult to compress point clouds compared with 2D images. Fortunately, several schemes for point cloud geometry and attribute compression such as voxel-based, image-based and tree-based algorithms have been proposed and applied in research works (Chou, Koroteev, and Krivokuća 2019; Shao et al. 2017, 2018) and standard specification (3DG 2021).

The MPEG point cloud geometry compression standard (G-PCC) (Schwarz et al. 2018) adopted a hand-crafted context-adaptive arithmetic encoder for bit allocation, which can be seen as a prediction for the currently encoding node based on the coded information. Recently, entropy encoders based on deep learning have been shown to outperform hand-crafted ones on rate-distortion performance. Among them, some methods partition point clouds into voxels, then adopt 3D convolution to learn and predict the occupancy of each voxel (Wang et al. 2021; Nguyen et al. 2021a; Quach, Valenzise, and Dufaux 2019; Que, Lu, and Xu 2021). Voxel-based models are capable of exploiting the local geometric patterns (e.g., planes, surfaces). However, they are not robust to point cloud resolution. These methods have to divide blocks into different scales for point clouds with varying point densities to find the optimal voxel size. Meanwhile, receptive fields are limited by the computational cost, i.e., they only extract features from voxels within a narrow range. Other works encode the point cloud into an octree, then encode the occupancy of octree nodes based on their ancestor nodes (Huang et al. 2020; Biswas et al. 2020). The octree-based model is robust to resolution, and it also utilizes a broader range of contexts than voxel-based ones. However, prior methods ignore that sibling nodes (i.e., nodes in the same octree level) provide low-level local geometry features, which are significant to exploit the geometry redundancy. In general, prior voxel-based and octree-based methods do not fully use much spatial context information.

In this work, we propose a point cloud compression method called OctAttention which generates and utilizes a large-scale context. Voxel is inefficient for representing sparse point clouds; thus, we encode the point cloud into an octree. Subsequently, we improve the predictability of the occupancy of each octree node based on large-scale context, which contains features from ancestor nodes of the current node, sibling nodes, and ancestors of sibling nodes. However, it should be noted that the side effect of expanding context is to introduce redundant and irrelevant information. For example, different sibling nodes may have the same ancestors, but they are repeated multiple times in the context; sibling nodes far from the current node may be worthless for prediction. To tackle this problem, we employ tree-structured attention to weight and explicitly express the contributions of different nodes in the prediction. The siblings in the context disable the parallelization strategy in prior works (Huang et al. 2020; Que, Lu, and Xu 2021), hence we propose a mask operation to encode multiple nodes in parallel.

We compare the proposed model with state-of-the-art methods on the 3D LiDAR dataset SemanticKITTI (Behley et al. 2019), object point cloud dataset MVUB (Charles et al. 2016) and MPEG 8i (Eugene et al. 2017). The experiments show that our method outperforms these state-of-the-art methods, which are only designed for a specific category of point clouds.

The contributions of our work can be summarized as:

We propose a tree-structured attention mechanism to model the dependency of nodes in a large-scale context, which is achieved by extending the receptive field of context and exploiting features from sibling nodes and their ancestors.

We employ a mask operation to encode octree in parallel to alleviate the drawbacks of introducing siblings in the large-scale context.

Our generic model of point cloud geometry compression for both LiDAR and object point clouds achieves state-of-the-art performance on several large-scale datasets.

Related Work

Voxel-based methods quantize the point cloud and classify the voxel occupancy by neural networks. Voxel-based methods outperform G-PCC (3DG 2021) on lossless geometric compression (Nguyen et al. 2021a; Quach, Valenzise, and Dufaux 2019), lossy geometric compression (Quach, Valenzise, and Dufaux 2019, 2020; Wang et al. 2021) and progressive compression (Guarda, Rodrigues, and Pereira 2020). Compared to an octree, geometric patterns can be naturally preserved in the voxel representation. Yet, the side effect is voxel-based networks are sensitive to the density variation and may fail for the sparse point clouds. All of the above methods are applied to dense point clouds (e.g. MPEG 8i) and may suffer tremendous computing and memory costs on sparse LiDAR point clouds. The proposed method directly processes the octree occupancy code to overcome the density variation problem.

Tree-Based Point Cloud Compression

Tree structures effectively reduce geometric redundancy by merging the common coordinates of point clouds. Numerous algorithms (Lasserre, Flynn, and Qu 2019; Zhang, Gao, and Liu 2020; Schwarz et al. 2018; Kammerl et al. 2012; Kathariya et al. 2018b; Gao et al. 2019) compressed point cloud based on tree structures such as octree (Schnabel and Klein 2006), quadtree (Kathariya et al. 2018a), KD tree(Devillers and Gandoin 2000), prediction tree (Gumhold et al. 2005), etc. Recently, many works have focused on designing octree context for arithmetic coding to compress bitstream. (Song et al. 2021) aggregated voxels in reduced space by removing free regions and acquires a compact context model for LiDAR compression. (Garcia and de Queiroz 2018) reordered the node sequences to improve the intra-frames lossless geometric compression.

All of the above methods model the context by hand-crafted features. (Lei, Akhtar, and Mian 2019; Riegler, Osman Ulusoy, and Geiger 2017; Wang et al. 2017) mainly introduce new convolution methods under the octree framework. OctSqueeze (Huang et al. 2020) is the first octree-based deep learning entropy model by modeling the dependency among node and its multiple ancestor nodes. MuSCLE (Biswas et al. 2020) reduced the temporal redundancy by exploiting spatio-temporal relationships across LiDAR sweeps. Both methods avoid high computational complexity, yet the strong dependency among sibling nodes is ignored. VoxelContext-Net (Que, Lu, and Xu 2021) partly solved this problem by employing voxel and octree hybrid structure to learn the context in the previous octree depth. However, features from higher resolution (i.e., from sibling nodes) are still ignored. Besides, as shown in Fig. 1, given a fixed-size voxel-based context, its receptive field shrinks with the increased octree depth. Due to the computational overhead, the receptive field of the voxel-based approach is limited, which restricts the ability to model the context. Our proposed method can acquire more than $500\times 500\times 500$ voxels receptive field ($9\times 9\times 9$ in VoxelContext-Net). Meanwhile, we introduce sibling nodes and their ancestors in the context. The extended context contains more potentially helpful information to model the distribution of octree nodes.

Methodology

We propose an extended context and a tree-structured attention mechanism, which is shown in Fig. 2. Intuitively, nodes with similar ancestors and siblings tend to follow a similar distribution. Therefore, the proposed context exploits features from siblings and their ancestors, which are beneficial for inference. To achieve accurate and flexible prediction with a large-scale context, we employ a tree-structured attention mechanism to determine the importance of each node in the context. Finally, we infer the occupancy of each octree node based on the attention context. We further propose a mask operation to counteract the increased coding time caused by the involvement of the sibling context.

Context Model

We propose an expanded large-scale context to achieve more accurate probability estimation. We first traverse the octree in a breadth-first order. Then for each currently encoding node $n_{i}$ in the sequence, we construct a context window $\left\{{n}_{i-N+1},\ldots,{n}_{k},\ldots,{n}_{i-1},{n}_{i}\right\}$ with the length of $N$ (see Fig. 2 Left). The currently encoding node is at the end of the window, and the local context window slides forward with the currently encoding node moving forward. In this manner, we select $N-1$ sibling nodes to exploit the strong dependency among the nodes at the same depth. Considering the dependency between the nodes and their ancestor nodes, we further introduce $K-1$ ancestors of the $N$ nodes in the context window, respectively. In summary, we integrate $N*K$ related nodes and greatly expand the contexts. Specifically, we factorize distribution $Q(\mathbf{x})$ into a product of conditional probabilities of each occupancy symbol $x_{i}$ as:

where $x_{i}\in\{1,2,\ldots,255\}$ is the occupancy symbol of currently encoding node $n_{i}$. High-dimensional vector $\boldsymbol{f}_{k}=[\boldsymbol{h}_{k}^{(0)},\boldsymbol{h}^{(1)}_{k},\boldsymbol{h}^{(2)}_{k},...,\boldsymbol{h}^{(K-1)}_{k}]$ denotes the concatenation of embedding features of nodes $n_{k}$ and features of its $K-1$ ancestors, which are defined as the embedding of their occupancy, depth and octant. $\mathbf{w}$ denotes the context model parameters.

Voxel-based methods naturally exploit the low-level geometry features preserved by a voxel representation. However, in an octree, these low-level features are hidden in the octree nodes. We employ an embedding for each node in the context window to reveal these features. Node information is embedded to a fixed-length vector respectively and then concatenated as $\boldsymbol{h}_{k}=\left[S_{k}W_{1},L_{k}W_{2},O_{k}W_{3}\right]$, where $S_{k},L_{k},O_{k}$ are one-hot coded occupancy, level index and octant index, and $W$ is their respective embedding matrix. Embedding also serves as normalization for the three different scale variables. It should be noted that the depth and octant of $n_{i}$ are already available while decoding $n_{i}$, yet its occupancy code $x_{i}$ is unknown, so we pad it with $x_{i-1}$.

In this manner, we utilize the features from sibling nodes and their ancestors, which are significant to prediction. Previous octree-based works (Huang et al. 2020; Biswas et al. 2020) ignored it. In these works, node $n_{i}$ is assumed to be conditional independence under the condition $\boldsymbol{f}_{i}$. We strengthen this condition by expanding the context. Thus we make a more general assumption. Meanwhile, we avoid the inefficient sparse context and tremendous incremental computations caused by expanding the context in previous voxel-based works (Que, Lu, and Xu 2021). The receptive field of the proposed context can exceed 1000 octree nodes, which is extremely difficult to be achieved in voxel-based methods.

Tree-Structured Attention

where $n\leq m,m={i-N+1,...,i}$. The summation for node $n_{m}$ ends at $m$ since a mask operation is applied to the attention map to achieve fast encoding, which is discussed in the next section. With attention mechanism, we can draw weighted context as:

In summary, the contexts are fed to 2 layers of multi-head self-attention $\text{MultiHead}(\boldsymbol{f},i)=\text{MLP}([\boldsymbol{C}_{i}^{(0)},...,\boldsymbol{C}_{i}^{(K-1)}])$ and multi-layer perception (MLP) successively, and finally outputs a 255-dimensional probability for $n_{i}$:

Mask operation

The estimated probability is adopted to guide the arithmetic coder, which codes the octree nodes sequentially in a lossless way. Previous methods (Huang et al. 2020; Que, Lu, and Xu 2021) excluded siblings from contexts and only depend on ancestors. Hence they can naturally parallelize encoding and decoding within each level. Without mask operation, as shown in Fig. 3(a), given a sliding window, we adopt all $N-1$ siblings in the context to predict the last node in the context window. It is difficult to achieve the same parallelization since we can only encode the last node in one propagation. To reduce encoding time, we introduce a mask operation in Eq. (5) which assigns a varied-length receptive field for each node. Each node is restricted to the access of the previous nodes in the context window at training and testing. As shown in Fig. 3(b), in this way, Eq. (7) can equally apply to the last $N_{0}$ nodes. Hence we are allowed to encode them simultaneously in one propagation. While encoding the $j^{\text{th}}$ node in the window, only $j$ nodes are available for inference. Compared to the way of maximum receptive field in Fig. 3(a), on average, the receptive field of each node shrinks from $N$ to $(2N-N_{0}+1)/2$. Nevertheless, the coding time reduces by $N_{0}$ times. The parameter $N_{0}$ balances the receptive field and coding time. Although the receptive field shrinks, we find there is negligible performance loss as the result of the mask operation during the training.

Learning

We optimize the cross-entropy between the predicted occupancy code and ground-truth, which is defined as:

Here, $q_{i}\left(x_{i}\mid\boldsymbol{f};\mathbf{w}\right)$ is the estimated probability of occupancy code $x_{i}$ at node $n_{i}$, which is defined in Eq. (7).

Experiments

SemanticKITTI (Behley et al. 2019) is a large sparse LiDAR dataset for self-driving. It is collected from a Velodyne HDL-64E sensor and contains 43552 scans with 4549 million points. Following VoxelContext-Net (Que, Lu, and Xu 2021), we normalize the raw data into $^{3}$ as reference point cloud and use sequences 00 to 10 (including 23201 scans) for training, and sequences 11 to 21 (including 20351 scans) for testing.

Object Point Cloud Dataset

Microsoft Voxelized Upper Bodies (MVUB) (Charles et al. 2016) is a dynamic voxelized point cloud dataset containing five half-body subjects sequences with 9 and 10-bit precision. 8i Voxelized Full Bodies (MPEG 8i) (Eugene et al. 2017) includes sequences of smooth surface and complete human shape point clouds with 10 and 12-bit precision. Following the setting of VoxelDNN (Nguyen et al. 2021a), we use Andrew10, David10, Sarah10 sequences in MVUB, Soldier10 and Longdress10 sequences in MPEG 8i for training. We select several point cloud sequences with different resolutions for testing. All testing point clouds were not used during training.

Experimental Details

Training and Testing Strategy

For the static LiDAR compression, we train a single model with the max octree depth of 12 as our model can learn the distribution of all layers in one model. While testing, we truncate the octree over 8-12 levels to evaluate our model at different bitrates. For object dataset compression, we train one model using the octree sequence data converted from point clouds at depths 9 and 10. We also evaluate our model on data with different geometry precision to verify robustness. We implement our model in PyTorch and perform the training/testing with Xeon E5-2637 CPU and one NVIDIA TITAN Xp GPU (12G memory). We use batch sizes of 32, epochs of 8 and Adam optimizer with a learning rate of 1e$-3$. It takes 2 days to train our model in each experiment. Occupancy, level index, and octant index are embedded into 128, 6, and 4 dimensions, respectively. We set $K=4$, $N=N_{0}=1024$ and use 2 layers and 4 heads in multi-head self-attention in experiments unless otherwise specified.

Evaluation Metrics

It is important to adopt the same evaluation metrics to make a fair comparison. Following the MPEG standards (Schwarz et al. 2018), we use two standard metrics which measure geometry reconstruction quality named point to point PSNR (D1 PSNR) and point to plane PSNR (D2 PSNR) in lossy geometry compression. Both of them can be calculated by the MPEG tool pc_error. We estimate the normal at each point using the MATLAB function pcnormals. We also report chamfer distance (CD) and set PSNR peak value $r=1$ following VoxelContext-Net. We correct its results from correspondence with the authors by eliminating inconsistencies in the PSNR formula. As for the object dataset, we adopt the official default configuration in TMC13. We use bits per point (bpp) to measure the performance. Unless otherwise specified, all distortion curves and bitrates are obtained by averaging over sequence.

Experiment Results

The rate-distortion curves of LiDAR compression are shown in Fig. 4. We compare our method with VoxelContext-Net without coordinate refinement model (i.e., VoxelContext-Net w/o CRM) for fairness, as post-processing is irrelevant to compression performance evaluation. Our method outperforms other baselines at all bitrates. On average, our approach (i.e., OctAttention) saves 25.4% bitrates on SemanticKITTI compared with G-PCC, while OctSqueeze only saves less than 4% bitrates. Our method achieves more than 11% relative reduction in bitrate versus the state-of-the-art method VoxelContext-Net at high bitrates. It may be due to the voxel-based method failing in the sparse scenario that lacks occupied voxels. The experiment results demonstrate the effectiveness of our large receptive field context model.

Results for Object Point Cloud Compression

In Table 1, we provide the bpp results for lossless compression on object point clouds. Our method outperforms VoxelDNN and achieves a 32.8% gain over G-PCC on average.

Ablation Study and Analysis

We perform an ablation experiment on SemanticKITTI to demonstrate the effectiveness of a large receptive field context. We set $N_{0}=N$ so that each node is predicted only once and the average receptive field is $(N+1)/2$. We then alter the context window size $N$, and the number of parameters in our model remains unchanged. As shown in Table 2, we can save 14% bitrates by enlarging the context window size from 8 to 1024. The encoding time decreases with increased context window size due to our mask operation, where we decrease I/O by $N_{0}$ times. Decoding time does not increase significantly since the time consumption is primarily in I/O. To balance the decoding time and performance, we set $N=1024$.

Effectiveness of Attention and Sibling Context

As the visualization in Fig. 5, the attention mechanism discovers the similarity among points in a context window according to geometry patterns such as line, plane, surface, and curvature. The node with the highest attention score to the currently encoding nodes (red points) is colored in yellow. It confirms that the attention mechanism can predict occupancy by integrating similar features from sibling nodes in a large-scale context. In Table 3, we set $N=256$ and removed the attention and sibling features respectively for the ablation study. The results further illustrate the effectiveness of attention mechanism and sibling-involved context.

Robustness to Varying Point Densities

To evaluate OctAttention robustness to varying point densities, we test point clouds with additional geometry precisions and point densities. See Fig. 6. VoxelDNN only outperforms G-PCC on the geometry precision of 10. Its performance drops on sparse point clouds due to the lack of occupied voxels in the context. Since we adopt an octree structure with a fixed-length context window, our method’s performance is shown to be stable with varying point densities.

Runtime

The number of parameters in our model is 2.67M. See table 4. Our method saves 95% encoding time and 91% decoding time compared with VoxelDNN (Nguyen et al. 2021a). Our approach can be applied to real-time point cloud encoding and offline point cloud decompression. We believe it is possible to speed up the decoding by dividing the octree into disjoint subtrees and developing a GPU-based algorithm like an arithmetic encoder.

Conclusion

We proposed a novel octree-based entropy model called OctAttention for sparse and dense point cloud geometry compression by exploiting large-scale contexts. Specifically, we extend the context and introduce sibling nodes in the octree. We employ the attention mechanism to emphasize the significant nodes to utilize these abundant features. We further propose a mask operation to achieve parallel encoding under the condition of introducing siblings in the context. We evaluate our method on both the LiDAR and object point cloud datasets. The results demonstrate that the proposed method achieves state-of-the-art on both types of datasets.

Acknowledgements

This work was supported by the National Natural Science Foundation of China (No. 62172021, 61801303, 62031013), the Shenzhen Fundamental Research Program (GXWD20201231165807007-20200806163656003), Guangdong Basic and Applied Basic Research Foundation (2019A1515012031), Shenzhen Science and Technology Plan Basic Research Project (JCYJ20190808161805519).

References

Appendix 1. Compression performance at high bitrates

Our model is trained on 12 depth octrees and tested on 11-16 depth data, which is almost non-existent in the training data. As shown in Fig. 7, our method outperforms OctSqueeze at all bitrates and saves 19.4% bitrates on average. At the highest bitrate, our method saves 13.5% bitrates compared with OctSqueeze and saves 8.5% bitrates compared with MuSCLE. It should be noted that temporal correlations are utilized in MuSCLE to encode dynamic point clouds, even though OctAttention does not make use of inter-frame information. Nevertheless, we achieve comparable compression performance at low bitrates and better performance at high bitrates. The experimental results demonstrate the effectiveness and robustness of our method at high bitrates.

Appendix 2. Visualization of embedding

Embedding is supposed to discover the local geometry patterns which are hidden in octree representation. To demonstrate its effectiveness, we visualize the embedded feature $\boldsymbol{f}_{k}=[\boldsymbol{h}_{k}^{(0)},\boldsymbol{h}^{(1)}_{k},\boldsymbol{h}^{(2)}_{k}$$,...,\boldsymbol{h}^{(K-1)}_{k}]$ of node $n_{k}$, where $\boldsymbol{h}^{(0)}_{k}$ represents the feature $\left[S_{k}W_{1},L_{k}W_{2},O_{k}W_{3}\right]$ of node $n_{k}$ and $\boldsymbol{h}^{(1)}_{k},\boldsymbol{h}^{(2)}_{k}$$,...,\boldsymbol{h}^{(K-1)}_{k}$ represent the features of its $K-1$ ancestors. $S_{k},L_{k},O_{k}$ are one-hot coded occupancy, level index and octant index, and $W$ is their respective embedding matrix. We employ Principal Component Analysis to reduce $\boldsymbol{f}_{k}$ dimension to 3 and then color corresponding points with the 3-dimension vectors. We also visualize the original feature of node $n_{k}$, which are defined as the concatenation of $[S_{k},L_{k},O_{k}]$ and features of its $K-1$ ancestors. We implement our experiment on ModelNet40 (Wu et al. 2015).

Fig. 8 shows that embedding can describe local features such as surfaces, orientation, and curvatures. For example, original features are unable to distinguish surfaces explicitly or discover parallel planes. Points within the same surface (see tent, car, bed, radio, and airplane) or parallel planes (see sink, bookshelf, and stairs) may have varying colors, resulting from the inconsistency between the occupancy code value and geometry patterns. On the contrary, similar geometry patterns are rendered in similar colors by embedding. More specifically, the colors of the upper and lower surfaces are blue and green, and the colors of the left and right sides are yellow and red. In summary, embedding can serve as an enhanced local geometry descriptor, which is helpful to model point cloud geometry distribution.