Underwater image-to-image translation using CycleGAN in PyTorch

Prerequisites

• Linux OS
• Python 3
• CPU or NVIDIA GPU + CUDA CuDNN

Getting Started

Installation

• Clone this repo:

  git clone https://github.com/ioannispol/UnderWaterGAN.git
  cd UnderWaterGAN

• Some sample scripts are provided in ./scripts. Before you can run them the first time, set execute permission via
  chmod +x -R ./scripts/

• Install dependencies via pip
  #!./scripts/install_deps.sh

  • PyTorch
  • visdom
  • For all other dependencies, use the command
    pip install -r requirements.txt.

Prepare Your Dataset

To train a model on your own dataset, you need to create a data folder with two subdirectories trainA and trainB that contain images from domain A and B (unaired). You can test your model on your training set by setting --phase train in test.py. You can also create subdirectories testA and testB if you have test data.

Training and Testing

• Prepare your dataset as described above
• To view training results and loss plots, run python -m visdom.server and click the URL http://localhost:8097.
• Train a model:

  #!./scripts/train_cyclegan.sh
  python train.py --dataroot /{path_to_train_data} --name {experiment_name}

  • To specify more train options, see the files in the directory ./options/train_options.py
  • To see more intermediate results, check out ./results/{experiment_name}/web/index.html.
• Test the model:

  #!./scripts/test_cyclegan.sh
  python test.py --dataroot /{path_to_test_data} --name {experiment_name} --model cycle_gan

  • To specify more test options, see the files in the directory ./options/test_options.py
  • The test results will be saved to a html file here: ./results/{experiment_name}/test_latest/index.html.
  • The (pre-)trained model is saved at ./results/{experiment_name}/latest_net_G_{A_or_B}.pth.
  • The option --model test is used for generating results of CycleGAN only for one side. This option will automatically set --dataset_mode single, which only loads the images from one set. In this case it might be neccesary to specify --model_suffix "_{A_or_B}" to find the correct models.
    On the contrary, using --model cycle_gan requires loading and generating results in both directions, which is sometimes unnecessary.
  • The results will be saved at ./results/. Use --results_dir {directory_path_to_save_result} to specify the results directory.

Network architecture For Cycle A to B

Citation

If you use this project for your research, please cite our [paper].

ToDo

• Add evaluation scripts
• Make evaluation scripts for FID and SSIM-py
• Paper TBR

Acknowledgments

The project is forked from m4ln/stainTransfer_CycleGAN_pytorch: Image-to-Image Translation in PyTorch. Only a few parts are modified or added. Please refer to this project for more details on the project files.

UnderWaterGAN DATASET

The dataset used in the project can be found at