PGP-SAM

Code release for "PGP-SAM: Prototype-Guided Prompt Learning for Efficient Few-Shot Medical Image Segmentation"

PGP-SAM is a novel prototype-based few-shot tuning approach which leverages limited samples to replace tedious manual prompts. Our key idea is using inter and intra-class prototypes to capture class-specific knowledge and relationships. We introduce two key components: (1) a plug-and-play contextual modulation module to integrate multi-scale information, (2) a class-guided cross-attention mechanism fusing prototypes and features for automatic prompt generation.

Installation

The code has been tested with python>=3.8 and pytorch==1.12.0. To prepare the conda environment please run the following:

conda create --name pgp-sam python=3.10 -y
conda activate pgp-sam

conda install pytorch==1.12.0 torchvision==0.13.0 torchaudio==0.12.0 cudatoolkit=11.3 -c pytorch
python -m pip install 'git+https://github.com/facebookresearch/detectron2.git'

git clone https://github.com/zzzyzh/PGP-SAM.git
cd PGP-SAM
pip install -r requirements.txt

Data

Download

1. Abdominal CT Synapse Multi-atlas Abdominal Segmentation dataset

Pre-processing

Please refer to Ouyang et al.

Code

Code will be released after paper accepted...