Implementación de neural style transfer con TensorFlow 2.0

Las redes neuronales convolucionales son el tipo de red más potente para la clasificación y el análisis de imágenes. Sus aplicaciones han superado muchas barreras y se han consolidado como el elemento clave de muchas soluciones de deep learning que vemos hoy. A muy alto nivel, las CNN pueden aprender representaciones internas a nivel de características de las imágenes que reciben como entrada. Esto es lo que las hace tan potentes. Y resulta que esta capacidad no solo sirve para tareas como la clasificación de imágenes, sino también para su construcción. Aplicaciones como Deep Dream y Neural Style Transfer componen imágenes a partir de las activaciones de capas dentro de CNN y de las características que extraen.

En este tutorial, verás cómo funciona Neural Style Transfer y cómo implementarlo con TensorFlow 2.0.

Nota: para seguir el tutorial, necesitas entender cómo funcionan las CNN. Si buscas recursos para refrescar conceptos rápidamente, prueba con los siguientes:

• El curso de DataCamp Convolutional Neural Networks for Image Processing
• CS231n: Convolutional Neural Networks for Visual Recognition

¿Qué es Neural Style Transfer?

Para quienes se inician en visión por computador y deep learning, el término neural style transfer puede imponer un poco. Para entender cada componente del término, fíjate en estas dos imágenes:

En este contexto, a la imagen de la izquierda la llamamos imagen de contenido y a la de la derecha, imagen de estilo. Lo que te interesa es aplicar el estilo de una imagen (la derecha) sobre otra (la izquierda). Eso explica las dos últimas palabras del término: style transfer o transferencia de estilo. Para llevar a cabo el proceso, se entrena una red neuronal (CNN) para optimizar una función de pérdida personalizada; de ahí la primera palabra: neural. Al fusionar las dos imágenes anteriores con neural style transfer, el resultado final sería este (derecha):

¿Para qué sirve neural style transfer? Imagina que tienes una obra original de Vincent van Gogh y quieres ver cómo se vería ese estilo aplicado a uno de tus dibujos. Ahí es donde la transferencia de estilo tiene sentido. Otro ejemplo son varias apps de filtros fotográficos como Prisma, que te permiten aplicar neural style transfer con una interfaz muy sencilla.

Separación entre contenido y estilo

Hasta ahora, la idea clave es que el proceso consiste en combinar el contenido de una imagen con el estilo de otra. Y es exactamente así. En esta sección, veamos qué significan contenido y estilo en el contexto de las CNN.

Una CNN suele estar compuesta por varias capas de convolución y de pooling. Las capas convolucionales extraen características complejas de una imagen, mientras que las capas de pooling descartan información espacial detallada que no es relevante para un problema de clasificación. El efecto es que la CNN aprende el contenido de una imagen más que aspectos específicos como color, textura, etc. A medida que profundizamos en la red, aumenta la complejidad de las características y las capas convolucionales más profundas suelen llamarse representaciones de contenido.

Por estilo entendemos cualquier rasgo específico de las propiedades de la imagen como la textura, el color, etc., por ejemplo, el uso predominante de un color en un cuadro. ¿Cómo se extrae el estilo de una imagen? Midiendo correlaciones entre las capas convolucionales. Las correlaciones nos indican hasta qué punto dos o más variables son similares o están relacionadas. Para entenderlo, piensa en una capa convolucional aprendida compuesta por varios mapas de características. Para cada mapa puedes medir cuánto se relacionan sus rasgos detectados con los de otros mapas de la misma capa. Así estimas cosas como: ¿un cierto color detectado en el primer mapa se parece a un color de otro mapa? ¿Hay formas comunes entre los mapas? Estos rasgos/similitudes definen el estilo de una imagen. Medir la similitud entre el contenido de varios mapas de características dentro de una capa ayuda a las redes a aprender una representación multiescala de la imagen centrada en rasgos espaciales como textura y color.

Un punto clave: al aplicar neural style transfer, también debes asegurarte de preservar el contenido de la imagen junto con el estilo deseado de la otra. Verás cómo hacerlo más adelante. En la próxima sección, aprenderás a extraer las características de una imagen (contenido) y a calcular la pérdida de contenido.

Cuantificar la imagen de contenido y calcular la pérdida de contenido

El algoritmo de neural style transfer fue presentado por Gatys et al. en 2015 en su artículo A Neural Algorithm of Artistic Style. Este tutorial, no obstante, toma como referencia Image Style Transfer Using Convolutional Neural Networks, una continuación del anterior.

