Transformer based 3D tooth segmentation via point cloud region partition

Abstract

Abstract Automatic and accurate tooth segmentation on 3D dental point clouds plays a pivotal role in computer-aided dentistry. Existing Transformer-based methods focus on aggregating local features, but fail to directly model global contexts due to memory limitations and high computational cost. In this paper, we propose a novel Transformer-based 3D tooth segmentation network, called PointRegion, which can process the entire point cloud at a low cost. Following a novel modeling of semantic segmentation that interprets the point cloud as a tessellation of learnable regions, we first design a RegionPartition module to partition the 3D point cloud into a certain number of regions and project these regions as embeddings in an effective way. Then, an offset-attention based RegionEncoder module is applied on all region embeddings to model global context among regions and predict the class logits for each region. Considering the irregularity and disorder of 3D point cloud data, a novel mechanism is proposed to build the point-to-region association to replace traditional convolutional operations. The mechanism, as a medium between points and regions, automatically learns the probabilities that each point belongs to its neighboring regions from the similarity between point and region features, achieving the goal of point-level segmentation. Since the number of regions is far less than the number of points, our proposed PointRegion model can leverage the capability of the global-based Transformer on large-scale point clouds with low computational cost and memory consumption. Finally, extensive experiments demonstrate the effectiveness and superiority of our method on our dental dataset.