Installation

To install simply run

conda env create -f environment.yml

Training and Sampling

Training information and additional README can be found in each model subfolder.

1. Diffusion model for generating mask(s)

To change directory:

export OPENAI_LOGDIR={OUTPUT_FOLDER}

Model parameters:

MODEL_FLAGS="--image_size 256 --num_channels 128 --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16"
DIFFUSION_FLAGS="--diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False"
TRAIN_FLAGS="--lr 1e-4 --batch_size 2"

Training:

python improved-diffusion/scripts/image_train.py --data_dir ./improved-diffusion/datasets/Kvasir-SEG/masks $MODEL_FLAGS $DIFFUSION_FLAGS $TRAIN_FLAGS

Sampling:

python improved-diffusion/scripts/image_sample.py --num_samples {SAMPLES} --model_path {MODEL_CHECKPOINT.pt} --output {OUTPUT_TYPE} --postprocess {POSTPROCESS} $MODEL_FLAGS $DIFFUSION_FLAGS

2. Conditional diffusion model for generating polyp images

Training:

python latent-diffusion/main.py --base latent-diffusion/configs/latent-diffusion/kvasir-ldm-vq4-.yaml -t --gpus 0,

Sampling: We can generate polyp images based on the condition, based on:

• Existing dataset

python latent-diffusion/scripts/inference_dataset.py

• Generated mask, mask(s) needs to be placed inside latent-diffusion/data/samples/masks:

python latent-diffusion/scripts/inference_mask.py {IMAGE_NAME} --samples {SAMPLES}

Results are stored inside latent-diffusion/results

Pipeline

To render multiple polyps we can use {SAMPLES} to sample multiple mask(s) and use them to generate polyp(s):

export OPENAI_LOGDIR='latent-diffusion/results/masks/'

MODEL_FLAGS="--image_size 256 --num_channels 128 --num_res_blocks 2 --num_heads 1 --learn_sigma True --use_scale_shift_norm False --attention_resolutions 16"
DIFFUSION_FLAGS="--diffusion_steps 1000 --noise_schedule linear --rescale_learned_sigmas False --rescale_timesteps False"
TRAIN_FLAGS="--lr 1e-4 --batch_size 2"

python improved-diffusion/scripts/image_sample.py --num_samples {SAMPLES} --model_path {MODEL_CHECKPOINT.pt} --output png --postprocess {POSTPROCESS} $MODEL_FLAGS $DIFFUSION_FLAGS
python latent-diffusion/scripts/inference_pipe.py

Utils

a) Mask comparator

To visually inspect overlap between generated mask(s) and the training dataset that was used for training the diffusion model, we can use:

python improved-diffusion/scripts/image_compare.py {KVASIR_PATH} {MASK_IMAGE_PATH}

b) Segmentation model

Follow the instruction in ./segmentation_experiments

Contact details:

[email protected] or [email protected]

Citation:

@inproceedings{10.1145/3592571.3592978,
author = {Mach\'{a}\v{c}ek, Roman and Mozaffari, Leila and Sepasdar, Zahra and Parasa, Sravanthi and Halvorsen, P\r{a}l and Riegler, Michael A. and Thambawita, Vajira},
title = {Mask-conditioned latent diffusion for generating gastrointestinal polyp images},
year = {2023},
isbn = {9798400701863},
publisher = {Association for Computing Machinery},
address = {New York, NY, USA},
url = {https://doi.org/10.1145/3592571.3592978},
doi = {10.1145/3592571.3592978},
booktitle = {Proceedings of the 4th ACM Workshop on Intelligent Cross-Data Analysis and Retrieval},
pages = {1–9},
numpages = {9},
keywords = {polyp segmentation, polyp generative model, generating synthetic data, diffusion model},
location = {Thessaloniki, Greece},
series = {ICDAR '23}
}