Underwater Ranker: Learn Which Is Better and How to Be Better (AAAI 2023, Oral Presentation)

This repository contains the official implementation of the following paper:

Underwater Ranker: Learn Which Is Better and How to Be Better
Chunle Guo^#, Ruiqi Wu^#, Xin Jin, Linghao Han, Zhi Chai, Weidong Zhang, Chongyi Li^*
Proceedings of the AAAI conference on artificial intelligence (AAAI), 2023

[Arxiv Paper] [中文版 (TBD)] [Project] [Dataset (pwd: nuin)]

News

• [07/05/2023] We released the Google Drive link of our model checkpoints.
• [03/30/2023] Our URanker has been added in Pytorch IQA.
• [01/11/2023] Our paper was awarded oral presentation qualification by AAAI.

Dependencies and Installation

1. Clone Repo

   git clone https://github.com/RQ-Wu/UnderwaterRanker.git
   cd UnderwaterRanker

2. Create Conda Environment

   conda env create -f environment.yaml
   conda activate underwater_ranker

Get Started

Prepare pretrained models & dataset

1. You are supposed to download our pretrained model first in the links below and put them in dir ./checkpoints/:

Model	SRCC/ KRCC (PSNR / SSIM)	🔗 Download Links
URanker	0.8655 / 0.7402	[Google Drive] [Baidu Disk (pwd: fu7y)]
NU2Net	22.669 / 0.9246

2. Two datasets used in our work can be downloaded in the links below:

   • URankerSet: [Google Drive (TBD)] [Baidu Disk (pwd: nuin)]
   • Underwater Image Enhancement Benchmark (UIEB): [Download Link]

   The data is put in dir ./data/.

The directory structure will be arranged as:

checkpoints
    |- URanker_ckpt.pth
    |- NU2Net_ckpt.pth
data
    |- UIEB
        |- raw-890
        |- reference-890
    |- UIERank

Quick demo

Run demos to process the images in dir ./examples/ by following commands:

python ranker_demo.py \
     --opt_path options/URanker.yaml \
     --checkpoint_path checkpoints/URanker_ckpt.pth \
     --input_path examples/ranker_example \
     --save_path results/ranker_result.txt

python uie_demo.py \
     --opt_path options/NU2Net.yaml \
     --checkpoint_path checkpoints/NU2Net_ckpt.pth \
     --input_path examples/uie_example \
     --save_path results

Training & Evaluation

Our training and evaluation configures are provided in options/URanker.yaml (for URanker) and options/NU2Net.yaml (for NU²Net)

Run the following commands for training:

python ranker_main_train.py --opt_path options/URanker.yaml
python uie_main_train.py --opt_path options.NU2Net.yaml

Run the following commands for evaluation:

python ranker_main_test.py --opt_path options/URanker.yaml --test_ckpt_path checkpoints/URanker_ckpt.pth
python uie_main_test.py --opt_path options.NU2Net.yaml --test_ckpt_path checkpoints/NU2Net_ckpt.pth --save_image

Citation

If you find our repo useful for your research, please cite us:

@inproceedings{guo2023uranker,
  title={Underwater Ranker: Learn Which Is Better and How to Be Better},
  author={Guo, Chunle and Wu, Ruiqi and Jin, Xin and Han, Linghao and Chai, Zhi and Zhang, Weidong and Li, Chongyi},
  booktitle={Proceedings of the AAAI Conference on Artificial Intelligence},
  year={2023}
}

License

Licensed under a Creative Commons Attribution-NonCommercial 4.0 International for Non-commercial use only. Any commercial use should get formal permission first.

Acknowledgement

This repository is maintained by Ruiqi Wu.