This repository provides implementation of some underwater image enhancement methods and datasets, including:
Supported Methods
| Method | Pub | Language | Paper | Reference Codes |
|---|---|---|---|---|
| Fusion | 2012 CVPR | MATLAB | Enhancing Underwater Images and Videos by Fusion | Fusion |
| UWCNN | 2020 PR | Pytorch | Underwater Scene Prior Inspired Deep Underwater Image and Video Enhancement | UWCNN |
| UIEC2Net | 2021 SPIC | Pytorch | UIEC^2-Net: CNN-based Underwater Image Enhancement Using Two Color Space | UIEC2Net |
| MLLE | 2022 TIP | MATLAB | Underwater Image Enhancement via Minimal Color Loss and Locally Adaptive Contrast Enhancement | MLLE |
| UTrans | 2023 TIP | Pytorch | U-Shape Transformer for Underwater Image Enhancement | UTrans |
| NU2Net | 2023 AAAI | Pytorch | Underwater Ranker: Learn Which Is Better and How to Be Better | NU2Net |
| FiveAPlus | 2023 BMVC | Pytorch | Five A+ Network: You Only Need 9K Parameters for Underwater Image Enhancement | FiveAPlus |
Supported Datasets
conda create -n uie python=3.9
conda activate uie
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 pyiqa
pip install pytorch_lightning==2.0.9.post0
Download the UIEB dataset in the ./data/UIEB/ folder, then you have three subfolders raw-890/, reference-890/, challenging-60/, just like:
./data/UIEB/
├── challenging-60/
├── challenging.txt
├── raw-890/
├── reference-890/
├── test.txt
└── train.txt
Download the LSUI dataset in the ./data/LSUI/ folder, then you have two subfolders GT/, input/, just like:
./data/LSUI/
├── GT/
├── input/
├── test.txt
└── train.txt
Train on the training set of UIEB dataset:
python train_UIEB.py --model_name UIEC2Net --batch_size 16 --epochs 100
We provide four models' checkpoints trained on UIEB dataset in the ./checkpoints/UIEB/ folder:
./checkpoints/UIEB/
├── FIVE_APLUS.ckpt
├── NU2Net.ckpt
├── UIEC2Net.ckpt
└── UWCNN.ckpt
After training, you can enhance the images in the test set and challenging set of UIEB dataset:
python test_UIEB.py --model_name UIEC2Net
The generated enhanced images are saved in the ./data/UIEB/All_Results/ folder.
The folder's structure is like:
./data/UIEB/All_Results/
├── FIVE_APLUSNet
│ ├── C60/
│ └── T90/
├── NU2Net
│ ├── C60/
│ └── T90/
├── UIEC2Net
│ ├── C60/
│ └── T90/
├── UTrans
│ ├── C60/
│ └── T90/
└── UWCNN
├── C60/
└── T90/
Each subfolder corresponds to the results of one method.
After testing, you can evaluate any method's perfermance on the test set of UIEB dataset:
python evaluate_UIEB.py --method_name UIEC2Net --folder T90
or you can evaluate perfermance on the challenging set of UIEB dataset:
python evaluate_UIEB.py --method_name UIEC2Net --folder C60
@article{du2023uiedp,
title={UIEDP: Underwater Image Enhancement with Diffusion Prior},
author={Du, Dazhao and Li, Enhan and Si, Lingyu and Xu, Fanjiang and Niu, Jianwei and Sun, Fuchun},
journal={arXiv preprint arXiv:2312.06240},
year={2023}
}