Counting 1,563 Big Data & Machine Learning Frameworks, Toolsets, and Examples...
Suggestion? Feedback? Tweet @stkim1

PyTorch-Style-Transfer

This repo provides PyTorh Implementation of MSG-Net (ours) and Neural Style (Gatys et al. CVPR 2016). Please install PyTorch with cuda support before running the code. The code has been tested on Ubuntu 16.04 with Titan X Pascal and Maxwell. If you would like to produce exactly the same result as in the MSG-Net paper, please visit our original Torch implementation.

Tabe of content

MSG-Net

Multi-style Generative Network for Real-time Transfer [arXiv] [project]
Hang Zhang, Kristin Dana
@article{zhang2017multistyle,
	title={Multi-style Generative Network for Real-time Transfer},
	author={Zhang, Hang and Dana, Kristin},
	journal={arXiv preprint arXiv:1703.06953},
	year={2017}
}

Stylize Images Using Pre-trained MSG-Net

  1. Download the pre-trained model
    git clone [email protected]:zhanghang1989/PyTorch-Style-Transfer.git
    cd PyTorch-Style-Transfer/experiments
    bash models/download_model.sh
  2. Test the model
    python main.py eval --content-image images/content/venice-boat.jpg --style-image images/9styles/candy.jpg --model models/9styles.model
  • If you don't have a GPU, simply set --cuda=0. For a different style, set --style-image path/to/style. If you would to stylize your own photo, change the --content-image path/to/your/photo. More options:

    • --content-image: path to content image you want to stylize.
    • --style-image: path to style image (typically covered during the training).
    • --model: path to the pre-trained model to be used for stylizing the image.
    • --output-image: path for saving the output image.
    • --content-size: the content image size to test on.
    • --cuda: set it to 1 for running on GPU, 0 for CPU.

Train Your Own MSG-Net Model

  1. Download the COCO dataset
    bash dataset/download_dataset.sh
  2. Train the model
    python main.py train --epochs 4
  • If you would like to customize styles, set --style-folder path/to/your/styles. More options:
    • --style-folder: path to the folder style images.
    • --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.

Neural Style

Image Style Transfer Using Convolutional Neural Networks by Leon A. Gatys, Alexander S. Ecker, and Matthias Bethge.

python main.py optim --content-image images/content/venice-boat.jpg --style-image images/9styles/candy.jpg
  • --content-image: path to content image.
  • --style-image: path to style image.
  • --output-image: path for saving the output image.
  • --content-size: the content image size to test on.
  • --style-size: the style image size to test on.
  • --cuda: set it to 1 for running on GPU, 0 for CPU.

Extending the Software

  • Extending a new network type:
    • define your own file under the folder net/
    • implement your own nn.Module in mynn.py if need
    • implement train and evaluate functions base on your need
    • set up the network as in main.py
  • Extending a new experiment (requires differet setting)
    • define the subcommand of the options in option.py
    • implement the experiment function like optimize()
    • set up the experiment as in main.py

Acknowledgement

The code benefits from outstanding prior work and their implementations including: