Home

Awesome

UIE-WD (ICASSP 2022)

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

<div align="center"><img src="./Images/examples.png" width="80%"></div>

Requirement

Dependency

  1. Clone repository

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

  2. Check the requirement.txt

    <code>pip install -r requirements.txt</code>

Dataset

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.

<img src="./Images/dataset.png">

Use the following code to generate training dataset.

python dataset_nyu.py nyu_rgbd_path datapath

Arguments:

Testing

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

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

Arguments:

Update:

2022.4.4

The pretrained model has been updated.

Training

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

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

Arguments:

Environment

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.

Citation