FAMNet

Official code for AAAI 2025 paper: FAMNet: Frequency-aware Matching Network for Cross-domain Few-shot Medical Image Segmentation

• [News!] 24-12-10: Our work is accepted by AAAI25. Arxiv Paper can be found here. 🎉

Abstract

Existing few-shot medical image segmentation (FSMIS) models fail to address a practical issue in medical imaging: the domain shift caused by different imaging techniques, which limits the applicability to current FSMIS tasks. To overcome this limitation, we focus on the cross-domain few-shot medical image segmentation (CD-FSMIS) task, aiming to develop a generalized model capable of adapting to a broader range of medical image segmentation scenarios with limited labeled data from the novel target domain. Inspired by the characteristics of frequency domain similarity across different domains, we propose a Frequency-aware Matching Network (FAMNet), which includes two key components: a Frequency-aware Matching (FAM) module and a Multi-Spectral Fusion (MSF) module. The FAM module tackles two problems during the meta-learning phase: 1) intra-domain variance caused by the inherent support-query bias, due to the different appearances of organs and lesions, and 2) inter-domain variance caused by different medical imaging techniques. Additionally, we design an MSF module to integrate the different frequency features decoupled by the FAM module, and further mitigate the impact of inter-domain variance on the model's segmentation performance. Combining these two modules, our FAMNet surpasses existing FSMIS models and Cross-domain Few-shot Semantic Segmentation models on three cross-domain datasets, achieving state-of-the-art performance in the CD-FSMIS task.

Getting started

Dependencies

Please install the 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

Datesets and Preprocessing

Please download:

1. Abdominal MRI: Combined Healthy Abdominal Organ Segmentation dataset
2. Abdominal CT: Multi-Atlas Abdomen Labeling Challenge
3. Cardiac LGE and b-SSFP: Multi-sequence Cardiac MRI Segmentation dataset
4. Prostate UCLH and NCI: Cross-institution Male Pelvic Structures

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

Training

1. Compile ./data/supervoxels/felzenszwalb_3d_cy.pyx with cython (python ./data/supervoxels/setup.py build_ext --inplace) and run ./data/supervoxels/generate_supervoxels.py
2. Download pre-trained ResNet-50 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_<direction>.sh, for example: ./script/train_ct2mr.sh

Inference

Run ./script/test_<direction>.sh

Acknowledgement

Our code is based on the works: SSL-ALPNet, ADNet and QNet.

Citation

If you use this code for your research or project, please cite:

@article{bo2024famnet,
  title={FAMNet: Frequency-aware Matching Network for Cross-domain Few-shot Medical Image Segmentation},
  author={Bo, Yuntian and Zhu, Yazhou and Li, Lunbo and Zhang, Haofeng},
  journal={arXiv preprint arXiv:2412.09319},
  year={2024}
}