AnatoMask: Enhancing Medical Image Segmentation with Reconstruction-guided Self-masking

TL;DR

AnatoMask introduces a self-supervised learning method for 3D medical image segmentation that dynamically masks important anatomical regions to enhance pretraining efficiency and performance across multiple imaging modalities.

Contribution

It proposes a novel reconstruction-guided self-masking strategy that adaptively identifies and masks significant regions, improving SSL effectiveness in medical imaging.

Findings

Outperforms existing SSL methods on multiple datasets

Effective across CT, MRI, and PET modalities

Demonstrates scalability and superior segmentation accuracy

Abstract

Due to the scarcity of labeled data, self-supervised learning (SSL) has gained much attention in 3D medical image segmentation, by extracting semantic representations from unlabeled data. Among SSL strategies, Masked image modeling (MIM) has shown effectiveness by reconstructing randomly masked images to learn detailed representations. However, conventional MIM methods require extensive training data to achieve good performance, which still poses a challenge for medical imaging. Since random masking uniformly samples all regions within medical images, it may overlook crucial anatomical regions and thus degrade the pretraining efficiency. We propose AnatoMask, a novel MIM method that leverages reconstruction loss to dynamically identify and mask out anatomically significant regions to improve pretraining efficacy. AnatoMask takes a self-distillation approach, where the model learns both how to find more significant regions to mask and how to reconstruct these masked regions. To avoid suboptimal learning, Anatomask adjusts the pretraining difficulty progressively using a masking dynamics function. We have evaluated our method on 4 public datasets with multiple imaging modalities (CT, MRI, and PET). AnatoMask demonstrates superior performance and scalability compared to existing SSL methods. The code is available at https://github.com/ricklisz/AnatoMask.