Stable Diffusion Segmentation for Biomedical Images with Single-step Reverse Process

MICCAI 2024 | |

By | Zhiguang Chen | | | Fudan Zheng

• 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 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).

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 .

mkdir src
cd src

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

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

cd ..
pip install -e .

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

We evaluate SDSeg on the following medical image datasets:

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

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

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/

STORAGE 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.

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

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.

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/

If you find our work useful, please cite:

@article{lin2024stable,
  title={Stable Diffusion Segmentation for Biomedical Images with Single-step Reverse Process},
  author={Lin, Tianyu and Chen, Zhiguang and Yan, Zhonghao and Zheng, Fudan and Yu, Weijiang},
  journal={arXiv preprint arXiv:2406.18361},
  year={2024}
}

@inproceedings{lin2024stable,
  title={Stable Diffusion Segmentation for Biomedical Images with Single-step Reverse Process},
  author={Lin, Tianyu and Chen, Zhiguang and Yan, Zhonghao and Yu, Weijiang and Zheng, Fudan},
  booktitle={International Conference on Medical Image Computing and Computer-Assisted Intervention},
  year={2024},
  organization={Springer}
}

• 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.