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

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

Keywords: Semi-supervised learning, U-Mamba2-SSL, CBCT Imaging, Tooth and Pulp Segmentation, STSR 2025 Challenge

Abstract:

Accurate segmentation of teeth and pulp in Cone-Beam Computed Tomography (CBCT) is vital for clinical applications like treatment planning and diagnosis. However, manual segmentation requires extensive expertise and is exceptionally time-consuming, highlighting the critical need for automated semi-supervised segmentation algorithms that can 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 obtained 0.917 DSC and 0.948 mIoU 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.

Changes Summary:

Response to Reviewer YQeT
> 1.Ambiguity in “average score”:
We have updated the abstract to report DSC and mIOU, which are more commonly known instead of the “average score”.
> 2.Discrepancy in test set metrics:
We have added a short discussion in Section 4.4 Limitations to discuss the discrepancy in the test set metrics.
> 3.Section 2.2 contains a run-on sentence beginning with “For an unlabeled training data, it is first passed through the model…” (line ~4 of Section 2.2);
Thank you for pointing out the grammatical mistake. We have amended the error in the revision.
> 4.In Table 3, the row labeling is inconsistent:
Thank you for pointing out the confusion. The variants of U-Mamba2-SSL in Table 3 refers to the different incremental ablation of our method. We have modified the first paragraph in Section 4.1 and the table caption to clarify the differences.

Response to Reviewer eA1Q
Thank you for the reviews. There were some repeated comments (i.e. Point 1 - 5) that we have combined and addressed as a single comment.
> Abstract (Lines 5–12): “highlighting the critical need for automated algorithms that can effectively utilize unlabeled data” is generic; specify the concrete limitation addressed.
Thank you for the comment. We have made the contributions concrete by specifying the application of the proposed method.
> Figure 1 caption (Page 3): Caption is overly long; standardize subfigure labels as “(a), (b), (c)” and reduce verbosity.
We have updated the subfigure labels and simplified the figure caption by moving the description to the main text.
> Section 3.2 (Implementation Details, Page 6): Rephrase “The input patch is randomly cropped by ensuring that at least 33% of the samples have a foreground label” for clarity (e.g., “Crop patches so ≥33% contain foreground”).
Thank you for pointing that out. We have rephrased and clarified the ambiguous sentence.
> Table 3 (Page 7): Harmonize abbreviation capitalization in the note (“CR denotes Consistency Regularization; PL denotes Pseudo Label”).
Thank you for the suggestion. We have amended the table caption to make the abbreviation consistent.
> Section 4.4 (Limitations, Page 8): “Future work should design data processing and augmentation techniques…” is verbose; streamline wording.
We have shortened the sentences in Section 4.4 to make the discussion more compact.
> Checklist Table 4 (Page 12): Remove the instruction “Delete this Table in the camera-ready submission” before final submission.
For the camera-ready revision, we have removed the checklist table. Thank you for the reminder.

Main Tex File: main.tex

Confirm Latex Only: true

Code Url: https://github.com/zhiqin1998/UMamba2

Authors Changed: false

Arxiv Update Plans: Yes, https://arxiv.org/abs/2509.20154

Submission Number: 21