- 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.
Download data:
- MHL test tank dataset: MHL.tar.gz
- Jamaica field dataset: Jamaica.tar.gz
- 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
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.
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}}