SelfMedHPM: Self Pre-training With Hard Patches Mining Masked Autoencoders For Medical Image Segmentation

TL;DR

This work tackles CT multi-organ segmentation by addressing a limitation of MAE-based pretraining: it does not focus on hard-to-reconstruct regions. It introduces SelfMedHPM, a self-training framework with hard patches mining, an auxiliary loss predictor, and an EMA teacher–student architecture to steer mask generation and patch reconstruction. The method achieves state-of-the-art average Dice scores on BTCV abdominal and SMWB body segmentation (e.g., approximately $85.8\%$ and $90.9\%$ DSC, respectively), outperforming MAE-based pretraining while maintaining a compact training overhead. By combining the loss-prediction objective and an easy-to-hard masking schedule, SelfMedHPM improves segmentation robustness in challenging regions and has potential for broader medical image analysis tasks and prognosis applications.

Abstract

In recent years, deep learning methods such as convolutional neural network (CNN) and transformers have made significant progress in CT multi-organ segmentation. However, CT multi-organ segmentation methods based on masked image modeling (MIM) are very limited. There are already methods using MAE for CT multi-organ segmentation task, we believe that the existing methods do not identify the most difficult areas to reconstruct. To this end, we propose a MIM self-training framework with hard patches mining masked autoencoders for CT multi-organ segmentation tasks (selfMedHPM). The method performs ViT self-pretraining on the training set of the target data and introduces an auxiliary loss predictor, which first predicts the patch loss and determines the location of the next mask. SelfMedHPM implementation is better than various competitive methods in abdominal CT multi-organ segmentation and body CT multi-organ segmentation. We have validated the performance of our method on the Multi Atlas Labeling Beyond The Cranial Vault (BTCV) dataset for abdomen mult-organ segmentation and the SinoMed Whole Body (SMWB) dataset for body multi-organ segmentation tasks.

Figures (4)

• Figure 1: Illustration of our proposed selfMedHPM, containing a student network and a teacher network, where the teacher is updated by the student in an exponential moving average (EMA) manner.
• Figure 2: Scans in the SMWB are carefully annotated with 41 organs. Each row from left to right is the scanning plane image, coronal plane image, axial plane image, and axial plane labels.
• Figure 3: Reconstruction results of BTCV. First row: Original image. Second row: Masked image where masked regions are colored with black. Third row: Reconstructed images from unmasked patches. Each column shows the slices of different depths.
• Figure 4: Qualitative Results of Segmentation. Results for BTCV are shown in the first two rows. Results for SMWB are shown in the last two rows. In the first and fourth row, note that there is no false positive segmentation (red asterisk). In the second and third rows, the segmentation created by the MAE pre-trained UNETR method (white asterisk) is incomplete compared to the HPM pre-trained UNETR.