WaterGAN

• This repository contains source code for WaterGAN developed in WaterGAN: Unsupervised Generative Network to Enable Real-time Color Correction of Monocular Underwater Images.
• This code is modified from Taehoon Kim's DCGAN-tensorflow (MIT-licensed). Our modifications are MIT-licensed.

Usage

Download data:

1. MHL test tank dataset: MHL.tar.gz
2. Jamaica field dataset: Jamaica.tar.gz
3. In air data: Any RGB-D dataset, e.g. Microsoft 7-Scenes, NYU Depth, UW RGB-D Object, B3DO
   Note: The current configuration expects 640x480 PNG images for in-air data.

Directory structure:

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

Train a model with the MHL dataset:

python mainmhl.py --water_dataset water_images --air_dataset air_images --depth_dataset air_depth

Train a model with the Jamaica dataset:

python mainjamaica.py --water_dataset water_images --air_dataset air_images --depth_dataset air_depth

Color Correction Network

WaterGAN outputs a dataset with paired true color, depth, and (synthetic) underwater images. We can use this to train an end-to-end network for underwater image restoration. Source code and pretrained models for the end-to-end network are available here. For more details, see the paper.

Citations

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

@article{Li:2017aa,
	Author = {Jie Li and Katherine A. Skinner and Ryan Eustice and M. Johnson-Roberson},
	Date-Added = {2017-06-12 22:07:13 +0000},
	Date-Modified = {2017-06-12 22:12:20 +0000},
	Journal = {IEEE Robotics and Automation Letters (RA-L)},
	Keywords = {jrnl},
	Note = {accepted},
	Title = {WaterGAN: Unsupervised Generative Network to Enable Real-time Color Correction of Monocular Underwater Images},
	Year = {2017}}

exp_2 runtime options

The modelmhl.py and modeljamaica.py entry points support both legacy MAT depth files and compact PNG depth files. Image and depth decoding can run in a shared thread pool without changing the configured batch size.

Runtime environment variables:

export CUDA_VISIBLE_DEVICES=2
export WATERGAN_IO_WORKERS=16
export WATERGAN_LOG_EVERY=10
export WATERGAN_THROTTLE_DIAGNOSTICS=1

• WATERGAN_IO_WORKERS controls parallel image/depth decoding.
• WATERGAN_LOG_EVERY controls the diagnostic interval.
• WATERGAN_THROTTLE_DIAGNOSTICS=1 avoids summary, loss, and parameter diagnostic fetches on every batch. It does not change the one-discriminator and two-generator optimizer update schedule.
• GPU selection is controlled by CUDA_VISIBLE_DEVICES; the Python entry points no longer override it.