Generative AI Enables Medical Image Segmentation in Ultra Low-Data Regimes

Requirements

Python 3.7+ and Pytorch 1.12.1+ and CUDA 11.3 are recommended. Docker can be used with the given Dockerfile to quickly setup the enviornment or a local conda env can be create using the following:

conda create -n semantic python=3.8
conda activate semantic
pip install -r requirements.txt

Datasets

We use JSRT as the in-domain dataset to train and evaluate the model. Further, we use NLM(MC) and NLM(SZ), as the out-of-domain datasets for model evaluation. (For some image examples, their lung segmentation masks are divided into the right and the left lung mask. For these, the masks need to be combined first.)

Download the data and place it in the data folder. Dataset tree structure example:

data/JSRT
├── Images
│   ├── JPCLN001.png
│   ├── JPCLN002.png
│   ├── ...
├── Masks
│   ├── JPCLN001.gif
│   ├── JPCLN002.gif
│   ├── ...
project code
├── ...

Training and Testing

We pre-train the GAN-based augmentation model on the train and val sets of the in-domain dataset followed by training both augmentation and semantic segmentation models end-to-end on the in-domain dataset. Finally, we test the trained models on the out-of-domain datasets. The results on the test set of both in-domain and out-of-domain datasets are shown using wandb during training. The training process needs about 1.5 hours on the NVIDIA A100 device with 40G memory.

To train the models from scratch, use the following command (Related configurations of model path should be changed mutually):

# Pre-train the augmentation model
bash scripts/train_pix2pix_jsrt.sh

# Train the segmentation based on our framework
bash scripts/train_end2end_jsrt.sh

# Inference the trained segmentation model
bash scripts/test_lung.sh

Pre-trained model

Models pre-trained on the JSRT dataset (trained with 9 labeled data examples) are available through the following links: Pix2Pix-generator | Pix2Pix-discriminator | U-Net

Citation

If you find this project useful in your research, please consider citing:

@article{zhang2024generative,
  title={Generative AI Enables Medical Image Segmentation in Ultra Low-Data Regimes},
  author={Zhang, Li and Jindal, Basu and Alaa, Ahmed and Weinreb, Robert and Wilson, David and Segal, Eran and Zou, James and Xie, Pengtao},
  journal={medRxiv},
  pages={2024--08},
  year={2024},
  publisher={Cold Spring Harbor Laboratory Press}
}

Code Dependencies

Our code is based on the following repositories: Pix2Pix model | U-Net | Betty framework

License

GenSeg is licensed under the Apache 2.0 License.