This repository contains the official implementation of the paper Highlighted Diffusion Model as Plug-in Priors for Polyp Segmentation, presented at JBHI-2024.

• Requirements
• Dataset Preparation
• Sampling with HDMs
• Training Your Own HDMs
• Downstream Evaluation
• Citations

To get started, set up your environment with the following dependencies:

conda create -n HDM python=3.10
conda activate HDM
conda install pytorch==2.0.1 torchvision==0.15.2 torchaudio==2.0.2 pytorch-cuda=11.7 -c pytorch -c nvidia
pip install pytorch-lightning==1.9.3 omegaconf einops tensorboard albumentations natsort prettytable timm==0.9.5 thop

• Download the 2D dataset from this repository.

• Download highlighted images from this link.

• Organize the dataset as follows:

  ├── ${data_root}
  │   ├── ${train_data_dir}
  │   │   ├── images
  │   │   │   ├── ***.png
  │   │   ├── masks
  │   │   │   ├── ***.png
  │   │   ├── highlighted_GT
  │   │   │   ├── ***.png
  │   ├── ${test_data_dir}
  │   │   ├── images
  │   │   │   ├── ***.png
  │   │   ├── masks
  │   │   │   ├── ***.png
  │   │   ├── highlighted_GT
  │   │   │   ├── ***.png

• Download the dataset from this repository.

• Download highlighted images from this link.

• Organize the dataset as follows:

  ├── ${data_root}
  │   ├── ${train_data_dir}
  │   │   ├── Frame
  │   │   │   ├── case***
  │   │   │   │   ├── ***.jpg
  │   │   ├── GT
  │   │   │   ├── case***
  │   │   │   │   ├── ***.png
  │   │   ├── Highlighted_GT
  │   │   │   ├── case***
  │   │   │   │   ├── ***.png
  │   ├── ${test_data_dir}
  │   │   ├── Frame
  │   │   │   ├── case***
  │   │   │   │   ├── ***.jpg
  │   │   ├── GT
  │   │   │   ├── case***
  │   │   │   │   ├── ***.png
  │   │   ├── Highlighted_GT
  │   │   │   ├── case***
  │   │   │   │   ├── ***.png

• For the SUN-SEG dataset, combine the Unseen test data from both the Easy and Hard cases into a new test set under test_data_dir.

Pre-trained HDM models are available for download:

To run inference using HDMs:

1. Set CONFIG_FILE_PATH and RESUME_PATH in the sample.py file.

2. Run the following command:

   python sample.py

Extracted features are available for download:

To train your own HDM models:

1. Set train_data_dir and test_data_dir in the corresponding configuration file (HDM_xxx.yaml) located in the configs folder.

2. Set CONFIG_FILE_PATH in the main.py file.

3. Run the following command:

   python main.py

For downstream evaluation, detailed instructions can be found in the HDM_Downstream folder. Please refer to this file for further details.

If you find HDM useful for your research, please cite our paper:

@ARTICLE{du2024highlighted,
  author={Du, Yuhao and Jiang, Yuncheng and Tan, Shuangyi and Liu, Si-Qi and Li, Zhen and Li, Guanbin and Wan, Xiang},
  journal={IEEE Journal of Biomedical and Health Informatics}, 
  title={Highlighted Diffusion Model as Plug-in Priors for Polyp Segmentation}, 
  year={2024},
  doi={10.1109/JBHI.2024.3485767}}