Según Image Style Transfer Using Convolutional Neural Networks, se emplea la arquitectura VGG-19 para extraer tanto las características de contenido como de estilo de las imágenes correspondientes. Para el contenido, se usa la segunda capa de convolución del cuarto bloque (de capas convolucionales), denominada conv4_2. Una vez tienes las características de contenido, debes compararlas con una imagen objetivo para medir la pérdida de contenido. ¿Qué es una imagen objetivo? ¿Por qué hace falta calcular esa pérdida? Demos un paso atrás y respondamos.

Para combinar contenido y estilo en una sola imagen, necesitas partir de una imagen objetivo, que puede ser un lienzo en blanco o una copia de la imagen de contenido. Para que una CNN aprenda eficazmente contenido y estilo, necesitas una función de pérdida personalizada que optimices para obtener una imagen estilizada y suave a partir de las imágenes de contenido y de estilo. Esta función de pérdida es, esencialmente, la suma de dos pérdidas distintas:

• Pérdida de contenido, que garantiza que se preserve el contenido global.
• Pérdida de estilo, que controla la cantidad de estilo que se transfiere a la imagen objetivo.

Volvamos ahora a la pérdida de contenido y definámosla.

donde $T_c$ es la imagen objetivo y $C_c$ la imagen de contenido. Esta función te da una medida de cuán alejadas están las características de las imágenes de contenido y objetivo. La red intentará minimizar esta pérdida. En la siguiente sección verás cómo se determina la pérdida de estilo.

La pérdida de estilo y la matriz de Gram

Para calcular la pérdida de estilo, el artículo indica que tomes las representaciones (números, básicamente) de las siguientes capas y obtengas las matrices de Gram de los mapas de características dentro de esas capas.

'conv1_1'
'conv2_1'
'conv3_1'
'conv4_1'
'conv5_1'

Usar varias capas para definir la pérdida de estilo permite aprender una representación multiescala del estilo presente en una imagen. Las matrices de Gram te ayudan a determinar la similitud entre las características de los mapas dentro de una misma capa convolucional. Esto captura información no localizada sobre la imagen que moldea su estilo. Una matriz de Gram en este contexto se calcula así.

Supón que la dimensión de una imagen convolucionada es 8x8x16, lo que indica que tiene 16 mapas de características. Quieres encontrar similitudes entre las características de esos mapas mediante una matriz de Gram. Para ello, aplanas las dos primeras dimensiones de la imagen convolucionada y las conviertes en un vector 1D. En este caso, el aplanado contendría 64 entradas en el vector 1D. Repites esto para cada mapa de características. Obtienes así una matriz final de dimensión 16x64 (llamémosla A). Transpones esa matriz (pasa a 64x16) y la multiplicas por A. La matriz de Gram resultante tendrá dimensión 16x16. Cada valor de la matriz de Gram indica la similitud entre mapas de características. Visualmente, el proceso se parece a esto:

Para calcular la pérdida de estilo entre la imagen objetivo y la de estilo, tomas las distancias cuadráticas medias entre sus matrices de Gram en cada uno de los bloques de capas convolucionales. Formalmente, se define así:

donde T_(s,i) es la matriz de Gram de la imagen objetivo calculada en el bloque i y S_(s,i) es la matriz de Gram de la imagen de estilo en el mismo bloque. Con w_i puedes asignar pesos personalizados a los diferentes bloques para conseguir una representación más afinada del estilo. Por último, $a$ es una constante que normaliza los valores de cada capa dentro de los bloques. Juntando ambas pérdidas, definimos la pérdida total que la red optimiza durante la transferencia de estilo.

Definir la pérdida total

La pérdida total es simplemente una suma ponderada:

donde $\alpha$ es el peso del contenido y $\beta$ el del estilo. Esto te ayuda a mantener un buen equilibrio entre el contenido y la cantidad de estilo que se transfiere a la imagen objetivo. En el artículo puedes ver el efecto de distintas combinaciones de $\alpha$ y $\beta$; te animamos a explorarlo. Ya tienes todas las piezas para implementar neural style transfer. Vamos a ello en la siguiente sección.

Neural Style Transfer en acción

Para la implementación usarás TensorFlow 2.0. Incluye muchas novedades y es una de las actualizaciones más completas de la biblioteca hasta la fecha. Si quieres conocer algunas de ellas, echa un vistazo a este artículo.

La versión más reciente de TensorFlow (en el momento de escribir este tutorial) es 2.0.0-beta0. Si aún no la tienes, instálala siguiendo las instrucciones aquí. Empezarás importando los paquetes necesarios y cargando las imágenes de contenido y estilo.

# Packages
import tensorflow as tf
from tensorflow.keras.applications.vgg19 import preprocess_input
from tensorflow.keras.models import Model
import matplotlib.pyplot as plt
import numpy as np
np.random.seed(7)

