# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     language: python
#     name: python3
# ---

# +
import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import time
from helpers import load_and_process_image, display_image

from tensorflow.python.keras.applications.vgg19 import VGG19
from tensorflow.python.keras.preprocessing.image import load_img, img_to_array
from tensorflow.python.keras.models import Model

# %matplotlib inline
# -

# ## Import the model

model = VGG19(include_top=False, weights='imagenet')
model.trainable = False
model.summary()

# ## Display content and style images

content_path = 'images/louvre.jpg'
style_path = 'images/monet.jpg'

content = load_and_process_image(content_path)
style = load_and_process_image(style_path)

display_image(content, 'Content image')

display_image(style, 'Style image')

# ## Content and style models

content_layer = 'block5_conv2'
style_layers = [
    'block1_conv1',
    'block3_conv1',
    'block5_conv1'
]

content_model = Model(
    inputs=model.input,
    outputs=model.get_layer(content_layer).output
)
style_models = [Model(
    inputs=model.input,
    outputs=model.get_layer(l).output
) for l in style_layers]


# ## Content cost

def content_cost(content, generated):
    a_C = content_model(content)
    a_G = content_model(generated)
    
    return tf.reduce_mean(tf.square(a_C - a_G))


# ## Style cost

def gram_matrix(A):
    n_C = int(A.shape[-1]) # number of channels
    # The result will have a shape of (sth, n_C) => for an image it will 
    # be (height x width, num of channels)
    a = tf.reshape(A, [-1, n_C])
    n = tf.shape(a)[0]
    G = tf.matmul(a, a, transpose_a=True)
    
    return G / tf.cast(n, tf.float32)


# Weight of each style layer to compute the style cost.
# We assume that each layer has the same weight
weight = 1 / len(style_models)


def style_cost(style, generated):
    J_style = 0
    
    for style_model in style_models:
        a_S = style_model(style)
        a_G = style_model(generated)
        GS = gram_matrix(a_S)
        GG = gram_matrix(a_G)
        current_cost = tf.reduce_mean(tf.square(GS - GG))
        J_style += current_cost * weight
        
    return J_style


# ## Train the model

# +
generated_images = []

def training_loop(iterations=20, alpha=10, beta=20, learning_rate=7):
    generated = tf.Variable(content, dtype=tf.float32)
    opt = tf.keras.optimizers.Adam(learning_rate)
    
    best_cost = 1e12
    best_image = None
    start_time = time.time()
    
    for i in range(iterations):
        with tf.GradientTape() as tape:
            J_content = content_cost(content, generated)
            J_style = style_cost(style, generated)
            J_total = alpha * J_content + beta * J_style
            
        grads = tape.gradient(J_total, generated)
        opt.apply_gradients([(grads, generated)])
        
        if J_total < best_cost:
            best_cost = J_total
            best_image = generated
            
        print('Cost at {}: {}. Time elapsed: {}'.format(i, J_total, time.time() - start_time))
        generated_images.append(generated.numpy())
        
    return best_image


# -

best_image = training_loop(iterations=100)

# ## Plot the results

display_image(best_image)

# +
plt.figure(figsize=(10,10))

for i in range(20):
    plt.subplot(5, 4, i+1)
    display_image(generated_images[i])
# -



StarcoderdataJupyterScriptsDedupFiltered