infrontofme / UWGAN_UIE

Source code for UWGAN: Underwater GAN for Real-world Underwater Color Restoration and Dehazing

• This repository contains source code for the method developed in UWGAN: Underwater GAN for real-world underwater color restoration and dehazing
• This code is modified from WaterGAN and UNet image segmentation

Network Structure

UWGAN for generating realistic underwater images

UWGAN takes color image and its depth map as input, then it synthesizes underwater realistic images based on underwater optical imaging model by learning parameters through generative adversarial training. You can find more details in the paper.

Synthetic underwater-style images through UWGAN. (a) are in-air sample images, (b)-(d) are synthetic underwater-style sample images of different water types.

Underwater image restoration based on UNet

Proposed U-net Architecture for underwater image restoration and enhancement. The effects of different loss functions in U-net are compared, the most suitable loss function for underwater image restoration is suggested based on the comparison, you can find more details about loss function in paper.

Dataset

Download data:

1. In-air RGBD data: NYU Depth Dataset V1, NYU Depth Dataset V2

2. Underwater images: [Baidu Cloud Link] [Google Drive]

3. UIEB Dataset for verification: [github link]

4. The NYU datasets we used to train UWGAN: [Baidu Cloud Link] [Google Drive]

5. fake water images generated from UWGAN: [Google Drive]

6. pretrained model:[Google Drive]

Data directory structure in UWGAN

.
├── ...
├── data                    
│   ├── air_images
│   │   └── *.png
│   ├── air_depth  
│   │   └── *.mat
│   └── water_images 
│       └── *.jpg
└── ...

Usage

• Train a UWGAN model - Firstly, change directory to UWGAN folder, then run python uwgan_mian.py, you can adjust learning parameters in uwgan_main.py.
• Train a UNet restoration model - Firstly, change directory to UNetRestoration folder, then run python train.py, you can adjust learning parameters and change loss functions in train.py. Run python test.py after training has been completed.

Results and Discussion

• Underwater image restoration through our approach. Our method is tested on both UIEB public benchmark and Our underwater images collected from marine organisms’ farms.

• High-level computer vision task: underwater target detection on underwater images before and after processing with our method.

• The effects of different loss functions in restoration network (UNet)

  

Citations

If you find this work useful for your research, please cite this article in your publications.

@misc{wang2019uwgan,
    title={UWGAN: Underwater GAN for Real-world Underwater Color Restoration and Dehazing},
    author={Nan Wang and Yabin Zhou and Fenglei Han and Haitao Zhu and Yaojing Zheng},
    year={2019},
    eprint={1912.10269},
    archivePrefix={arXiv},
    primaryClass={eess.IV}
}