Yunfan Ye, Yuhang Huang, Renjiao Yi, [Zhiping Cai](), Kai Xu.
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
pip install -r requirement.txt
accelerate config
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
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 |
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.
accelerate launch train_vae.py --cfg ./configs/first_stage_d4.yaml
./configs/BSDS_train.yaml
(line 42), then train latent diffusion-edge model:
accelerate launch train_cond_ldm.py --cfg ./configs/BSDS_train.yaml
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.
Environment | Version |
---|---|
TensorRT | 8.6.1 |
cuda | 11.6 |
cudnn | 8.7.0 |
pycuda | 2024.1 |
Please follow this link to install TensorRT.
python demo_trt.py --input_dir $your input dir$ --pre_weight $the downloaded weight path$ --out_dir $the path saves your results$
If you have some questions, please contact with huangai@nudt.edu.cn.
Thanks to the base code DDM-Public.
@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}
}