Tied Prototype Model for Few-Shot Medical Image Segmentation

Code for our MICCAI 2025 paper "Tied Prototype Model for Few-Shot Medical Image Segmentation" (camera-ready version)

• Abstract
• Methods
• Datasets
• Reproducing the results
• Citation

Common prototype-based medical image few-shot segmentation (FSS) methods model foreground and background classes using class-specific prototypes. However, given the high variability of the background, a more promising direction is to focus solely on foreground modeling, treating the background as an anomaly—an approach introduced by ADNet. Yet, ADNet faces three key limitations: dependence on a single prototype per class, a focus on binary classification, and fixed thresholds that fail to adapt to patient and organ variability. To address these shortcomings, we propose the Tied Prototype Model (TPM), a principled reformulation of ADNet with tied prototype locations for foreground and background distributions. Building on its probabilistic foundation, TPM naturally extends to multiple prototypes and multi-class segmentation while effectively separating non-typical background features. Notably, both extensions lead to improved segmentation accuracy. Finally, we leverage naturally occurring class priors to define an ideal target for adaptive thresholds, boosting segmentation performance. Taken together, TPM provides a fresh perspective on prototype-based FSS for medical image segmentation. The code can be found at https://github.com/hjk92g/TPM-FSS.

(a) SPM

(b) TPM

(c) TPM-MP

(d) TPM-MC

(a) Standard prototype model (SPM):

• Assigns distinct prototypes (centers) for different classes
• Directional separation (class regions: Voronoi cells)
• Not good for modeling the highly heterogeneous background class from only a few prototypes

(b) Tied prototype model (TPM):

• Use a tied prototype location for foreground and background classes
• Inside-outside (IO) classification of feature vectors
• Ideal for modeling the highly heterogeneous background class
• Theorem 1: In a single prototype setting with spherical embedding, TPM is equivalent to ADNet.

(c) TPM-MP:

• Multi-prototype extension of TPM (multiple prototypes per class)
• Increase the representation diversity of the foreground class

(d) TPM-MC:

• Multi-foreground class extension of TPM (segment multiple foreground classes simultaneously)
• Enables training to separate different foreground features into their respective prototype positions while pushing background features away from the prototypes

Ideal class priors (ICPs):

• Address class imbalance in medical image segmentation
• ICP makes: Number of predicted foreground pixel counts = Number of true foreground pixel counts
• To estimate ICPs of unseen images, we consider two estimation methods from training data: average estimation (AvgEst) and linear estimation (LinEst). LinEst uses a simple linear model that takes only 2d information: support foreground size and query slice location.

ABD-MRI dataset: CHAOST2

ABD-CT dataset: SABS (Download Abdomen-RawData.zip)

Requirements can be found in the Dockerfile.

We follow the preprocessing steps described by Ouyang et al. (you can ignore their "pseudolabel generation" and following steps). For this, you need to use their code.

We follow the supervoxel generation step by Hansen et al. (you can ignore their "downloading pretrained 3D ResNeXt" and following steps). For this, you need to use their code. Specifically, run ./supervoxels/generate_supervoxels.py with some modification. (Uncomment lines 90-91. Adjust base_dir and seg_fid, accordingly. Set n_sv parameter to 5000 for CHAOST2 and 2000 for SABS.)

Before running codes, put your Weights & Biases key into wandb.login(key='**********') in main_...py files. (The same applies to analysis and inference codes.)

Run ./scripts/train_<abd, sabs>_2D.sh for binary segmentation (single-foreground class) training. (abd: ABD-MRI dataset. sabs: ABD-CT dataset.)

Run ./scripts/train_multiclass_<abd, sabs>_2D.sh for multi-class segmentation training.

Run ./scripts/analyze_<abd, sabs>_2D.sh to estimate (to train) ICPs for binary segmentation from training data.

Run ./scripts/analyze_multiclass_<abd, sabs>_2D.sh to estimate (to train) ICPs for multi-class segmentation from training data.

(You need to run the analyze step before running this step.)

Run ./scripts/test_<abd, sabs>_2D.sh for binary segmentation evaluation.

Run ./scripts/test_multiclass_<abd, sabs>_2D.sh for multi-class segmentation evaluation.

Note that this repo is based on ADNet repository and our implemented ADNet++ is different from its original paper. We didn't implement "3.2.2. Uncertainty estimation" and "3.2.3. Supervoxel-informed feature refinement module".

If you use our research, we kindly ask you to cite our paper:

@inproceedings{kim2025tied,
        title = {Tied Prototype Model for Few-Shot Medical Image Segmentation},
        author = {Kim, Hyeongji and Hansen, Stine and Kampffmeyer, Michael},
        booktitle = {International Conference on Medical Image Computing and Computer-Assisted Intervention},
        year = {2025},
        publisher = {Springer Nature Switzerland},
        volume = {LNCS 15962},
        month = {September},
        pages = {651 -- 661}
}