This is the PyTorch implementation for "A Wavelet-based Dual-stream Network for Underwater Image Enhancement", ICASSP, 2022. We present a wavelet-based dual-stream network that addresses color cast and blurry details in underwater images. More results can be found in our webpage. paper | webpage

Dependency

• Python 3.7
• PyTorch 0.4.1
• numpy
• skimage
• matplotlib
• tqdm
• cv2
• Pillow

1. Clone repository

   git clone https://github.com/ZZiyin/UIE-WD_Code.git

2. Check the requirement.txt

   pip install -r requirements.txt

The synthesized training dataset is generated from NYU Depth V2 following Anwar et al. (2018), which consists of 1449 images. The dataset is augmented by generating 6 images of each class using random parameters, thus for each ground truth image, we have corresponding 36 images of different water types. The following are the examples of generated images.

Use the following code to generate training dataset.

python dataset_nyu.py nyu_rgbd_path datapath

Arguments:

• nyu_rgbd_path: The path of NYU Depth V2 dataset (.mat)
• datapath: The path of folder where you would like to store the genereated synthesized training images.

Pretrained models can be found in the google drive, place the pretrained model in ./checkpoints folder. Use the following code to generate results. And the enhanced images can be found in ./results.

python test_multi.py multi --test_dataset UIEB --data_path ../UIEB/raw-890 --model_load_path checkpoints/multi/model.pth

Arguments:

• --test_dataset: The name of dataset you would like to use for testing.
• --data_path: The path of dataset you would like to use for testing.
• --model_load_path: The path of pretrained model.
• --test_size (optional) : The size of test dataset, default = 890.

Update:

2022.4.4

The pretrained model has been updated.

Use the following code to train the model. The weights will be saved every 2 epoch, which can be found in ./checkpoints.

python train_multi.py multi --data_path ../image/data --label_path ../image/label 

Arguments:

• --data_path: The path of training input data.
• --label_path: The path of training label data.
• --batch_size (optional): Batch_size of training model. default = 4.
• --save_interval (optional): Save models after this many epochs, default = 2.

We conduct training and testing on Intel Core i5-7200 CPU and NVIDIA Geforce RTX 2070 GPU. Noted that the provided model is retrianed on Tesla T4 as the previous server is not accessible at the moment, we will update the pretrained model soon.