A Lightweight Polarization-Guided Plug-in for Underwater Image Enhancement

Introduction

In this project, we use Ubuntu 22.04.3 LTS, Python 3.10.13, Pytorch 2.1.1 cuda11.8 and one NVIDIA RTX 3090 GPU.

datasets

in the paper, we mainly used two polarization datasets

1. RGBP-UIE
2. UCPD

you can get them by follow

RGBP-UIE

UCPD

In the generalization experiment, we used two classic datasets and performed random splitting. LSUI and UIEBD

Ushape+PPM

Test

First, you need to download the trained model weights, or retrain the model weights yourself. Baidu Netdisk(zhth )

In the article we tested a total of three datasets, if you need to test the indoor dataset, run test.py, outdoor dataset, run test_real2.py, ucpd dataset run test_real_opt_data.py. They all load the same weight file, only slightly different in the data processing part.

Training

When you train your model, First, download the pre-training model, then you only need to run train_domain.py (with Domain-adversarial Training). Or you can use the simplified version train.py

PUIE+PPM

You just need to download our code; the other files that need to be executed for the train and test methods are the same as the original project code.

ucolor+PPM

You just need to download our code; the other files that need to be executed for the train and test methods are the same as the original project code.

Citation

@article{ju2025lightweight,
  title={A Lightweight Polarization-Guided Plug-in for Underwater Image Enhancement},
  author={Ju, Guangyao and Zhang, Jiqing and Zang, Jingqi and Li, Yafeng and Li, Yudong and Mi, Zetian and Yang, Xin and Wang, Huibing and Fan, Jiarui and Fu, Xianping},
  journal={IEEE Transactions on Circuits and Systems for Video Technology},
  year={2025},
  publisher={IEEE}
}

Acknowledgement

Our code is based on Ushape, PUIE, and Ucolor. Thanks for their outstanding contributions.