Twin-Adversarial-Contrastive-Learning-for-Underwater-Image-Enhancement-and-Beyond

This is an implement of the TACL, “Twin-Adversarial-Contrastive-Learning-for-Underwater-Image-Enhancement-and-Beyond”, Risheng Liu*, Zhiying Jiang, Shuzhou Yang, Xin Fan, IEEE Transactions on Image Processing (TIP), 2022.

Overview

Prerequisites

• Linux or macOS
• Python 3
• CPU or NVIDIA GPU + CUDA CuDNN

🔑 Installation

Type the command:

pip install -r requirements.txt

🤖 Download

Download the pre-trained model and put it in ./checkpoints

• Google Drive
• Baidu Yun
  code: ​ YSZD ​

🚀 Quick Run

• Create directories ./dataset/testA and ./dataset/testB. Put your test images in ./dataset/testA (And you should keep whatever one image in ./dataset/testB to make sure program can start.)
• To test the pre-trained models for Underwater Enhancement on your own images, run

python test.py --dataroot ./datasets/[YOUR-DATASETS] --name underwater --model cycle_gan

Results will be shown in results folder.

Train Backbone

• First, you need to train a base backbone:

python train.py --dataroot ./datasets/[YOUR-DATASETS] --name chinamm_train --model cycle_gan

Training TAF

• Second, you need to train a TAF module (here we adopt SSD):

  • Download an Underwater Detection Dataset (Chinamm).
  • Run this to make Chinamm in VOC format:

  python makeTXT.py

  • Use the trained backbone to enhance JPEGImages of chinamm.
  • cd ./ssd.pytorch-master

• Download the fc-reduced VGG-16 PyTorch base network weights at: https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth

• By default, we assume you have downloaded the file in the ssd.pytorch/weights dir:

mkdir weights
cd weights
wget https://s3.amazonaws.com/amdegroot-models/vgg16_reducedfc.pth

• To train SSD using the train script simply specify the parameters listed in train.py as a flag or manually change them.

python train.py

• Note:

  • For training, an NVIDIA GPU is strongly recommended for speed.
  • For instructions on Visdom usage/installation, see the Installation section.
  • You can pick-up training from a checkpoint by specifying the path as one of the training parameters (again, see train.py for options)

• Evaluation To evaluate a trained network:

python eval.py

You can specify the parameters listed in the eval.py file by flagging them or manually changing them.

Training

cd ./ssd.pytorch-master
Run

python trainall.py

• Test final version:

python visual.py

📌 Citation

If you find this code useful for your research, please use the following BibTeX entry.

@ARTICLE{9832540,
  author={Liu, Risheng and Jiang, Zhiying and Yang, Shuzhou and Fan, Xin},
  journal={IEEE Transactions on Image Processing}, 
  title={Twin Adversarial Contrastive Learning for Underwater Image Enhancement and Beyond}, 
  year={2022},
  volume={31},
  number={},
  pages={4922-4936},
  doi={10.1109/TIP.2022.3190209}}