DiffWater: Underwater Image Enhancement Based on Conditional Denoising Diffusion Probabilistic Model

[paper]

This Repo includes the training and testing codes of our DiffWater. (Pytorch Version)
If you use our code, please cite our paper and hit the star at the top-right corner. Thanks!

@ARTICLE{10365517,
  author={Guan, Meisheng and Xu, Haiyong and Jiang, Gangyi and Yu, Mei and Chen, Yeyao and Luo, Ting and Zhang, Xuebo},
  journal={IEEE Journal of Selected Topics in Applied Earth Observations and Remote Sensing}, 
  title={DiffWater: Underwater Image Enhancement Based on Conditional Denoising Diffusion Probabilistic Model}, 
  year={2024},
  volume={17},
  number={},
  pages={2319-2335},
  keywords={Image color analysis;Colored noise;Noise reduction;Image restoration;Image enhancement;Visualization;Lighting;Color compensation;conditional denoising diffusion probabilistic model (DDPM);underwater image enhancement (UIE)},
  doi={10.1109/JSTARS.2023.3344453}}

Environment

pip install -r requirement.txt

Train and Test sets:

To make use of the train.py and test.py the dataset folder names should be lower-case and structured as follows.

└──── <data directory>/
    ├──── UIEB_R90/
    |   ├──── input_256/
    |   |   ├──── 01.png/
    |   |   ├──── ...
    |   |   └──── 90.png/
    |   ├──── target_256/
		├──── 01.png/
		├──── ...
		└──── 90.png/

Dataset

(1) LSUI : Data

(2) UIEB : Data

(3) SQUID : Data

(4) U45 : Data

The test results of the paper

[U90,C60,U45,S16,L504]Google and BaiduYun) password:gms1

Evaluation

To resume from a checkpoint file, simply use the --resume argument in test.py to specify the checkpoint.

For your convenience, we provide the pre-trained model in our paper. BaiduYun password:gms1 and Google

Training

To resume from a checkpoint file, simply use the --resume argument in train.py to specify the checkpoint.

Acknowledgement

Our code is adapted from the original SR3 repository. We thank the authors for sharing their code.

Contact

If you have any questions, please contact Meisheng Guan at 1971306417@qq.com.