Skip to content
Source code for UWGAN: Underwater GAN for Real-world Underwater Color Restoration and Dehazing
Python
Branch: master
Clone or download
infrontofme Update README.md
Add Google Drive Link
Latest commit 81c71b9 Jan 4, 2020
Permalink
Type Name Latest commit message Commit time
Failed to load latest commit information.
Figures Add files via upload Dec 24, 2019
UNetRestoration Add files via upload Dec 24, 2019
UWGAN Add files via upload Dec 24, 2019
LICENSE Create LICENSE Dec 24, 2019
README.md Update README.md Jan 4, 2020

README.md

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.

UWGAN structure

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

Synthetic underwater-style images

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.

UNet structure

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]

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

public dataset enhancement our dataset enhancement

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

underwater target detection

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

    different loss functions

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}
}
You can’t perform that action at this time.