jzy2017 / tacl

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

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

tacl's Issues

How to divide the training&testing set?

Hi, thanks for sharing your method! Could you please give a suggestion about how to divide the 890 images into train set and test set? How to make a fair comparison with other methods?

When the train gonna release?

File "/root/TACL/train.py", line 60, in <module> objs.update(loss.data.item(), opt.batch_size) AttributeError: 'NoneType' object has no attribute 'data'

When I run “python train.py --dataroot ./datasets/UIBED --name chinamm_train --model cycle_gan” ，the error occured ：File "/root/TACL/train.py", line 60, in
objs.update(loss.data.item(), opt.batch_size)
AttributeError: 'NoneType' object has no attribute 'data'

no such file or directory:"./checkpoints/underwater/latest_net_G.pth'

作者好！当我使用readme中的指令“python test.py --dataroot ./datasets/[YOUR-DATASETS] --name underwater --model cycle_gan”进行测试时，程序报错“no such file or directory:"./checkpoints/underwater/latest_net_G.pth'”。
我注意到你们提供了latest_net_G_A.pth和latest_net_G_B.pth，但是没有测试所需的latest_net_G.pth，请问能够提供呢？在此表示感谢！

python train.py出错

Test premodel

When testing a single image, there was no pre trained model provided for testing. Could you please share the model

这是我的运行命令：python test.py --dataroot /data1/Water/SyreaNet/samples/ --name underwater --model test

Weights not found in Google Drive?

Hi, thanks for your work.
Can you provide the weights for the model?

dataset not found!

Python train.py - dataroot/ When it comes to datasets/[YOU DATASETS] -- name Chinamm_train -- model cycl_gan

Thank you for sharing your method. May I ask if I am working on Python train.py - dataroot/ When it comes to datasets/[YOU DATASETS] -- name Chinamm_train -- model cycl_gan, there is none. Save the PT file, can you tell me what the problem is

Download the pre-trained model and put it in ./checkpoints google drive The URL can't go in, do you have this file?

Download the pre-trained model and put it in ./checkpoints google drive The URL can't go in, do you have this file?