/Style_Transfer_PyTorch

A PyTorch implement of the paper "Perceptual losses for real-time style transfer and super-resolution".

Primary LanguagePython

Fast Neural Style Transfer implement by PyTorch!

Author: Willie Chen

This repository contains a PyTorch implementation of an algorithm for fast style transfer. The algorithm can be used to mix the content of the input image with the style of the style image.

Table of Contents

Network structure overview

Network architecture detail: LINK

Network Overview

Image Stylization Result

Original Image (Source: https://www.pixiv.net/artworks/75323963)

Style Image Stylize image of the original image

Usage

Training the Style-Transfer Network

  1. Download the coco dataset
bash download_dataset.sh
  1. Train the Style Transfer model

train.py: Train the Transform Network that learn the style from the style_image and retain the semantic-information about the input_image.

python neural_style/train.py --content-dir ./images/content_images --style-img-path ./images/style_images/mosaic.jpg --epochs 1 --batch-size 4

Arguments (Optional)

usage: train_success.py [-h] [--epochs EPOCHS] [--lr LR]
                        [--batch-size BATCH_SIZE] [--img-size IMG_SIZE]
                        [--content-weight CONTENT_WEIGHT]
                        [--style-weight STYLE_WEIGHT]
                        [--save-interval SAVE_INTERVAL]
                        [--content-dir CONTENT_DIR]
                        [--style-img-path STYLE_IMG_PATH]
                        [--save-img-path SAVE_IMG_PATH]
                        [--save-model-path SAVE_MODEL_PATH] [--seed SEED]
                        [--gpu-id GPU_ID]
                        [--style-model-path STYLE_MODEL_PATH]

Style Transfer Project

optional arguments:
  -h, --help            show this help message and exit
  --epochs EPOCHS       Number of the training epochs
  --lr LR               Learning rate (default: 0.001)
  --batch-size BATCH_SIZE
                        Batch size of trainign and evaluation
  --img-size IMG_SIZE   Training image size
  --content-weight CONTENT_WEIGHT
                        Content weight for the final loss
  --style-weight STYLE_WEIGHT
                        Style weight for the final loss
  --save-interval SAVE_INTERVAL
                        Save model when every update save-interval times
  --content-dir CONTENT_DIR
                        Path for the content image root (default:
                        ./images/content_images
  --style-img-path STYLE_IMG_PATH
                        Path for the style image path (default:
                        ./images/style_images/starry-night-cropped.jpg
  --save-img-path SAVE_IMG_PATH
                        Path for the content image root (default:
                        ./images/result_images
  --save-model-path SAVE_MODEL_PATH
                        Path for the model weight (default: ../weights)
  --seed SEED           Set the random seed (default: 1)
  --gpu-id GPU_ID       Select the sepcific GPU card (default: 0)
  --style-model-path STYLE_MODEL_PATH
                        Specific the final file name of the model weight
                        (default: style_transform)

Neural Style Transfer

  1. Download the pretrained weight

Pretrained Weight: https://drive.google.com/drive/folders/1Iy-JGUA-KFjY0OgRmzhl2HQmaaXrHYjh?usp=sharing

  1. Run the fast style transfer
python neural_style/stylize_inference.py

Hint:

  • images/input_images: The picture you want to transfer style, please put it in this folder

TODOs

  • Video Style Transfer
  • Webcam

Reference

  1. J. Johnson, A. Alahi, and L. Fei-Fei. Perceptual losses for real-time style transfer and super-resolution. ECCV 2016
  2. J. Johnson, A. Alahi, and L. Fei-Fei. Perceptual Losses for Real-Time Style Transfer and Super-Resolution: Supplementary Material

Acknowledgement

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