The official repository for paper "PREGAN: Pose Randomization and Estimation for Weakly Paired Image Style Translation", accepted for IEEE Robotics & Automation Letters (RAL).
There are a few dependencies required to run the code. They are listed below:
Python 3.7
CUDA 10.1
CuDnn
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.
Download the model here:
https://v2.fangcloud.com/share/09582731e63b1008d819c199eb
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
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
└── ...
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
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/
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 |
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.