Style Transfer

# The content and style inputs are generated deterministically in cell 7.

import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras.applications.vgg16 import preprocess_input
import IPython.display as display

img_size = 128
rng = np.random.default_rng(2026)
y, x = np.mgrid[0:img_size, 0:img_size]
circle = ((x - img_size * 0.5) ** 2 + (y - img_size * 0.5) ** 2) < (img_size * 0.3) ** 2
img_content = np.stack([0.15 + 0.75 * circle, 0.2 + 0.45 * x / img_size, 0.25 + 0.55 * y / img_size], axis=-1).astype("float32")
checker = ((x // 8 + y // 8) % 2).astype("float32")
waves = (np.sin(x / 5.0) + 1.0) / 2.0
img_style = np.stack([0.2 + 0.7 * checker, 0.15 + 0.7 * waves, 0.8 - 0.5 * checker], axis=-1).astype("float32")
img_result = rng.uniform(0, 1, size=(img_size, img_size, 3)).astype("float32")

matplotlib.rcParams["figure.figsize"] = (12, 4)
matplotlib.rcParams["axes.grid"] = False
fig, axes = plt.subplots(1, 3)
for axis, image, title in zip(axes, [img_content, img_style, img_result], ["Content", "Style", "Initial"]):
    axis.imshow(image)
    axis.set_title(title)
    axis.axis("off")
plt.show()

vgg = tf.keras.applications.VGG16(include_top=False, weights=None)
vgg.trainable = False

vgg.summary()

def layer_extractor(layers):
    outputs = [vgg.get_layer(name).output for name in layers]
    return tf.keras.Model(vgg.input, outputs)

def model_input(image):
    image = tf.cast(image, tf.float32)
    return preprocess_input(tf.expand_dims(image * 255.0, axis=0))

content_layers = ["block4_conv2"]
content_extractor = layer_extractor(content_layers)
content_target = tf.stop_gradient(content_extractor(model_input(img_content)))

def content_loss(image):
    current = content_extractor(model_input(image))
    return 0.5 * tf.reduce_sum(tf.square(current - content_target))

img = tf.Variable(img_result)
opt = tf.optimizers.Adam(learning_rate=0.002, beta_1=0.99, epsilon=1e-1)

clip = lambda value: tf.clip_by_value(value, clip_value_min=0.0, clip_value_max=1.0)

def optimize(image, loss_fn):
    with tf.GradientTape() as tape:
        loss = loss_fn(image)
    gradient = tape.gradient(loss, image)
    opt.apply_gradients([(gradient, image)])
    image.assign(clip(image))

def train(image, loss_fn, epochs=2, steps_per_epoch=10):
    for _ in range(epochs):
        for _ in range(steps_per_epoch):
            optimize(image, loss_fn=loss_fn)
    plt.imshow(clip(image).numpy())
    plt.axis("off")
    plt.show()

train(img, content_loss)

def gram_matrix(features):
    result = tf.linalg.einsum("bijc,bijd->bcd", features, features)
    shape = tf.shape(features)
    locations = tf.cast(shape[1] * shape[2], tf.float32)
    return result / locations

style_layers = ["block1_conv1", "block2_conv1", "block3_conv1", "block4_conv1"]

def style_extractor(image):
    return [gram_matrix(value) for value in layer_extractor(style_layers)(model_input(image))]

style_target = [tf.stop_gradient(value) for value in style_extractor(img_style)]

def style_loss(image):
    current = style_extractor(image)
    return tf.add_n([
        tf.reduce_mean(tf.square(value - target))
        for value, target in zip(current, style_target)
    ]) / len(style_layers)

def total_loss(img):
    return 2*content_loss(img)+style_loss(img)

img.assign(img_result)

train(img,loss_fn=total_loss)

def variation_loss(image):
    return tf.reduce_mean(tf.image.total_variation(tf.expand_dims(image, axis=0)))

def total_loss_var(image):
    return content_loss(image) + 150 * style_loss(image) + 30 * variation_loss(image)

start = np.clip(
    img_content + rng.normal(0.0, 0.03, size=img_content.shape),
    0.0,
    1.0,
).astype("float32")
img.assign(start)
train(img, loss_fn=total_loss_var)

keras.utils.save_img("result.png", tf.clip_by_value(img, 0.0, 1.0).numpy())