https://github.com/ddz16/uiedp
This repository provides offical Pytorch implementation of the paper "UIEDP: Boosting underwater image enhancement with diffusion prior".
https://github.com/ddz16/uiedp
diffusion-models image-processing underwater-image-enhancement underwater-image-restoration underwater-images
Last synced: about 1 year ago
JSON representation
This repository provides offical Pytorch implementation of the paper "UIEDP: Boosting underwater image enhancement with diffusion prior".
- Host: GitHub
- URL: https://github.com/ddz16/uiedp
- Owner: ddz16
- Created: 2023-11-28T08:18:56.000Z (over 2 years ago)
- Default Branch: main
- Last Pushed: 2025-02-12T01:33:13.000Z (over 1 year ago)
- Last Synced: 2025-02-12T02:37:13.667Z (over 1 year ago)
- Topics: diffusion-models, image-processing, underwater-image-enhancement, underwater-image-restoration, underwater-images
- Language: Python
- Homepage: https://www.sciencedirect.com/science/article/abs/pii/S0957417424021389
- Size: 20.7 MB
- Stars: 0
- Watchers: 1
- Forks: 0
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
Awesome Lists containing this project
README
# UIEDP
This repository provides offical Pytorch implementation of the paper "[UIEDP: Boosting underwater image enhancement with diffusion prior](https://www.sciencedirect.com/science/article/abs/pii/S0957417424021389)".
## 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](https://openaipublic.blob.core.windows.net/diffusion/jul-2021/256x256_diffusion_uncond.pt) 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](https://pan.baidu.com/s/1nZI734Hauh6ze3J_JgCY3Q?pwd=data), 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](https://github.com/ddz16/UIE_Benckmark). 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](https://github.com/openai/guided-diffusion/tree/main) and [GDP](https://github.com/Fayeben/GenerativeDiffusionPrior). 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}
}
```