The Pytorch Implementation of ''Adaptive Uncertainty Distribution in Deep Learning for Unsupervised Underwater Image Enhancement''.

In this project, we use Linux host with a single NVidia GeForce RTX 2080 Ti GPU with 11 GB of memory, Python 3.7, Pytorch 1.9.1.

Download the pretrained model pretrained model.

Check the model and image pathes in Inference.py, and then run:

python Inference.py  

To train the model, you need to download Underwater Dataset dataset.

Check the dataset path in Train.py, and then run:

python Train.py   

If you find UDnet is useful in your research, please cite our paper:

@article{Saleh2022udnet,
    title = {{Adaptive Uncertainty Distribution in Deep Learning for Unsupervised Underwater Image Enhancement}},
    year = {2022},
    author = {Saleh, Alzayat and Sheaves, Marcus and Jerry, Dean and Azghadi, Mostafa Rahimi},
    month = {12},
    url = {https://arxiv.org/abs/2212.08983v1},
    doi = {10.48550/arxiv.2212.08983},
    arxivId = {2212.08983},
    keywords = {Convolutional Neural Net-works, Deep Learning, Index Terms-Computer Vision, Machine Learning, Underwater Image Enhancement, Variational Autoen-coder}
}

The code is made available for academic research purpose only. This project is open sourced under MIT license.

https://github.com/zhenqifu/PUIE-Net
https://github.com/deepxzy/USLN
https://github.com/xueleichen/PyTorch-Underwater-Image-Enhancement