Semi-supervised 3D Tooth Segmentation Using nn-UNet with Axial Attention and Positional Correction

Qiupu Chen, Yimou Wang, Jun Xu, Qiankun Li

Abstract: Accurate 3D tooth segmentation plays a pivotal role in the realm of computer-aided dental diagnosis and treatment. Despite its significance, the scarcity of labeled 3D tooth data poses a substantial challenge. This paper details our contribution to the MICCAI 2023 semi-supervised teeth segmentation challenge, aimed at enhancing the precision of tooth segmentation through the application of deep learning within a semi-supervised learning framework. Our approach leverages the nn-UNet architecture, incorporating innovative modifications to improve segmentation performance. Notably, we introduce two novel components: an axial attention mechanism module and a positional correction module. The axial attention mechanism enhances the model’s ability to capture contextual information in axial slices, contributing to improved segmentation accuracy. Simultaneously, the positional correction module solves the problem of incorrect segmentation of bony structures similar to tooth morphology and density. In the context of semi-supervised learning, where labeled data is limited, we propose a robust selection methodology for pseudo labels. This methodology considers the stability of pseudo labels across re-training iterations, ensuring the reliability of the learning process. The integration of these components and methodologies collectively enhances the model’s adaptability to the challenges posed by limited labeled data. The proposed model ranked fifth in the final ranking on unseen test data, with a score of 81.47%. The code is available at https://github.com/qpuchen/nnUNet_att_position_correction.