Skip to content
/ TACL Public

TIP 2022 | Twin Adversarial Contrastive Learning for Underwater Image Enhancement and Beyond.

Notifications You must be signed in to change notification settings

Jzy2017/TACL

Repository files navigation

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

avatar

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

🚀 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}}

About

TIP 2022 | Twin Adversarial Contrastive Learning for Underwater Image Enhancement and Beyond.

Topics

Resources

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages