Stable Diffusion Segmentation (SDSeg)

The official implementation of Stable Diffusion Segmentation for Biomedical Images with Single-step Reverse Process at MICCAI 2024.

📣 News

⚠️⚠️⚠️ WARNING: for previous users, please set increase_log_steps: False in the *.yaml setting files, this will reduce meaningless logging process and increase training speed!!!

• 10/07: The final published version of the paper is available! See:

• 09/29: The MICCAI poster of SDSeg is released: See you in Marrakesh!

• 07/14: We release a for you to understand our work better. Check it out!

• 06/27: The paper of SDSeg has been pre-released on

• 06/17: 🎉🥳 SDSeg has been accepted by MICCAI2024! Our paper will be available soon.

📌 SDSeg Framework

SDSeg is built on Stable Diffusion (V1), with a downsampling-factor 8 autoencoder, a denoising UNet, and trainable vision encoder (with the same architecture of the encoder in the f=8 autoencoder).

⚙️ Requirements

A suitable conda environment named sdseg can be created and activated with:

conda env create -f environment.yaml
conda activate sdseg

Then, install some dependencies by:

pip install -e git+https://github.com/CompVis/taming-transformers.git@master#egg=taming-transformers
pip install -e git+https://github.com/openai/CLIP.git@main#egg=clip
pip install -e .

Solve GitHub connection issues when downloading taming-transformers or clip

After creating and entering the sdseg environment:

1. create an src folder and enter:

mkdir src
cd src

2. download the following codebases in *.zip files and upload to src/:
   • https://github.com/CompVis/taming-transformers, taming-transformers-master.zip
   • https://github.com/openai/CLIP, CLIP-main.zip
3. unzip and install taming-transformers:

unzip taming-transformers-master.zip
cd taming-transformers-master
pip install -e .
cd ..

4. unzip and install clip:

unzip CLIP-main.zip
cd CLIP-main
pip install -e .
cd ..

5. install latent-diffusion:

cd ..
pip install -e .

Then you're good to go!

🩻 Dataset Settings

Note

The image data should be place at ./data/, while the dataloaders are at ./ldm/data/

We evaluate SDSeg on the following medical image datasets:

Dataset	URL	Preprocess
BTCV	This URL, download the Abdomen/RawData.zip.	Use the code in ./data/synapse/nii2format.py
STS-3D	This URL, download the labelled.zip.	Use the code in ./data/sts3d/sts3d_preprocess.py
REFUGE2	This URL	Following this repo
CVC-ClinicDB	This URL	None
Kvasir-SEG	This URL	None

📦 Model Weights

Pretrained Models

SDSeg uses pre-trained weights from SD to initialize before training.

For pre-trained weights of the autoencoder and conditioning model, run

bash scripts/download_first_stages_f8.sh

For pre-trained wights of the denoising UNet, run

bash scripts/download_models_lsun_churches.sh

Trained SDSeg Models

The model weights trained on medical image datasets will be available soon.

📄 Scripts

Training Scripts

Take CVC dataset as an example, run

nohup python -u main.py --base configs/latent-diffusion/cvc-ldm-kl-8.yaml -t --gpus 0, --name experiment_name > nohup/experiment_name.log 2>&1 &

You can check the training log by

tail -f nohup/experiment_name.log

Also, tensorboard will be on automatically. You can start a tensorboard session with --logdir=./logs/

Note

If you want to use parallel training, the code trainer_config["accelerator"] = "gpu" in main.py should be changed to trainer_config["accelerator"] = "ddp". However, parallel training is not recommended.

Warning

A single SDSeg model ckeckpoint is around 5GB. By default, save only the last model and the model with the highest dice score. If you have tons of storage space, feel free to save more models by increasing the save_top_k parameter in main.py.

Testing Scripts

After training an SDSeg model, you should manually modify the run paths in scripts/slice2seg.py, and begin an inference process like

python -u scripts/slice2seg.py --dataset cvc

Stability Evaluaition

To conduct an stability evaluation process mentioned in the paper, you can start the test by

python -u scripts/slice2seg.py --dataset cvc --times 10 --save_results

This will save 10 times of inference results in ./outputs/ folder. To run the stability evaluation, open scripts/stability_evaluation.ipynb, and modify the path for the segmentation results. Then, click Run All and enjoy.

‼️ Important Files and Folders to Focus on

Training related:

• SDSeg model: ./ldm/models/diffusion/ddpm.py in the class LatentDiffusion.
• Experiment Configurations: ./configs/latent-diffusion

Inference related:

• Inference starting scripts: ./scripts/slice2seg.py,
• Inference implementation: ./ldm/models/diffusion/ddpm.py, under the log_dice method of LatentDiffusion.

Dataset related:

• Dataset storation: ./data/
• Dataloader files: ./ldm/data/

📝 Citation

If you find our work useful, please cite:

@InProceedings{lin2024stable,
author="Lin, Tianyu
  and Chen, Zhiguang
  and Yan, Zhonghao
  and Yu, Weijiang
  and Zheng, Fudan",
title="Stable Diffusion Segmentation for Biomedical Images with Single-Step Reverse Process",
booktitle="Medical Image Computing and Computer Assisted Intervention -- MICCAI 2024",
year="2024",
publisher="Springer Nature Switzerland",
address="Cham",
pages="656--666",
isbn="978-3-031-72111-3"
}

🔜 TODO List

• Organizing the inference code. (Toooo redundant right now.)
• Reimplement SDSeg in OOP. (Elegance is the key!)
• Add README for multi-class segmentation.
• Release model weights.
• Reimplement using diffusers.
• Reduce model checkpoint size (no need to save autoencoder's weights).