This is the official repository of the paper Annotator Consensus Prediction for Medical Image Segmentation with Diffusion Models

The code is based on Improved Denoising Diffusion Probabilistic Models and SegDiff: Image Segmentation with Diffusion Probabilistic Models.

Installation

Conda environment

To create the environment use the conda environment command

conda env create -f environment.yml

Project structure and data preparations

Our project need to be arranged in the following format

Annotator-Consensus-Prediction/ # git clone the source code here

data/ # the root of the data folders
    brain_growth/
    brain_tumor/
    kidney/
    prostate/

All datasets were taken from Quantification of Uncertainties in Biomedical Image Quantification Challenge (QUBIQ) 2021. Download the datasets from the following link.

The datasets should have the following format

<dataset_name>/
    train/*
    val/*

Train and Evaluate

Execute the following commands (multi gpu is supported for training, set the gpus with CUDA_VISIBLE_DEVICES and -n for the actual number)

Most important training options:

# Training
--batch-size    Batch size
--lr            Learning rate

# Architecture
--rrdb_blocks       Number of rrdb blocks
--dropout           Dropout
--diffusion_steps   Number of steps for the diffusion model

# Ablation variants (described in the paper)
--no_annotator_training
--annotators_training
--soft_label_training
--consensus_training (ours)

# Misc
--save_interval     interval for saving model weights

Training script example (different script for each dataset):

CUDA_VISIBLE_DEVICES=0,1,2,3 mpiexec -n 4 python image_train_diff_brain_tumor_1.py --save_interval 2500 --batch_size 1 --lr 0.00002 --diffusion_steps 100 --no_annotator_training True

Evaluation script example (different script for each dataset):

CUDA_VISIBLE_DEVICES=0,1 mpiexec -n 2 python image_sample_diff_brain_tumor_1.py --model_path <path-for-model-weights>