search

rrmina / fast-neural-style-pytorch

Fast Neural Style Transfer implementation in PyTorch :art: :art: :art:
315 stars 77 forks source link
convolutional-neural-networks neural-network pytorch style-transfer

readme

fast-neural-style: Fast Style Transfer in Pytorch! :art:

An implementation of fast-neural-style in PyTorch! Style Transfer learns the aesthetic style of a style image, usually an art work, and applies it on another content image. This repository contains codes the can be used for:

  1. fast image-to-image aesthetic style transfer,
  2. image-to-video aesthetic style transfer, and for
  3. training style-learning transformation network

This implemention follows the style transfer approach outlined in Perceptual Losses for Real-Time Style Transfer and Super-Resolution paper by Justin Johnson, Alexandre Alahi, and Fei-Fei Li, along with the supplementary paper detailing the exact model architecture of the mentioned paper. The idea is to train a separate feed-forward neural network (called Transformation Network) to transform/stylize an image and use backpropagation to learn its parameters, instead of directly manipulating the pixels of the generated image as discussed in A Neural Algorithm of Artistic Style aka neural-style paper by Leon A. Gatys, Alexander S. Ecker, and Matthias Bethge. The use of feed-forward transformation network allows for fast stylization of images, around 1000x faster than neural style.

This implementation made some modifications in Johnson et. al.'s proposed architecture, particularly:

  1. The use of reflection padding in every Convolutional Layer, instead of big single reflection padding before the first convolution layer
  2. Ditching of the Tanh output. The generated image are the raw outputs of the convolutional layer. While the Tanh model produces visually pleasing results, the model fails to transfer the vibrant and loud colors of the style image (i.e. generated images are usually darker). This however makes for a good retro style effect.
  3. Use of Instance Normalization, instead of Batch Normalization after Convolutional and Deconvolutional layers, as discussed in Instance Normalization: The Missing Ingredient for Fast Stylization paper by Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky.

The original caffe pretrained weights of VGG16 were used for this implementation, instead of the pretrained VGG16's in PyTorch's model zoo.

Image Stylization

It took about 1.5 seconds for a GTX 1060 to stylize University of the Philippines Diliman - Oblation (1400×936) by LeAnne Jazul/Rappler. From Top to Right: Udnie Style, Mosaic Style, Tokyo Ghoul Style, Original Picture, Udnie Style with Original Color Preservation

Video Stylization

It took 6 minutes and 43 seconds to stylize a 2:11 minute-24 fps-1280x720 video on a GTX 1080 Ti.

More videos in this Youtube playlist. Unfortunately, Youtube's compression isn't friendly with style transfer videos, possibily because each frame is shaky with respect to its adjacent frames, hence obvious loss in video quality. Raw and lossless output video can be downloaded in my Dropbox folder, or Gdrive Folder

Webcam Demo

Webcam Demo

webcam.py can output 1280x720 videos at a rate of at least 4-5 frames per second on a GTX 1060.

Requirements

Most of the codes here assume that the user have access to CUDA capable GPU, at least a GTX 1050 ti or a GTX 1060

Data Files

Dependecies

Usage

All arguments, parameters and options are hardcoded inside these 5 python files. Before using the codes, please arrange your files and folders as defined below.

Training Style Transformation Network

train.py: trains the transformation network that learns the style of the style image. Each model in transforms folder was trained for roughly 23 minutes, with single pass (1 epoch) of 40,000 training images, and a batch size of 4, on a GTX 1080 Ti. 

python train.py

Options

transformer.py: contains the architecture definition of the trasnformation network. It includes 2 models, TransformerNetwork() and TransformerNetworkTanh(). TransformerNetwork doesn't have an extra output layer, while TransformerNetworkTanh, as the name implies, has for its output, a Tanh layer and a default output multiplier of 150. TransformerNetwork faithfully copies the style and colorization of the style image, while Tanh model produces images with darker color; which brings a retro style effect.

Options

experimental.py: contains the model definitions of the experimental transformer network architectures. These experimental transformer networks largely borrowed ideas from the papers Aggregated Residual Transformations for Deep Neural Networks or more commonly known as ResNeXt, and Densely Connected Convolutional Networks or more commonly known as DenseNet. These experimental networks are designed to be lightweight, with the goal of minimizing the compute and memory needed for better real-time performance.

See table below for the comparison of different transformer networks.

See transforms folder for some pretrained weights. For more pretrained weights, see my Gdrive or Dropbox.

Stylizing Images

stylize.py: Loads a pre-trained transformer network weight and applies style (1) to a content image or (2) to the images inside a folder

python stylize.py

Options

Stylizing Videos

video.py: Extracts all frames of a video, apply fast style transfer on each frames, and combine the styled frames into an output video. The output video doesn't retain the original audio. Optionally, you may use FFmpeg to merge the output video and the original video's audio.

python video.py

Options

Stylizing Webcam

webcam.py: Captures and saves webcam output image, perform style transfer, and again saves a styled image. Reads the styled image and show in window. 

python webcam.py

Options

Files and Folder Structure

master_folder
 ~ dataset 
    ~ train2014
        coco*.jpg
        ...
 ~ frames
    ~ content_folder
        frame*.jpg
        ...
 ~ images
    ~ out
        *.jpg
      *.jpg
 ~ models
    *.pth
 ~ style_frames
    frames*.jpg
 ~ transforms
    *.pth
 *.py

Comparison of Different Transformer Networks

Network size (Kb) no. of parameters final loss (million)
transformer/TransformerNetwork 6,573 1,679,235 9.88
experimental/TransformerNetworkDenseNet 1,064 269,731 11.37
experimental/TransformerNetworkUNetDenseNetResNet 1,062 269,536 12.32
experimental/TransformerNetworkV2 6,573 1,679,235 10.05
experimental/TransformerResNextNetwork 1,857 470,915 10.31
experimental/TransformerResNextNetwork_Pruned(0.3) 44 8,229 19.29
experimental/TransformerResNextNetwork_Pruned(1.0) 260 63,459 12.72

TransformerResNextNetwork and TransformerResNextNetwork_Pruned(1.0) provides the best tradeoff between compute, memory size, and performance.

Todo!

Citation

  @misc{rusty2018faststyletransfer,
    author = {Rusty Mina},
    title = {fast-neural-style: Fast Style Transfer in Pytorch!},
    year = {2018},
    howpublished = {\url{https://github.com/iamRusty/fast-neural-style-pytorch}},
    note = {commit xxxxxxx}
  }

Attribution

This implementation borrowed some implementation details from:

Related Work

License

Copyright (c) 2018 Rusty Mina. Free for academic or research use, as long as proper attribution is given and this copyright notice is retained.