Skip to content

Implementation of the fast neural style transfer algorithm on Keras. Includes Jupyter notebooks, python script and web app.

License

Notifications You must be signed in to change notification settings

Gooogr/Fast_Neural_Style_Transfer

Repository files navigation

Fast Neural Style Transfer project

This project is my attempt to solve the problem of quick style transfer by myself. While creating it, I moved step by step, gradually studying new papers.
Jupyter notebooks for each step you can find in the research_notebook page.
Also you can find all notebooks in my Google Collab project folder. The source of the development was there.

Original article: "Perceptual Losses for Real-Time Style Transfer and Super-Resolution"
Supplementary Material: "Perceptual Losses for Real-Time Style Transfer and Super-Resolution: Supplementary Material"

Result

Environment (main packages):

  • Python==3.6
  • Keras==2.2.5
  • opencv-python==4.1.2.30
  • Pillow
  • tensorflow==1.15.2
  • streamlit

How to use it

There are two ways to use this neural network. You can train or predict in this google collab file or use a python script.

If you want to train your version of neural network, you have to download train part of the COCO 2014 dataset.

For training the network:

python main.py train --conf [path to config json file, optional]

For predicting:

python main.py predict -w [path to the pre-trained weights] -i [path to the content image] -r [directory where to save result, optional]

Explanation some of the config.json parametres:

  • "net_name" - It will be the name of saved weights after training
  • "height", "weight" - Size for re-shaped images during the training. It were 256 x 256 in the original article, but I used 512x512 for better results
  • "verbose_iter" - Determine how often you will print training info and save test image

Default files structure

project
│   README.md
|   LICENSE
│   main.py
|   model_functions.py
|   model_zoo.py  
|   utilities.py
│
└───img
│   │   style_img.jpg
│   │   content_img.jpg
|   |   test_content_img.jpg
│   │
│   └───iteration_results
│   
└───dataset
|   |
|   └───train2014
|       |   # 82 000 images
|
└───saved_weights
    |   fst_draft_512_weights.h5
    |   fst_kandinskiy_512_weights.h5
    |   fst_night_256_weights.h5
    |   fst_night_512_weights.h5

Perfomance

It takes about 10 hours to train a network on a Nvidia K80 GPU (Google Collab).

Web app

I have also added simple web app with Streamlit. You can upload your content image and select one of three style images.
For running locally:

streamlit run streamlit_app.py

License

MIT © Grigoriy Gusarov

About

Implementation of the fast neural style transfer algorithm on Keras. Includes Jupyter notebooks, python script and web app.

Topics

Resources

License

Stars

Watchers

Forks