UIEDP

This repository provides offical Pytorch implementation of the paper "UIEDP: Boosting underwater image enhancement with diffusion prior".

Environment

conda create -n uiedp python=3.9
conda activate uiedp

pip install torch==1.13.0+cu117 torchvision==0.14.0+cu117 torchaudio==0.13.0 -f https://download.pytorch.org/whl/torch_stable.html

pip install tb_nightly==2.14.0a20230808 -i https://mirrors.aliyun.com/pypi/simple

pip install pyiqa

pip install blobfile

pip install mpi4py

Pretrained Diffusion Model

You can download the pretrained diffusion model provided by OpenAI. Then put it in the ./models/ folder.

Prepare Data

You need to download the raw images and pseudo-label images of three test datasets (T90|C60|U45) from Baidu Netdisk, and the passward is 'data'. Then unzip the data.zip archive in the ./data/ folder:

unzip ./data/data.zip -d ./data/

Then you can obtain the following folder structure:

./data/
├── pseudo_label
│   ├── C60/
│   ├── T90/
│   └── U45/
├── raw
│   ├── C60/
│   ├── T90/
│   └── U45/
└── reference
    └── T90/

• The subfolder pseudo_label/ contains the pseudo-label images of T90|C60|U45 datasets generated by pretrained UIEC2Net. Here we utilize UIEC2Net as the UIE algorithm which generates pseudo-label images in UIEDP. If you want to use any other UIE algorithm, please put its enhanced results in the pseudo_label/ folder, just like above.

• The subfolder raw/ contains the raw images of T90|C60|U45 datasets.

• The subfolder reference/ contains the reference images of T90 dataset, because C60 and U45 have no reference images.

Conditional Generation

You can directly run the script:

sh run_ddpm.sh

the script contains:

SAMPLE_FLAGS="--batch_size 8"
MODEL_FLAGS="--attention_resolutions 32,16,8 --class_cond False --diffusion_steps 1000 --image_size 256 --learn_sigma True --noise_schedule linear --num_channels 256 --num_head_channels 64 --num_res_blocks 2 --resblock_updown True --use_fp16 True --use_scale_shift_norm True"
CUDA_VISIBLE_DEVICES=0 python uie_sample.py $MODEL_FLAGS --uie_dataset U45 --classifier_scale 4000.0 --model_path models/256x256_diffusion_uncond.pt $SAMPLE_FLAGS

You can change the dataset into one of T90|C60|U45, and change the classifier_scale (guidance scale). The generated enhanced images are saved in the ./data/UIEDP/T90|C60|U45/ folder.

Evaluation

After sampling, you can evaluate perfermance of any dataset:

python uie_test.py --dataset T90|C60|U45

Acknowledgment

We adapted the code of guided-diffusion and GDP. Thanks to the original authors for their work!

Citation

@article{du2025uiedp,
  title={UIEDP: Boosting underwater image enhancement with diffusion prior},
  author={Du, Dazhao and Li, Enhan and Si, Lingyu and Zhai, Wenlong and Xu, Fanjiang and Niu, Jianwei and Sun, Fuchun},
  journal={Expert Systems with Applications},
  volume={259},
  pages={125271},
  year={2025},
  publisher={Elsevier}
}