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

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

News

• We release a real-time model trained on BSDS, please see Real-time DiffusionEdge.
• We create a WeChat Group for flexible discussion. Please use WeChat APP to scan the QR code.
• 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 activate 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:

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.

VI. Real-time DiffusionEdge.

1. We now only test in the following environment, and more details will be released soon.

Please follow this link to install TensorRT.

2. Download the pretrained weight. Real-time, qi~dong!

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

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}
}