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_epochThen your test images will be output to './outputs/' folder as default.