Fast Style Transfer in TensorFlow 2

This is an implementation of Fast-Style-Transfer on Python 3 and Tensorflow 2. The neural network is a combination of Gatys' A Neural Algorithm of Artistic Style, Johnson's Perceptual Losses for Real-Time Style Transfer and Super-Resolution, and Ulyanov's Instance Normalization.

Image Stylization 🎨

Added styles from various paintings to a photo of Chicago. Check the ./images/results folder to see more images.

All the models were trained on the same default settings.

Implementation Details

• The feed-forward network is roughly the same as described in Johnson, except that batch normalization is replaced with Ulyanov's instance normalization, and the scaling/offset of the output tanh layer is slightly different (for better convergence), also use Resize-convolution layer to replace the regular transposed convolution for better upsampling (to avoid checkerboard artifacts)
• The loss network used in this implementation follows Logan Engstrom , which is similar to the one described in Gatys , using VGG19 instead of VGG16 and typically using "shallower" layers than in Johnson's implementation, for larger scale style features in transformation (e.g. use relu1_1 rather than relu1_2).

Training Style Transfer Networks

Use main.py to train a new style transfer network. Training takes 5~6 hours on a GTX 1060 3GB (when batch size is 2). Before you run this, you should run setup.sh to download the dataset.

Example usage:

python main.py train    \
  --style ./path/to/style/image.jpg \
  --dataset ./path/to/dataset \
  --weights ./path/to/weights \
  --batch 2    

Evaluating Style Transfer Networks

Use main.py to evaluate a style transfer network. Evaluation takes 2s per frame(712x474) on a GTX 1060 3GB. Models for evaluation are located here.

Example usage:

python main.py evaluate    \
  --weights ./path/to/weights \
  --content ./path/to/content/image.jpg(video.mp4) \
  --result ./path/to/save/results/image.jpg

Requirements

You will need the following to run the above:

• TensorFlow >= 2.0
• Python 3.7.5, Pillow 7.0.0, Numpy 1.18, Opencv 4.1.2
• If you want to train (and don't want to wait too long):
  • A decent GPU
  • All the required NVIDIA software to run TF on a GPU (cuda, etc)

Attributions/Thanks

• Some readme/docs formatting was borrowed from Logan Engstrom's fast-style-transfer
• Some code was borrowed from TensorFlow documentation Neural style transfer