%matplotlib inline

print(tf.__version__)

2.0.0-beta0

# Load the content and style images
content = plt.imread('Content.jpeg')
style = plt.imread('Style.jpg')

# Display the images
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 10))

# Content and style images side-by-side
ax1.imshow(content)
ax1.set_title('Content Image')
ax2.imshow(style)
ax2.set_title('Style Image')
plt.show()

En el siguiente paso, escribirás una función auxiliar que cargue ambas imágenes en arrays numéricos (los ordenadores solo entienden números) y las reajuste para que sean compatibles con el modelo.

def load_image(image):
  image = plt.imread(image)
  img = tf.image.convert_image_dtype(image, tf.float32)
  img = tf.image.resize(img, [400, 400])
  # Shape -> (batch_size, h, w, d)
  img = img[tf.newaxis, :]
  return img

# Use load_image of content and style images
content = load_image('Content.jpeg')
style = load_image('Style.jpg')

# Verify the shapes
content.shape, style.shape

(TensorShape([1, 400, 400, 3]), TensorShape([1, 400, 400, 3]))

Las imágenes ya están redimensionadas. Ahora cargarás un modelo VGG19 preentrenado para extraer las características. Como lo usarás solo para extracción, no necesitas la parte del clasificador.

vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')
vgg.trainable = False

Downloading data from https://github.com/fchollet/deep-learning-models/releases/download/v0.1/vgg19_weights_tf_dim_ordering_tf_kernels_notop.h5
80142336/80134624 [==============================] - 5s 0us/step

Los autores especificaron los bloques usados para extraer contenido y estilo. El VGG19 que acabas de cargar usa la misma convención de nombres. Para aprovecharlo, primero imprime los nombres de todas las capas de la red.

# Print the layer names for convenience
for layer in vgg.layers:
  print(layer.name)

input_1
block1_conv1
block1_conv2
block1_pool
block2_conv1
block2_conv2
block2_pool
block3_conv1
block3_conv2
block3_conv3
block3_conv4
block3_pool
block4_conv1
block4_conv2
block4_conv3
block4_conv4
block4_pool
block5_conv1
block5_conv2
block5_conv3
block5_conv4
block5_pool

Te interesan las siguientes para obtener las características de estilo:

'conv1_1'
'conv2_1'
'conv3_1'
'conv4_1'
'conv5_1'

Para el contenido, necesitarás conv4_2. Guárdalas en variables.

# Content layer
content_layers = ['block4_conv2']

# Style layer
style_layers = ['block1_conv1',
                'block2_conv1',
                'block3_conv1',
                'block4_conv1',
                'block5_conv1']


num_content_layers = len(content_layers)
num_style_layers = len(style_layers)

Crea ahora un modelo VGG personalizado compuesto por esas capas. Así podrás hacer pases hacia delante sobre las imágenes y extraer las características necesarias.

def mini_model(layer_names, model):

  outputs = [model.get_layer(name).output for name in layer_names]

  model = Model([vgg.input], outputs)
  return model

Definir una matriz de Gram es muy sencillo en TensorFlow; puedes hacerlo así:

# Gram matrix
def gram_matrix(tensor):
  temp = tensor
  temp = tf.squeeze(temp)
  fun = tf.reshape(temp, [temp.shape[2], temp.shape[0]*temp.shape[1]])
  result = tf.matmul(temp, temp, transpose_b=True)
  gram = tf.expand_dims(result, axis=0)
  return gram

Ahora definirás un modelo personalizado con la función mini_model(). Servirá para devolver las características de contenido y de estilo de las imágenes correspondientes.

 class Custom_Style_Model(tf.keras.models.Model):
  def __init__(self, style_layers, content_layers):
    super(Custom_Style_Model, self).__init__()
    self.vgg =  mini_model(style_layers + content_layers, vgg)
    self.style_layers = style_layers
    self.content_layers = content_layers
    self.num_style_layers = len(style_layers)
    self.vgg.trainable = False

  def call(self, inputs):
    # Scale back the pixel values
    inputs = inputs*255.0
    # Preprocess them with respect to VGG19 stats
    preprocessed_input = preprocess_input(inputs)
    # Pass through the mini network
    outputs = self.vgg(preprocessed_input)
    # Segregate the style and content representations
    style_outputs, content_outputs = (outputs[:self.num_style_layers],
                                      outputs[self.num_style_layers:])

    # Calculate the gram matrix for each layer
    style_outputs = [gram_matrix(style_output)
                     for style_output in style_outputs]

    # Assign the content representation and gram matrix in
    # a layer by layer fashion in dicts
    content_dict = {content_name:value
                    for content_name, value
                    in zip(self.content_layers, content_outputs)}

    style_dict = {style_name:value
                  for style_name, value
                  in zip(self.style_layers, style_outputs)}

    return {'content':content_dict, 'style':style_dict}

Como ya tienes el modelo personalizado, úsalo sobre las imágenes para obtener las características de contenido y estilo:

# Note that the content and style images are loaded in
# content and style variables respectively
extractor = Custom_Style_Model(style_layers, content_layers)
style_targets = extractor(style)['style']
content_targets = extractor(content)['content']


In the paper, optimization was done using the L-BFGS algorithm, but you can use Adam also.

opt = tf.optimizers.Adam(learning_rate=0.02)

Definamos ahora los pesos globales de contenido y estilo, y también los pesos de cada representación de estilo como vimos antes. Nota: son hiperparámetros; te conviene experimentar con ellos.

 # Custom weights for style and content updates
style_weight=100
content_weight=10

# Custom weights for different style layers
style_weights = {'block1_conv1': 1.,
                 'block2_conv1': 0.8,
                 'block3_conv1': 0.5,
                 'block4_conv1': 0.3,
                 'block5_conv1': 0.1}

Llega la parte crucial que hace que neural style transfer sea tan interesante: la función de pérdida.

# The loss function to optimize
def total_loss(outputs):
    style_outputs = outputs['style']
    content_outputs = outputs['content']
    style_loss = tf.add_n([style_weights[name]*tf.reduce_mean((style_outputs[name]-style_targets[name])**2)
                           for name in style_outputs.keys()])
    # Normalize
    style_loss *= style_weight / num_style_layers

    content_loss = tf.add_n([tf.reduce_mean((content_outputs[name]-content_targets[name])**2)
                             for name in content_outputs.keys()])
    # Normalize
    content_loss *= content_weight / num_content_layers
    loss = style_loss + content_loss
    return loss

Ahora escribirás otra función que:

• Calcule los gradientes de la función de pérdida que acabas de definir.
• Use esos gradientes para actualizar la imagen objetivo.

Con GradientTape puedes aprovechar la diferenciación automática, que calcula los gradientes de una función en función de su composición. También usarás el decorador tf.function para acelerar las operaciones. Lee más aquí.

@tf.function()
def train_step(image):
  with tf.GradientTape() as tape:
    # Extract the features
    outputs = extractor(image)
    # Calculate the loss
    loss = total_loss(outputs)
  # Determine the gradients of the loss function w.r.t the image pixels
  grad = tape.gradient(loss, image)
  # Update the pixels
  opt.apply_gradients([(grad, image)])
  # Clip the pixel values that fall outside the range of [0,1]
  image.assign(tf.clip_by_value(image, clip_value_min=0.0, clip_value_max=1.0))

El último paso antes de entrenar es definir la imagen objetivo. Usaremos la imagen de contenido.

target_image = tf.Variable(content)

Ya estás listo para entrenar la red.

epochs = 10
steps_per_epoch = 100

step = 0
for n in range(epochs):
  for m in range(steps_per_epoch):
    step += 1
    train_step(target_image)
  plt.imshow(np.squeeze(target_image.read_value(), 0))
  plt.title("Train step: {}".format(step))
  plt.show()

WARNING: Logging before flag parsing goes to stderr.
W0617 16:21:34.491543 140709216896896 deprecation.py:323] From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/math_grad.py:1205: add_dispatch_support.<locals>.wrapper (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.where in 2.0, which has the same broadcast rule as np.where 

Te recomendamos jugar con los pesos de contenido y estilo para observar los cambios en la imagen objetivo.

Conclusión

El artículo de referencia del tutorial no menciona cómo optimizar más la calidad de la imagen generada. Como habrás visto en los resultados, sería deseable cierta suavidad adicional en la relación espacial entre el contenido y el estilo. Para abordar este problema se introdujo la total_variational_loss, análoga a una regularización. En términos generales, total_variational_loss penaliza los artefactos de alta frecuencia introducidos por el algoritmo original de neural style transfer. Consulta este tutorial si te interesa implementarla. A continuación tienes algunos recursos que se usaron como referencia para escribir este tutorial:

• Hands-On Transfer Learning with Python
• Neural style transfer
• Lección sobre Style Transfer de Udacity
• Deep Learning for Computer Vision de Adrian Rosebrock (Parte II)

Si quieres afinar tus habilidades de deep learning, echa un vistazo a estos cursos y tutoriales de DataCamp:

• Deep Learning in Python
• Advanced Deep Learning with Keras in Python
• Convolutional Neural Networks with TensorFlow