Fast Style Transfer for Arbitrary Styles

View on Run in Google Colab View on GitHub Download notebook See TF Hub model

Based on the model code in magenta and the publication:

Exploring the structure of a real-time, arbitrary neural artistic stylization network. Golnaz Ghiasi, Honglak Lee, Manjunath Kudlur, Vincent Dumoulin, Jonathon Shlens, Proceedings of the British Machine Vision Conference (BMVC), 2017.


Let's start with importing TF-2 and all relevant dependencies.

import functools
import os

from matplotlib import gridspec
import matplotlib.pylab as plt
import numpy as np
import tensorflow as tf
import tensorflow_hub as hub

print("TF Version: ", tf.__version__)
print("TF-Hub version: ", hub.__version__)
print("Eager mode enabled: ", tf.executing_eagerly())
print("GPU available: ", tf.test.is_gpu_available())
TF Version:  2.3.1
TF-Hub version:  0.9.0
Eager mode enabled:  True
WARNING:tensorflow:From <ipython-input-1-1b40332be70f>:13: is_gpu_available (from tensorflow.python.framework.test_util) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.config.list_physical_devices('GPU')` instead.
GPU available:  True

# @title Define image loading and visualization functions  { display-mode: "form" }

def crop_center(image):
  """Returns a cropped square image."""
  shape = image.shape
  new_shape = min(shape[1], shape[2])
  offset_y = max(shape[1] - shape[2], 0) // 2
  offset_x = max(shape[2] - shape[1], 0) // 2
  image = tf.image.crop_to_bounding_box(
      image, offset_y, offset_x, new_shape, new_shape)
  return image

def load_image(image_url, image_size=(256, 256), preserve_aspect_ratio=True):
  """Loads and preprocesses images."""
  # Cache image file locally.
  image_path = tf.keras.utils.get_file(os.path.basename(image_url)[-128:], image_url)
  # Load and convert to float32 numpy array, add batch dimension, and normalize to range [0, 1].
  img = plt.imread(image_path).astype(np.float32)[np.newaxis, ...]
  if img.max() > 1.0:
    img = img / 255.
  if len(img.shape) == 3:
    img = tf.stack([img, img, img], axis=-1)
  img = crop_center(img)
  img = tf.image.resize(img, image_size, preserve_aspect_ratio=True)
  return img

def show_n(images, titles=('',)):
  n = len(images)
  image_sizes = [image.shape[1] for image in images]
  w = (image_sizes[0] * 6) // 320
  plt.figure(figsize=(w  * n, w))
  gs = gridspec.GridSpec(1, n, width_ratios=image_sizes)
  for i in range(n):
    plt.imshow(images[i][0], aspect='equal')
    plt.title(titles[i] if len(titles) > i else '')

Let's get as well some images to play with.

# @title Load example images  { display-mode: "form" }

content_image_url = ''  # @param {type:"string"}
style_image_url = ''  # @param {type:"string"}
output_image_size = 384  # @param {type:"integer"}

# The content image size can be arbitrary.
content_img_size = (output_image_size, output_image_size)
# The style prediction model was trained with image size 256 and it's the 
# recommended image size for the style image (though, other sizes work as 
# well but will lead to different results).
style_img_size = (256, 256)  # Recommended to keep it at 256.

content_image = load_image(content_image_url, content_img_size)
style_image = load_image(style_image_url, style_img_size)
style_image = tf.nn.avg_pool(style_image, ksize=[3,3], strides=[1,1], padding='SAME')
show_n([content_image, style_image], ['Content image', 'Style image'])
Downloading data from
65536/58102 [=================================] - 0s 0us/step
Downloading data from
2686976/2684586 [==============================] - 0s 0us/step


Import TF-Hub module

# Load TF-Hub module.

hub_handle = ''
hub_module = hub.load(hub_handle)

The signature of this hub module for image stylization is:

outputs = hub_module(content_image, style_image)
stylized_image = outputs[0]

Where content_image, style_image, and stylized_image are expected to be 4-D Tensors with shapes [batch_size, image_height, image_width, 3].

In the current example we provide only single images and therefore the batch dimension is 1, but one can use the same module to process more images at the same time.

The input and output values of the images should be in the range [0, 1].

The shapes of content and style image don't have to match. Output image shape is the same as the content image shape.

Demonstrate image stylization

# Stylize content image with given style image.
# This is pretty fast within a few milliseconds on a GPU.

outputs = hub_module(tf.constant(content_image), tf.constant(style_image))
stylized_image = outputs[0]
# Visualize input images and the generated stylized image.

show_n([content_image, style_image, stylized_image], titles=['Original content image', 'Style image', 'Stylized image'])


Let's try it on more images

# @title To Run: Load more images { display-mode: "form" }

content_urls = dict(
style_urls = dict(

content_image_size = 384
style_image_size = 256
content_images = {k: load_image(v, (content_image_size, content_image_size)) for k, v in content_urls.items()}
style_images = {k: load_image(v, (style_image_size, style_image_size)) for k, v in style_urls.items()}
style_images = {k: tf.nn.avg_pool(style_image, ksize=[3,3], strides=[1,1], padding='SAME') for k, style_image in style_images.items()}

Downloading data from
3178496/3170828 [==============================] - 0s 0us/step
Downloading data from
409600/406531 [==============================] - 0s 0us/step
Downloading data from
65536/61306 [================================] - 0s 0us/step
Downloading data from
196608/195196 [==============================] - 0s 0us/step
Downloading data from
46931968/46930988 [==============================] - 2s 0us/step
Downloading data from
401408/396423 [==============================] - 0s 0us/step
Downloading data from
147456/144340 [==============================] - 0s 0us/step
Downloading data from
11403264/11403121 [==============================] - 0s 0us/step
Downloading data from
2908160/2905099 [==============================] - 0s 0us/step
Downloading data from
1236992/1234199 [==============================] - 0s 0us/step
Downloading data from
122880/120288 [==============================] - 0s 0us/step
Downloading data from
139264/131604 [===============================] - 0s 0us/step
Downloading data from
32768/32390 [==============================] - 0s 0us/step
Downloading data from
1916928/1914618 [==============================] - 0s 0us/step
Downloading data from
40960/40620 [==============================] - 0s 0us/step
Downloading data from
73728/66306 [=================================] - 0s 0us/step

content_name = 'sea_turtle'  # @param ['sea_turtle', 'tuebingen', 'grace_hopper']
style_name = 'munch_scream'  # @param ['kanagawa_great_wave', 'kandinsky_composition_7', 'hubble_pillars_of_creation', 'van_gogh_starry_night', 'turner_nantes', 'munch_scream', 'picasso_demoiselles_avignon', 'picasso_violin', 'picasso_bottle_of_rum', 'fire', 'derkovits_woman_head', 'amadeo_style_life', 'derkovtis_talig', 'amadeo_cardoso']

stylized_image = hub_module(tf.constant(content_images[content_name]),

show_n([content_images[content_name], style_images[style_name], stylized_image],
       titles=['Original content image', 'Style image', 'Stylized image'])