Skip to content

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

License

Notifications You must be signed in to change notification settings

infrontofme/UWGAN_UIE

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

21 Commits
 
 
 
 
 
 
 
 
 
 

Repository files navigation

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]

  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

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}
}

About

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

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

No packages published

Languages