gopirajumatta / uw-net Goto Github PK

[Project] [Paper]

This repository contains the tensorflow implementation for UW-Net and includes the scripts to train and test the network. The code is written and maintained by Honey Gupta.

• The following packages are required to run the codes. Note: the version of the packages were the ones we used and are suggestive, not mandatory.
  • python = 3.6
  • click = 7.0
  • tensorflow-gpu = 1.11
  • scikit-image = 0.15

• You can download any existing underwater image dataset or create your own by collecting images from the internet (while considering their copyrights).
  • Put images from each domain i.e RGBD terrestrial images in the form of .npy files and 2D underwater images in different folders
• Create the csv file as input to the data loader.

  • Edit the uwnet_datasets.py file.

    For example, if your dataset contains 1449 RGBD images and <1449 underwater images as .png files, you can just edit the uwnet_datasets.py as following

   DATASET_TO_SIZES = {
   'air2water_train': 1449
   }
   PATH_TO_CSV = {
   'air2water_train': 'input/air2water_train.csv'
   }
   DATASET_TO_IMAGETYPE = {
   'air2water_train': '.png'
   }

  • Run create_uwnet_dataset.py

   python create_uwnet_dataset --image_path_a=/path/to/folder/with/rgbd/npy/files \
    --image_path_b=/path/to/folder/containing/underwater/images --dataset_name="air2water_train"

• Create the configuration file. The configuration file contains basic information for training/testing. An example of the configuration file could be fond at configs/exp_01.json.

• Start training:

  python main.py \
      --to_train=1 \
      --log_dir=out/exp_01 \
      --config_filename=configs/exp_01.json

• Check the intermediate results.

  • Tensorboard

   tensorboard --port=6006 --logdir=out/exp_01/#timestamp# 

  • Check the html visualization at out/exp_01/#timestamp#/epoch_#id#.html

python main.py \
    --to_train=2 \
    --log_dir=out/exp_01 \
    --config_filename=configs/exp_01.json \
    --checkpoint_dir=out/exp_01/#timestamp#

You can test the models in two ways: (1) using the test.py file that does not require the RGBD dataset (2) Using the main.py file. I recommend using the test.py as it does not require the RGBD dataset.

• Create the testing dataset

  • Edit the uwnet_datasets.py file the same way as training
    • A sample entry for the dataset by Berman et al. can be found in the script
  • Create the csv file as the input to the data loader. To create the csv file containing only the underwater image list:

  python create_uwnet_dataset.py \
       --image_path_b=/path/to/folder/containing/underwater/images --dataset_name="hazelines" --mode="test"

  To create a csv file containg both UW and RGBD file names:

  python create_uwnet_dataset.py --image_path_a=/path/to/folder/with/rgbd/npy/files \
       --image_path_b=/path/to/folder/containing/underwater/images --dataset_name="hazelines"

• Run testing

  python test.py \
      --log_dir=out/exp_01_test \
      --config_filename=configs/exp_01_test.json \
      --checkpoint_dir=out/exp_01/#timestamp# 

  OR

  python main.py \
      --to_train=0 \
      --log_dir=out/exp_01_test \
      --config_filename=configs/exp_01_test.json \
      --checkpoint_dir=out/exp_01/#timestamp# 

The result is saved in out/exp_01_test/#timestamp#.

Download the pre-trained models from here and extract the files in the working folder.

In the provided checkpoints, the folder checkpoints/pre-trained/ contains the model trained on our collected underwater dataset and checkpoints/finetune contains the model from checkpoints/pre-trained/ fine-tuned on Berman et al.'s dataset for 20 epochs.

To test on the pre-trained models, change the entry for --checkpoint_dir in the above command line script.

Kindly cite our paper if this repository is useful for your research.

@article{gupta2019unsupervised, 
  title={Unsupervised Single Image Underwater Depth Estimation}, 
  author={Gupta, Honey and Mitra, Kaushik}, 
  journal={arXiv preprint arXiv:1905.10595}, 
  year={2019}
}

This project is licensed under the MIT License - see the LICENSE file for details.

This repository is partially built upon the CycleGAN repository written by Harry Yang and Nathan Silberman.