zhttp / PyTorch-Underwater-Image-Enhancement

The repo for "Underwater Image Enhancement based on Deep Learning and Image Formation Model"

PyTorch-Underwater-Image-Enhancement

This is the repo for "Underwater Image Enhancement based on Deep Learning and Image Formation Model"[arXiv]

The current code works with NVIDIA GPU on Ubuntu. You can do testing on CPU.

Requirements

pip install -r requirements.txt

Train the model

$ python train.py TRAIN_RAW_IMAGE_FOLDER TRAIN_REFERENCE_IMAGE_FOLDER

Test the model

$ python test.py --checkpoint CHECKPOINTS_PATH

For convenience, you can run the following command to quickly see the results using the trained model reported in our paper.

$ python test.py --checkpoint ./checkpoints/model_best_2842.pth.tar

Citation

If you use this code in your research, please consider citing the following paper:

@misc{chen2021underwater,
  title={Underwater Image Enhancement based on Deep Learning and Image Formation Model}, 
  author={Xuelei Chen and Pin Zhang and Lingwei Quan and Chao Yi and Cunyue Lu},
  year={2021},
  eprint={2101.00991},
  archivePrefix={arXiv},
  primaryClass={eess.IV}

or the Chinese journal version:

[1]陈学磊,张品,权令伟,易超,鹿存跃.融合深度学习与成像模型的水下图像增强算法[J/OL].计算机工程:1-9[2021-03-15].https://doi.org/10.19678/j.issn.1000-3428.0060653.

Acknowledgment

Sponsored by the Oceanic Interdisciplinary Program of Shanghai Jiao Tong University (project number SL2020ZD103)