Paint Transformer: Feed Forward Neural Painting with Stroke Prediction

[Paper] [PyTorch Implementation] [Paddle Implementation] [Huggingface Gradio Demo]

Overview

This repository contains the official PyTorch implementation of paper:

Paint Transformer: Feed Forward Neural Painting with Stroke Prediction,

Songhua Liu*, Tianwei Lin*, Dongliang He, Fu Li, Ruifeng Deng, Xin Li, Errui Ding, Hao Wang (* indicates equal contribution)

ICCV 2021 (Oral)

Prerequisites

Linux or macOS
Python 3
PyTorch 1.7+ and other dependencies (torchvision, visdom, dominate, and other common python libs)

Getting Started

Clone this repository:

git clone https://github.com/Huage001/PaintTransformer
cd PaintTransformer

Download pretrained model from Google Drive and move it to inference directory:
```
mv [Download Directory]/model.pth inference/
cd inference
```
Inference:
```
python inference.py
```
- Input image path, output path, and etc can be set in the main function.
- Notably, there is a flag serial as one parameter of the main function:
  - If serial is True, strokes would be rendered serially. The consumption of video memory will be low but it requires more time.
  - If serial is False, strokes would be rendered in parallel. The consumption of video memory will be high but it would be faster.
  - If animated results are required, serial must be True.
Train:
- Before training, start visdom server:
```
python -m visdom.server
```
- Then, simply run:
```
cd train
bash train.sh
```
- You can monitor training status at http://localhost:8097/ and models would be saved at checkpoints/painter folder.
You may feel free to try other training options written in train.sh.

More Results

Input	Animated Output

Citation

If you find ideas or codes useful for your research, please cite:

@inproceedings{liu2021paint,
  title={Paint Transformer: Feed Forward Neural Painting with Stroke Prediction},
  author={Liu, Songhua and Lin, Tianwei and He, Dongliang and Li, Fu and Deng, Ruifeng and Li, Xin and Ding, Errui and Wang, Hao},
  booktitle={Proceedings of the IEEE International Conference on Computer Vision},
  year={2021}
}

Acknowledgments

This implementation is developed based on the code framework of pytorch-CycleGAN-and-pix2pix by Junyan Zhu et al.

dandelin / PaintTransformer