/FastStyle

Fast-Style-Transfer in TensorFlow 2 :art:

Primary LanguagePythonMIT LicenseMIT

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