fast-neural-style (expanded with ONNX-Caffe2 integration) 🌇 🚀

This is an example of how abhiskk/fast-neural-style can be expanded to export trained model to ONNX and later consume in any of the supported backends.

As of today, you'd need to install onnx, onnx-caffe2 and Caffe2 from source.

Exporting the model is the matter of add --export_model <filename.onnx>

python neural_style/neural_style.py eval --content-image images/content-images/amber.jpg --model saved-models/starry-night.pth --output-image styled.jpg --cuda 0 --content-scale 2 --export_onnx starry-night.onnx

After that eval can be invoked with .onnx file directly and would invoke Caffe2 backend (see [neural_style/neural_style.py](neural_style/neural_style.py for how it's invoked):

python neural_style/neural_style.py eval --content-image images/content-images/amber.jpg --model starry-night.onnx --output-image styled.jpg --cuda 1 --content-scale 2

Sample converted ONNX model is also checked into the repo.

Original description

This repository contains a pytorch implementation of an algorithm for artistic style transfer. The algorithm can be used to mix the content of an image with the style of another image. For example, here is a photograph of a door arch rendered in the style of a stained glass painting.

The model uses the method described in Perceptual Losses for Real-Time Style Transfer and Super-Resolution along with Instance Normalization. The saved-models for examples shown in the README can be downloaded from here.

Requirements

The program is written in Python, and uses pytorch, scipy. A GPU is not necessary, but can provide a significant speed up especially for training a new model. Regular sized images can be styled on a laptop, desktop using saved models.

Usage

Stylize image

python neural_style/neural_style.py eval --content-image </path/to/content/image> --model </path/to/saved/model> --output-image </path/to/output/image> --cuda 0

• --content-image: path to content image you want to stylize.
• --model: saved model to be used for stylizing the image (eg: mosaic.pth)
• --output-image: path for saving the output image.
• --content-scale: factor for scaling down the content image if memory is an issue (eg: value of 2 will halve the height and width of content-image)
• --cuda: set it to 1 for running on GPU, 0 for CPU.

Train model

python neural_style/neural_style.py train --dataset </path/to/train-dataset> --style-image </path/to/style/image> --vgg-model-dir </path/to/vgg/folder> --save-model-dir </path/to/save-model/folder> --epochs 2 --cuda 1

There are several command line arguments, the important ones are listed below

• --dataset: path to training dataset, the path should point to a folder containing another folder with all the training images. I used COCO 2014 Training images dataset [80K/13GB] (download).
• --style-image: path to style-image.
• --vgg-model-dir: path to folder where the vgg model will be downloaded.
• --save-model-dir: path to folder where trained model will be saved.
• --cuda: set it to 1 for running on GPU, 0 for CPU.

Refer to neural_style/neural_style.py for other command line arguments.

DISCLAIMER: This implementation is also a part of the pytorch examples repository. Implementation in this repository uses pretrained Caffe2 VGG whereas the pytorch examples repository implementation uses pretrained Pytorch VGG. The two VGGs have different preprocessings which results in different --content-weight and --style-weight parameters. The styled output images also look slightly different.

Models

Models for the examples shown below can be downloaded from here or by running the script download_styling_models.sh.