Awesome
<!-- #region -->Polyp-DDPM: Diffusion-Based Semantic Polyp Synthesis for Enhanced Segmentation
Overview
Welcome to the official PyTorch implementation of our recent study, "Polyp-DDPM: Diffusion-Based Semantic Polyp Synthesis for Enhanced Segmentation". Building on our previous research, "Med-DDPM: Conditional Diffusion Models for Semantic 3D Brain MRI Synthesis", we extend our methodology to the semantic synthesis of 2D colored medical images. We introduce Polyp-DDPM, a novel diffusion-based method specifically designed to generate images of gastrointestinal (GI) tract polyps from segmentation masks. This method seeks to enhance polyp segmentation by addressing the significant challenges of data scarcity, high annotation costs, and privacy issues commonly encountered in medical imaging.
Getting Started
Prerequisites
- Python 3.8.10
- PyTorch 1.13.1+cu116
- NumPy 1.21.0
Training
Our training script is highly flexible, offering a variety of parameters to customize your training process:
-i
,--inputfolder
: Specify the folder containing mask images. Default:dataset/train/masks
.-l
,--targetfolder
: Specify the folder containing target images. Default:dataset/train/images
.--input_size
: Set the input image size. Default:256
.--num_channels
: Define the number of channels for convolutional layers. Default:64
.--num_res_blocks
: Set the number of residual blocks. Default:2
.--batchsize
: Choose the batch size. Default:32
.--epochs
: Specify the number of training epochs. Default:100000
.--timesteps
: Determine the number of timesteps for the diffusion process. Default:250
.-r
,--resume_weight
: Provide a path to resume training from a pre-trained model. Default:""
(indicating start from scratch).
Example Command
Start training with default parameters by running: python3 train.py
Or customize your training setup with your preferred parameters.
Model Weights
Download our pretrained model weights using the link below:
Download Model Weights
Once downloaded, place the file in the models
directory.
Inference
Customize your inference process with these command-line arguments:
-i
,--inputfolder
: Folder path for mask images. Default:dataset/test/masks
.-e
,--exportfolder
: Export folder for generated images. Default:exports
.-w
,--weightfile
: Path to the trained model weights. Default:models/model.pt
.-d
,--device
: Choose the device (cuda
orcpu
). Default:cuda
.--input_size
: Input image size. Default:256
.--num_channels
: Convolutional layer channels. Default:64
.--num_res_blocks
: Number of residual blocks. Default:2
.--timesteps
: Timesteps for diffusion during inference. Default:250
.--batchsize
: Inference batch size. Default:4
.--num_samples
: Samples per mask. Default:4
.
Example Command
For default parameter inference, execute: python3 inference.py To tailor the inference process, adjust the flags accordingly, for example: python3 inference.py -i dataset/custom_test/masks -e custom_exports -w models/custom_model.pt --input_size 256
References
If you find our work useful for your research, please consider citing our papers:
@misc{dorjsembe2024polypddpm,
title={Polyp-DDPM: Diffusion-Based Semantic Polyp Synthesis for Enhanced Segmentation},
author={Zolnamar Dorjsembe and Hsing-Kuo Pao and Furen Xiao},
year={2024},
eprint={2402.04031},
archivePrefix={arXiv},
primaryClass={cs.CV}
}
@ARTICLE{10493074,
author={Dorjsembe, Zolnamar and Pao, Hsing-Kuo and Odonchimed, Sodtavilan and Xiao, Furen},
journal={IEEE Journal of Biomedical and Health Informatics},
title={Conditional Diffusion Models for Semantic 3D Brain MRI Synthesis},
year={2024},
volume={28},
number={7},
pages={4084-4093},
doi={10.1109/JBHI.2024.3385504}
}
Our code is based on our previous work: Med-DDPM. For more details, see the paper https://arxiv.org/abs/2305.18453 and the repository https://github.com/mobaidoctor/med-ddpm.
Acknowledgement:
Our initial source code is based on denoising-diffusion-pytorch and guided-diffusion.
<!-- #endregion -->