Underwater Image Enhancement with Cascaded Contrastive Learning

This repository is the official PyTorch implementation of CCL-Net.

Dataset Preparation

You should prepare the structure of datasets folder as follows:

├──path_to_data
    ├── train
            ├── raw
                ├── im1.png
                ├── im2.png
                └── ...
            ├── ref
                ├── im1.png
                ├── im2.png
                └── ...
    ├── test
            ├── raw
                ├── im1.png
                ├── im2.png
                └── ...

Running Environment

1. python 3.8
2. pip install -r requirements.txt
3. pip install torch==2.0.0 torchvision==0.15.1 --index-url https://download.pytorch.org/whl/cu118

Test with the pre-trained models

1. rename directory pre-trained to checkpoints
2. run command 'python test.py --dataroot ./imgs --model HRNet'

CC-Net

Train

python train.py --dataroot /the_abs_path_of_data --model CCNet --lr 0.0005

Test

python test.py --dataroot /the_abs_path_of_data --model CCNet

HR-Net

Notice

Training or Testing the HR-Net need load the pre-trained CC-Net. Therefore, you should have trained the CC-Net before you start to train or test the HR-Net.

Train

python train.py --dataroot /the_abs_path_of_data --model HRNet --lr 0.001

Test

python test.py --dataroot /the_abs_path_of_data --model HRNet

Citation

TODO: paste the avalible link at here

Acknowledgements

• https://github.com/trentqq/SGUIE-Net_Simple
• https://github.com/Li-Chongyi/Ucolor
• https://github.com/GlassyWu/AECR-Net
• https://github.com/zhilin007/FFA-Net
• https://github.com/swz30/MIRNetv2