Code for paper: Learning De-biased Prototypes for Few-Shot Medical Image Segmentation

Prototypical networks have emerged as the dominant method for Few-shot Medical image Segmentation (FSMIS). Despite their success, the commonly used Masked Average Pooling (MAP) approach in prototypical networks computes the mean of the masks, resulting in imprecise and inadequate prototypes that fail to capture the subtle nuances and variations in the data. To address this issue, we propose a simple yet effective module called De-biasing Masked Average Pooling (DMAP) to generate more accurate prototypes from filtered foreground support features. Specifically, our approach introduces a Learnable Threshold Generation (LTG) module that adaptively learns a threshold based on the extracted features from both support and query images, and then choose partial foreground pixels that have larger similarity than the threshold to generate prototypes. Our proposed method is evaluated on three popular medical image segmentation datasets, and the results demonstrate the superiority of our approach over some state-of-the-art methods.

Please install following essential dependencies:

dcm2nii
json5==0.8.5
jupyter==1.0.0
nibabel==2.5.1
numpy==1.22.0
opencv-python==4.5.5.62
Pillow>=8.1.1
sacred==0.8.2
scikit-image==0.18.3
SimpleITK==1.2.3
torch==1.10.2
torchvision=0.11.2
tqdm==4.62.3

Download:

1. Combined Healthy Abdominal Organ Segmentation data set
2. Multi-sequence Cardiac MRI Segmentation data set (bSSFP fold)
3. Multi-Atlas Abdomen Labeling Challenge

Pre-processing is performed according to Ouyang et al. and we follow the procedure on their github repository.

1. Compile ./supervoxels/felzenszwalb_3d_cy.pyx with cython (python ./supervoxels/setup.py build_ext --inplace) and run ./supervoxels/generate_supervoxels.py
2. Download pre-trained ResNet-101 weights vanilla version or deeplabv3 version and put your checkpoints folder, then replace the absolute path in the code ./models/encoder.py.
3. Run ./script/train.sh

Run ./script/test.sh

@article{zhu2024learning,
  title={Learning De-biased prototypes for few-shot medical image segmentation},
  author={Zhu, Yazhou and Cheng, Ziming and Wang, Shidong and Zhang, Haofeng},
  journal={Pattern Recognition Letters},
  year={2024},
  publisher={Elsevier}
}