Home

Awesome

STGAN in Pytorch

This is a Pytorch implementation of the STGAN model from the paper "Cloud Removal in Satellite Images Using Spatiotemporal Generative Models," Sarukkai<sup>*</sup>, Jain<sup>*</sup>, Uzkent, and Ermon (https://arxiv.org/abs/1912.06838). The STGAN is accepted to IEEE WACV 2020.

This implementation is built on a clone of the implementation of Pix2Pix from https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix. The readme also is based on the readme from the CycleGAN/pix2pix repo. That repo contains the code corresponding to the following two papers:

Unpaired Image-to-Image Translation using Cycle-Consistent Adversarial Networks.<br> Jun-Yan Zhu*, Taesung Park*, Phillip Isola, Alexei A. Efros. In ICCV 2017. (* equal contributions) [Bibtex]

Image-to-Image Translation with Conditional Adversarial Networks.<br> Phillip Isola, Jun-Yan Zhu, Tinghui Zhou, Alexei A. Efros. In CVPR 2017. [Bibtex]

Prerequisites

Getting Started

Installation

git clone https://github.com/ermongroup/STGAN.git
cd stgan

STGAN dataset

Download the STGAN dataset from https://doi.org/10.7910/DVN/BSETKZ. This download includes two datasets (singleImage and multipleImage), described both in the paper and on the linked page. Before training a model, these images must be split into train/val/test splits--for instance, for both singleImage/clear and singleImage/cloudy, create three subfolders called "train","val","test", and assign the images from the main directory into the corresponding partitions either using a script that assigns splits at random or a split that keeps images from the same "tile" together (accounting for the possibility that these images are correlated). In our case, we used a 80-10-10 train-val-test split while keeping image crops from the same "tile" together. Make sure that corresponding clear and cloudy images are assigned to the same split.

Once the data is formatted this way, call:

python datasets/combine_A_and_B.py --fold_A /path/to/data/cloudy --fold_B /path/to/data/clear --fold_AB /path/to/data/combined

This will combine each pair of images (cloudy,clear) into a single image file, ready for training.

STGAN training

python train.py --dataroot ./path/to/data/combined  --name stgan --model temporal_branched_ir --netG unet_256_independent --input_nc 4
python test.py --dataroot ./path/to/data/combined  --name stgan --model temporal_branched_ir --netG unet_256_independent --input_nc 4
<img src="./results/STGAN.png" alt="WAMI_Positives" style="width: 300px;"/>

You can cite our paper as :

@article{sarukkai2019cloud,
  title={Cloud Removal in Satellite Images Using Spatiotemporal Generative Networks},
  author={Sarukkai, Vishnu and Jain, Anirudh and Uzkent, Burak and Ermon, Stefano},
  journal={arXiv preprint arXiv:1912.06838},
  year={2019}
}