PREGAN: Pose Randomization and Estimation for Weakly Paired Image Style Translation

The official repository for paper "PREGAN: Pose Randomization and Estimation for Weakly Paired Image Style Translation", accepted for IEEE Robotics & Automation Letters (RAL).

Dependencies

There are a few dependencies required to run the code. They are listed below:

System Environments:

Python 3.7

CUDA 10.1

CuDnn

Pip dependencies:

torch>=1.6.0

torchvision>=0.7.0

dominate>=2.4.0

visdom>=0.1.8.8

Kornia

TensorboardX

opencv-python

matplotlib

torchviz

You can install these dependencies by changing your current directory to the PRoGAN directory and running:

pip install -r requirements.txt  

Or manually install the dependencies, and Conda or Virtualenv is recommended to use.

Using Pre-trained model

Download the model here:

https://v2.fangcloud.com/share/09582731e63b1008d819c199eb

Getting Started

Datasets

If you want to train the Aeroground dataset or Carla dataset:

• Aeroground: https://github.com/ZJU-Robotics-Lab/OpenDataSet

• WeakPaired: https://github.com/ZJU-Robotics-Lab/OpenDataSet

Prepare your own dataset:

mkdir datasets
cd datasets
# The dataset structure is shown as following:
├──Your dataset
  ├── trainA (realA images)
      ├── Image0001.jpg 
      └── ...
  ├── trainB (realB images)
      ├── Image0001.jpg
      └── ...
  ├── testA (testing realA images)
      ├── Image3000.jpg
      └── ... 
  ├── testb (testing realB images)
      ├── Image3000.jpg
      └── ... 

Example Usage

Training

If you want to train the PRoGAN network, then run:

# on aeroground
python main.py --phase train --aeroground
# on weakPaired
python main.py --phase train --carla
# on other dataset
python main.py --phase train --name training_name --batch_size 1  --netG resnet_9blocks --load_size 256 --dataset your_dataset --input_nc your_image_channel --output_nc your_image_channel 

Training visualize

To visualize the training process, you can run:

# use visdom
python -m visdom.server # http://localhost:8097
# use tensorboard
python tensorboard --logdir checkpoints/log/your_training_name/

Training options

By default, this will train the network on the stereo_to_aerial dataset with the batch size of 1, learning rate of 0.00015 and run on GPU 0. There are several settings you can change by adding arguments below:

Arguments	What it will trigger	Default
--gpu_ids	The ids of gpu to use	0
--checkpoints_dir	The place to save models	'./checkpoints/'
--input_nc	The channel of input image	3
--output_nc	The channel of output image	3
--batch_size	The batch size of input	1
--load_size	The size of input image for network (128 / 256)	128
--continue_train	Continue to train
--epoch	The start epoch for continuing to train	'latest'
--phase	Choose to train or validate (train / val)	'train'
--lr	The learning rate for training	0.00015
--train_writer_path	Where to write the Log of training	'./checkpoints/log/'
--val_writer_path	Where to save the images of validating	'./outputs/'
--aeroground	Use dataset: AeroGround
--carla	Use dataset: WeakPaired

Validating

To validate on the dataset, you could run:

# on aeroground
python main.py --phase val --aeroground
# on your dataset
python main.py --phase val --name training_name --batch_size 1  --netG resnet_9blocks --load_size 256 --dataset your_dataset --input_nc your_image_channel --output_nc your_image_channel --epoch your_test_epoch

Then your test images will be output to './outputs/' folder as default.