Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

U-Mamba2-SSL for Semi-Supervised Tooth and Pulp Segmentation in CBCT

Zhi Qin Tan, Xiatian Zhu, Owen Addison, Yunpeng Li

TL;DR

To address the scarcity of labeled CBCT data for precise tooth and pulp segmentation, this work introduces U-Mamba2-SSL, a three-stage semi-supervised framework built on U-Mamba2. It first pre-trains the model with a disruptive autoencoder using all data, then trains with a mix of labeled supervision and unlabeled consistency regularization through carefully designed input and feature perturbations, and finally refines the model with pseudo labeling on unlabeled data under a confidence threshold. The approach achieves state-of-the-art performance in STSR 2025 Task 1, with an average score of $0.789$ and a Dice score of $0.917$ on the hidden test set, and demonstrates substantial gains in Identification Accuracy over baselines. The combination of self-supervised pre-training, robust perturbation-based regularization, and selective pseudo labeling offers a practical path to high-precision 3D tooth/pulp segmentation in CBCT while reducing labeling requirements and improving generalization for clinical workflows.

Abstract

Accurate segmentation of teeth and pulp in Cone-Beam Computed Tomography (CBCT) is vital for clinical applications like treatment planning and diagnosis. However, this process requires extensive expertise and is exceptionally time-consuming, highlighting the critical need for automated algorithms that can effectively utilize unlabeled data. In this paper, we propose U-Mamba2-SSL, a novel semi-supervised learning framework that builds on the U-Mamba2 model and employs a multi-stage training strategy. The framework first pre-trains U-Mamba2 in a self-supervised manner using a disruptive autoencoder. It then leverages unlabeled data through consistency regularization, where we introduce input and feature perturbations to ensure stable model outputs. Finally, a pseudo-labeling strategy is implemented with a reduced loss weighting to minimize the impact of potential errors. U-Mamba2-SSL achieved an average score of 0.789 and a DSC of 0.917 on the hidden test set, achieving first place in Task 1 of the STSR 2025 challenge. The code is available at https://github.com/zhiqin1998/UMamba2.

U-Mamba2-SSL for Semi-Supervised Tooth and Pulp Segmentation in CBCT

TL;DR

To address the scarcity of labeled CBCT data for precise tooth and pulp segmentation, this work introduces U-Mamba2-SSL, a three-stage semi-supervised framework built on U-Mamba2. It first pre-trains the model with a disruptive autoencoder using all data, then trains with a mix of labeled supervision and unlabeled consistency regularization through carefully designed input and feature perturbations, and finally refines the model with pseudo labeling on unlabeled data under a confidence threshold. The approach achieves state-of-the-art performance in STSR 2025 Task 1, with an average score of and a Dice score of on the hidden test set, and demonstrates substantial gains in Identification Accuracy over baselines. The combination of self-supervised pre-training, robust perturbation-based regularization, and selective pseudo labeling offers a practical path to high-precision 3D tooth/pulp segmentation in CBCT while reducing labeling requirements and improving generalization for clinical workflows.

Abstract

Accurate segmentation of teeth and pulp in Cone-Beam Computed Tomography (CBCT) is vital for clinical applications like treatment planning and diagnosis. However, this process requires extensive expertise and is exceptionally time-consuming, highlighting the critical need for automated algorithms that can effectively utilize unlabeled data. In this paper, we propose U-Mamba2-SSL, a novel semi-supervised learning framework that builds on the U-Mamba2 model and employs a multi-stage training strategy. The framework first pre-trains U-Mamba2 in a self-supervised manner using a disruptive autoencoder. It then leverages unlabeled data through consistency regularization, where we introduce input and feature perturbations to ensure stable model outputs. Finally, a pseudo-labeling strategy is implemented with a reduced loss weighting to minimize the impact of potential errors. U-Mamba2-SSL achieved an average score of 0.789 and a DSC of 0.917 on the hidden test set, achieving first place in Task 1 of the STSR 2025 challenge. The code is available at https://github.com/zhiqin1998/UMamba2.

Paper Structure

This paper contains 24 sections, 2 equations, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Method
  3. Problem Formulation.
  4. First Stage: Pre-training with Disruptive Autoencoder
  5. Second Stage: Consistency Regularization Training
  6. Input Perturbations.
  7. Feature Perturbations.
  8. Semi-Supervised Learning Schedule.
  9. Third Stage: Pseudo Labeling
  10. Experiments
  11. Dataset and Evaluation Metrics
  12. Implementation Details
  13. Preprocessing.
  14. Environment settings.
  15. Training protocols.
  16. ...and 9 more sections

Figures (3)

  • Figure 1: Overall diagram of the proposed U-Mamba2-SSL framework. (a) U-Mamba2 is first pre-trained by reconstructing inputs corrupted with noise, downsampled, and masked; (b) The second stage involves a combination of supervised loss for the labeled data and consistency regularization between the unperturbed output and the perturbed output of the unlabeled data; (c) The final stage introduces pseudo labeling on top of the training objectives in (b). Only pseudo labels with confidence above a certain threshold contribute to the training loss.
  • Figure 2: Qualitative results of U-Mamba2-SSL on the internal validation set. The 3D render and a representative 2D slice are shown for: (Top) the best scoring case and (Bottom) the worst scoring case.
  • Figure 3: (Left): Effect of the tile size on the metrics with '1,2' mirror axes in TTA. (Right): Effect of various mirror axes combinations in TTA on the metrics when tile size is set to 0.9. Axis definition: '0' is superior/inferior, '1' is anterior/posterior, and '2' is left/right.