DiffusionEdge: Diffusion Probabilistic Model for Crisp Edge Detection (arxiv)

Yunfan Ye, Yuhang Huang, Renjiao Yi, Zhiping Cai, Kai Xu.

News

• Update demo.py for faster inference.
• 2023-12-09: The paper is accepted by AAAI-2024.
• Upload the pretrained first stage checkpoint download.
• Upload pretrained weights and pre-computed results.
• We now update a simple demo, please see Quickly Demo
• First Committed.

I. Before Starting.

1. install torch

conda create -n diffedge python=3.9
conda avtivate diffedge
pip install torch==1.12.1+cu113 torchvision==0.13.1+cu113 torchaudio==0.12.1 --extra-index-url https://download.pytorch.org/whl/cu113

2. install other packages.

pip install -r requirement.txt

3. prepare accelerate config.

accelerate config

II. Prepare Data.

The training data structure should look like:

|-- $data_root
|   |-- image
|   |-- |-- raw
|   |-- |-- |-- XXXXX.jpg
|   |-- |-- |-- XXXXX.jpg
|   |-- edge
|   |-- |-- raw
|   |-- |-- |-- XXXXX.png
|   |-- |-- |-- XXXXX.png

The testing data structure should look like:

|-- $data_root
|   |-- XXXXX.jpg
|   |-- XXXXX.jpg

III. Quickly Demo !

1. download the pretrained weights:

Dataset	ODS (SEval/CEval)	OIS (SEval/CEval)	AC	Weight	Pre-computed results
BSDS	0.834 / 0.749	0.848 / 0.754	0.476	download	download
NYUD	0.761 / 0.732	0.766 / 0.738	0.846	download	download
BIPED	0.899	0.901	0.849	download	download

2. put your images in a directory and run:

python demo.py --input_dir $your input dir$ --pre_weight $the downloaded weight path$ --out_dir $the path saves your results$ --bs 8

The larger --bs is, the faster the inference speed is and the larger the CUDA memory is.

IV. Training.

1. train the first stage model (AutoEncoder):

accelerate launch train_vae.py --cfg ./configs/first_stage_d4.yaml

2. you should add the final model weight of the first stage to the config file ./configs/BSDS_train.yaml (line 42), then train latent diffusion-edge model:

accelerate launch train_cond_ldm.py --cfg ./configs/BSDS_train.yaml

V. Inference.

make sure your model weight path is added in the config file ./configs/BSDS_sample.yaml (line 73), and run:

python sample_cond_ldm.py --cfg ./configs/BSDS_sample.yaml

Note that you can modify the sampling_timesteps (line 11) to control the inference speed.

Contact

If you have some questions, please contact with huangai@nudt.edu.cn.

Thanks

Thanks to the base code DDM-Public.

Citation

@inproceedings{ye2024diffusionedge,
      title={DiffusionEdge: Diffusion Probabilistic Model for Crisp Edge Detection}, 
      author={Yunfan Ye and Kai Xu and Yuhang Huang and Renjiao Yi and Zhiping Cai},
      year={2024},
      booktitle={AAAI}
}