Ching-Heng Cheng, Jen-Wei Lee, Chia-Ming Lee, Chih-Chung Hsu

Advanced Computer Vision LAB, National Cheng Kung University and National Yang Ming Chiao Tung University.

This is a PyTorch implementation for "WWE-UIE: A Wavelet & White Balance Efficient Network for Underwater Image Enhancement." This document summarizes the environment requirements, dataset layout, and common commands you need to run the project.

| | |

| | | |

• Python 3.10 (recommended)
• PyTorch 2.4.0 + CUDA 11.8 (recommended)
• Additional packages listed in requirements.txt

pip install -r requirements.txt

WWE-UIE/
├── model.py                 # Wavelet-based backbone
├── train.py                 # paired training entry point
├── test.py                  # paired evaluation (PSNR/SSIM)
├── test_nr.py               # non-reference evaluation (NIQE/UCIQE/URanker)
└── utils/
    ├── dataset.py           # paired & non-ref dataloaders
    ├── loss_funcs.py        # perceptual / SSIM / edge-aware loss
    ├── metrics.py           # PSNR, SSIM, NIQE, UCIQE, URanker
    ├── utils.py             # plotting & visualization helpers
    └── ...                  # CIDNet & URanker configs / weights

All dataset paths default to UnderWaterDataset/. Feel free to relocate them, but keep the directory layout identical. Paired datasets require matching filenames under input/ and GT/. Non-reference datasets only need raw images.

⬇ We provide our splitted dataset on Google drive.

UnderWaterDataset/
├── EUVP-Dark/
│   ├── train/
│   │   ├── input/xxx.png
│   │   └── GT/xxx.png
│   ├── val/
│   │   ├── input/xxx.png
│   │   └── GT/xxx.png
│   └── test/
│       ├── input/xxx.png
│       └── GT/xxx.png
├── EUVP-Scene/...
├── LSUI/...
├── UFO-120/...
└── UIEB/...

Add your own dataset by creating train/<name>/input, train/<name>/GT, and matching folders under val/ and test/. Update train.py (dataset switch) or override the CLI arguments to point at the new folders.

UnderWaterDataset/
├── Challenging-60/test/
│   └── *.png
├── EUVP/test/
│   └── *.png
├── U45/test/
│   └── *.png
└── UCCS/test/
    ├── blue/*.png
    ├── blue-green/*.png
    └── green/*.png

[UIEB]

Training from scratch is efficient as our light-weight design.

python train.py \
  --dataset [dataset] \
  --epoch 100 \
  --train_batch_size 24 \
  --model_name WWE-UIE

Checkpoints, logs, and predictions are written to ./output/<model_name>/<dataset>/<timestamp>/:

• best_model.pth: weights from the best PSNR checkpoint
• records.txt: validation PSNR/SSIM per epoch

python test.py \
  --dataset [dataset] \
  --ckpt [checkpoint.pth] \

Note on non-reference metrics: we re-implement NIQE and UCIQE in Python, so the values may differ slightly from those reported in the original papers (previous measured by public github repo. in Matlab). However, our method still achieves top-ranked relative scores.

If you would like to include URanker as one of the evaluation metrics, please download the URanker pretrained weights and place them under utils/uranker/.

python test_nr.py \
  --dataset [dataset] \
  --ckpt [checkpoint.pth] \

If you use this code for your research, please cite our papers.

@misc{cheng2025wweuiewaveletwhite,
      title={WWE-UIE: A Wavelet & White Balance Efficient Network for Underwater Image Enhancement}, 
      author={Ching-Heng Cheng and Jen-Wei Lee and Chia-Ming Lee and Chih-Chung Hsu},
      year={2025},
      eprint={2511.16321},
      archivePrefix={arXiv},
      primaryClass={cs.CV},
      url={https://arxiv.org/abs/2511.16321}, 
}