style_transfer.ipynb

Kernel: Python 3

Copyright 2018 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Transferencia de estilo neuronal

Ver en TensorFlow.org

Ejecutar en Google Colab

Ver en GitHub

Descargar el bloc de notas

Ver modelo de TF Hub

En este tutorial se usa aprendizaje profundo para componer una imagen en el estilo de otra imagen (¿le gustaría pintar como Picasso o Van Gogh?). Esto se conoce como transferencia de estilo neuronal y esta técnica se explica en Un algoritmo de estilo artístico (Gatys et al.).

Nota: En este tutorial se muestra el algoritmo original de transferencia de estilo. Optimiza el contenido de la imagen a un estilo específico. Los enfoques modernos entrenan un modelo para generar imágenes estilizadas directamente (similar a CycleGAN). Este enfoque es mucho más rápido (hasta 1000x).

Para ver una aplicación simple de la transferencia de estilo con un modelo preentrenado de TensorFlow Hub, vea el tutorial Fast style transfer for arbitrary styles que usa un modelo de estilización de imagen. Pare ver un ejemplo de transferencia de estilo con TensorFlow Lite, vea Artistic style transfer with TensorFlow Lite.

La transferencia neuronal de estilo es una técnica de optimización que se usa para tomar dos imágenes, una imagen de contenido y una imagen de referencia de estilo (como la pintura de un artista famoso), y las combina para que la imagen de salida se vea como la imagen de contenido pero "pintada" en el estilo de la imagen de referencia.

La forma de implementarlo es optimizando la imagen de salida para que coincida con las estadísticas de la imagen de contenido y las estadísticas de la imagen de referencia. Estas estadísticas se extraen de las imágenes mediante la red neuronal convolucional.

Por ejemplo, usemos la imagen de este perro y la Composition 7 de Wassily Kandinsky.

Yellow Labrador Looking, de Wikimedia Commons de Elf. Licencia CC BY-SA 3.0

Entonces, ¿cómo se vería si Kandisky pintara la imagen de este perro en este estilo exclusivamente? ¿Algo así?

Preparación

Importar y configurar módulos

import os
import tensorflow as tf
# Load compressed models from tensorflow_hub
os.environ['TFHUB_MODEL_LOAD_FORMAT'] = 'COMPRESSED'

import IPython.display as display

import matplotlib.pyplot as plt
import matplotlib as mpl
mpl.rcParams['figure.figsize'] = (12, 12)
mpl.rcParams['axes.grid'] = False

import numpy as np
import PIL.Image
import time
import functools

def tensor_to_image(tensor):
  tensor = tensor*255
  tensor = np.array(tensor, dtype=np.uint8)
  if np.ndim(tensor)>3:
    assert tensor.shape[0] == 1
    tensor = tensor[0]
  return PIL.Image.fromarray(tensor)

Descarga imágenes y selecciona una imagen de estilo y una imagen de contenido:

content_path = tf.keras.utils.get_file('YellowLabradorLooking_new.jpg', 'https://storage.googleapis.com/download.tensorflow.org/example_images/YellowLabradorLooking_new.jpg')
style_path = tf.keras.utils.get_file('kandinsky5.jpg','https://storage.googleapis.com/download.tensorflow.org/example_images/Vassily_Kandinsky%2C_1913_-_Composition_7.jpg')

Visualizar la entrada

Defina una función para cargar la imagen y delimite la dimensión máxima a 512 píxeles.

def load_img(path_to_img):
  max_dim = 512
  img = tf.io.read_file(path_to_img)
  img = tf.image.decode_image(img, channels=3)
  img = tf.image.convert_image_dtype(img, tf.float32)

  shape = tf.cast(tf.shape(img)[:-1], tf.float32)
  long_dim = max(shape)
  scale = max_dim / long_dim

  new_shape = tf.cast(shape * scale, tf.int32)

  img = tf.image.resize(img, new_shape)
  img = img[tf.newaxis, :]
  return img

Cree una función simple para que se muestre la imagen.

def imshow(image, title=None):
  if len(image.shape) > 3:
    image = tf.squeeze(image, axis=0)

  plt.imshow(image)
  if title:
    plt.title(title)

content_image = load_img(content_path)
style_image = load_img(style_path)

plt.subplot(1, 2, 1)
imshow(content_image, 'Content Image')

plt.subplot(1, 2, 2)
imshow(style_image, 'Style Image')

Transferencia rápida de estilo con TF-Hub

En este tutorial se demuestra el algoritmo original de transferencia de estilo, que optimiza la imagen de contenido en un estilo en particular. Antes de profundizar en los detalles, veamos como lo hace el modelo de TensorFlow Hub:

import tensorflow_hub as hub
hub_model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
stylized_image = hub_model(tf.constant(content_image), tf.constant(style_image))[0]
tensor_to_image(stylized_image)

Definir las representaciones de contenido y estilo

Use las capas intermedias del modelo para obtener las representaciones de contenido y estilo de la imagen. Si empezamos por la capa de entrada de la red, las primeras activaciones de la capa representan características de poca importancia, como bordes y texturas. A medida que se avanza en la red, las últimas capas representan características de mayor importancia (partes de los objetos como ruedas u ojos). En este caso, está usando una arquitectura de red VGG19, una red preentrenada de clasificación de imágenes. Las capas intermedias son necesarias para definir la representación de contenido y estilo de las imágenes. Para obtener una imagen de salida, intente hacer coincidir el estilo las representaciones del estilo correspondiente y el contenido de destino en las capas intermedias.

Cargue una VGG19 y ejecútela en modo de prueba con nuestra imagen para asegurarse de que se pueda usar correctamente:

x = tf.keras.applications.vgg19.preprocess_input(content_image*255)
x = tf.image.resize(x, (224, 224))
vgg = tf.keras.applications.VGG19(include_top=True, weights='imagenet')
prediction_probabilities = vgg(x)
prediction_probabilities.shape

predicted_top_5 = tf.keras.applications.vgg19.decode_predictions(prediction_probabilities.numpy())[0]
[(class_name, prob) for (number, class_name, prob) in predicted_top_5]

Ahora cargue una VGG19 sin el encabezado de clasificación y enumere los nombres de las capas

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

print()
for layer in vgg.layers:
  print(layer.name)

Escoja capas intermedias en la red para representar el estilo y el contenido de la imagen:

content_layers = ['block5_conv2'] 

style_layers = ['block1_conv1',
                'block2_conv1',
                'block3_conv1', 
                'block4_conv1', 
                'block5_conv1']

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

Capas intermedias para el estilo y el contenido

¿Por qué las salidas intermedias de la red preentrenada de clasificación de imágenes nos permiten definir las representaciones de estilo y contenido?

A un nivel alto, para que la red pueda realizar la clasificación de la imagen (que para eso fue entrenada), debe comprender la imagen. Para eso debe tomar la imagen sin procesar como píxeles de entrada y construir una representación interna que convierta los píxeles de la imagen sin procesar en una comprensión compleja de las características presentes en la imagen.

Esta es otra de las razones por las que las redes neuronales convolucionales son capaces de generalizar bien: son capaces de captar las invarianzas y las características definitorias dentro de las clases (por ejemplo, gatos o perros) que son agnósticas al ruido de fondo y otras molestias. Así, entre el momento en que se ingresa la imagen sin procesar en el modelo y la etiqueta de clasificación de salida, el modelo actúa como un complejo extractor de características. Al acceder a las capas intermedias del modelo, es capaz de describir el contenido y el estilo de las imágenes de entrada.

Construir el modelo

Las redes en tf.keras.applications están diseñadas para que usted pueda extraer fácilmente los valores de la capa intermedia con la API funcional de Keras.

Para definir un modelo con la API funcional, especifique las entradas y las salidas:

model = Model(inputs, outputs)

La siguiente función, construye un modelo VGG19 que devuelve una lista de salidas de la capa intermedia:

def vgg_layers(layer_names):
  """ Creates a VGG model that returns a list of intermediate output values."""
  # Load our model. Load pretrained VGG, trained on ImageNet data
  vgg = tf.keras.applications.VGG19(include_top=False, weights='imagenet')
  vgg.trainable = False
  
  outputs = [vgg.get_layer(name).output for name in layer_names]

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

Y para crear el modelo:

style_extractor = vgg_layers(style_layers)
style_outputs = style_extractor(style_image*255)

#Look at the statistics of each layer's output
for name, output in zip(style_layers, style_outputs):
  print(name)
  print("  shape: ", output.numpy().shape)
  print("  min: ", output.numpy().min())
  print("  max: ", output.numpy().max())
  print("  mean: ", output.numpy().mean())
  print()

Calcular el estilo

El contenido de una imagen se representa mediante los valores de los mapas de características intermedias.

Lo que pasa es que el estilo de una imagen puede describirse por las medias y las correlaciones en los mapas de características diferentes. Calcule una matriz de Gram que incluya esta información. Para esto debe tomar el producto externo del vector de la característica con sí mismo en cada ubicación y promediar ese producto externo en todas las ubicaciones. La matriz de Gram puede calcularse para una capa en particular de la siguiente manera:

$$G^l_{cd} = \frac{\sum_{ij} F^l_{ijc}(x)F^l_{ijd}(x)}{IJ}$$

Puede implementarse de forma concisa con la función tf.linalg.einsum:

def gram_matrix(input_tensor):
  result = tf.linalg.einsum('bijc,bijd->bcd', input_tensor, input_tensor)
  input_shape = tf.shape(input_tensor)
  num_locations = tf.cast(input_shape[1]*input_shape[2], tf.float32)
  return result/(num_locations)

Extraer el estilo y el contenido

Construya un modelo que devuelva los tensores de estilo y de contenido.

class StyleContentModel(tf.keras.models.Model):
  def __init__(self, style_layers, content_layers):
    super(StyleContentModel, self).__init__()
    self.vgg = vgg_layers(style_layers + content_layers)
    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):
    "Expects float input in [0,1]"
    inputs = inputs*255.0
    preprocessed_input = tf.keras.applications.vgg19.preprocess_input(inputs)
    outputs = self.vgg(preprocessed_input)
    style_outputs, content_outputs = (outputs[:self.num_style_layers],
                                      outputs[self.num_style_layers:])

    style_outputs = [gram_matrix(style_output)
                     for style_output in style_outputs]

    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}

Cuando se lo llama en una imagen, el modelo devuelve la matriz de Gram (estilo) de style_layers y el contenido de content_layers:

extractor = StyleContentModel(style_layers, content_layers)

results = extractor(tf.constant(content_image))

print('Styles:')
for name, output in sorted(results['style'].items()):
  print("  ", name)
  print("    shape: ", output.numpy().shape)
  print("    min: ", output.numpy().min())
  print("    max: ", output.numpy().max())
  print("    mean: ", output.numpy().mean())
  print()

print("Contents:")
for name, output in sorted(results['content'].items()):
  print("  ", name)
  print("    shape: ", output.numpy().shape)
  print("    min: ", output.numpy().min())
  print("    max: ", output.numpy().max())
  print("    mean: ", output.numpy().mean())

Ejecutar el descenso de la gradiente

Con este extractor de estilo y contenido, ahora puede implementar el algoritmo de transferencia de estilo. Para hacer esto, calcule el error cuadrático medio de la salida de su imagen en relación con cada destino, y luego, tome la suma ponderada de estas pérdidas.

Establezca los valores de destino del estilo y del contenido:

style_targets = extractor(style_image)['style']
content_targets = extractor(content_image)['content']

Defina una tf.Variable donde se guarde la imagen que se va a optimizar. Para hacerlo rápido, inícielo con una imagen de contenido (la tf.Variable debe tener la misma forma que la imagen de contenido):

image = tf.Variable(content_image)

Ya que esta es una imagen flotante, defina una función para los valores de los píxeles permanezcan entre 0 y 1:

def clip_0_1(image):
  return tf.clip_by_value(image, clip_value_min=0.0, clip_value_max=1.0)

Cree un optimizador. En el artículo se recomienda usar LBFGS, pero Adam funciona bien también:

opt = tf.keras.optimizers.Adam(learning_rate=0.02, beta_1=0.99, epsilon=1e-1)

Para optimizar esto, use una combinación ponderada de dos pérdidas para obtener la pérdida total:

style_weight=1e-2
content_weight=1e4

def style_content_loss(outputs):
    style_outputs = outputs['style']
    content_outputs = outputs['content']
    style_loss = tf.add_n([tf.reduce_mean((style_outputs[name]-style_targets[name])**2) 
                           for name in style_outputs.keys()])
    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()])
    content_loss *= content_weight / num_content_layers
    loss = style_loss + content_loss
    return loss

Use tf.GradientTape para actualizar la imagen.

@tf.function()
def train_step(image):
  with tf.GradientTape() as tape:
    outputs = extractor(image)
    loss = style_content_loss(outputs)

  grad = tape.gradient(loss, image)
  opt.apply_gradients([(grad, image)])
  image.assign(clip_0_1(image))

Ahora, ejecute algunos pasos para probarlo:

train_step(image)
train_step(image)
train_step(image)
tensor_to_image(image)

Como sí funciona, realice una optimización más larga:

import time
start = time.time()

epochs = 10
steps_per_epoch = 100

step = 0
for n in range(epochs):
  for m in range(steps_per_epoch):
    step += 1
    train_step(image)
    print(".", end='', flush=True)
  display.clear_output(wait=True)
  display.display(tensor_to_image(image))
  print("Train step: {}".format(step))
  
end = time.time()
print("Total time: {:.1f}".format(end-start))

Pérdida total de variación

Una desventaja de esta implementación básica es que produce muchos artefactos de alta frecuencia. Puede reducir esta cantidad con un término de regularización explícito en los componentes de alta frecuencia de la imagen. En la transferencia de estilo, esto se conoce como total variation loss:

def high_pass_x_y(image):
  x_var = image[:, :, 1:, :] - image[:, :, :-1, :]
  y_var = image[:, 1:, :, :] - image[:, :-1, :, :]

  return x_var, y_var

x_deltas, y_deltas = high_pass_x_y(content_image)

plt.figure(figsize=(14, 10))
plt.subplot(2, 2, 1)
imshow(clip_0_1(2*y_deltas+0.5), "Horizontal Deltas: Original")

plt.subplot(2, 2, 2)
imshow(clip_0_1(2*x_deltas+0.5), "Vertical Deltas: Original")

x_deltas, y_deltas = high_pass_x_y(image)

plt.subplot(2, 2, 3)
imshow(clip_0_1(2*y_deltas+0.5), "Horizontal Deltas: Styled")

plt.subplot(2, 2, 4)
imshow(clip_0_1(2*x_deltas+0.5), "Vertical Deltas: Styled")

Aquí se muestra cómo se redujeron los componentes de alta frecuencia.

También, este componente de alta frecuencia es un detector de bordes, básicamente. Puedes obtener una salida similar con el detector de bordes Sobel, por ejemplo:

plt.figure(figsize=(14, 10))

sobel = tf.image.sobel_edges(content_image)
plt.subplot(1, 2, 1)
imshow(clip_0_1(sobel[..., 0]/4+0.5), "Horizontal Sobel-edges")
plt.subplot(1, 2, 2)
imshow(clip_0_1(sobel[..., 1]/4+0.5), "Vertical Sobel-edges")

La pérdida de regularización asociada es la suma de los cuadrados de los valores:

def total_variation_loss(image):
  x_deltas, y_deltas = high_pass_x_y(image)
  return tf.reduce_sum(tf.abs(x_deltas)) + tf.reduce_sum(tf.abs(y_deltas))

total_variation_loss(image).numpy()

Esto demostró como lo hace. Pero no hace falta que usted lo implemente, pues TensorFlow incluye una implementación estándar:

tf.image.total_variation(image).numpy()

Ejecutar la optimización de nuevo

Escoja una ponderación para total_variation_loss:

total_variation_weight=30

Ahora, agréguelo a la función train_step:

@tf.function()
def train_step(image):
  with tf.GradientTape() as tape:
    outputs = extractor(image)
    loss = style_content_loss(outputs)
    loss += total_variation_weight*tf.image.total_variation(image)

  grad = tape.gradient(loss, image)
  opt.apply_gradients([(grad, image)])
  image.assign(clip_0_1(image))

Reinicialice la variable imagen y el optimizador:

opt = tf.keras.optimizers.Adam(learning_rate=0.02, beta_1=0.99, epsilon=1e-1)
image = tf.Variable(content_image)

Y ejecute la optimización

import time
start = time.time()

epochs = 10
steps_per_epoch = 100

step = 0
for n in range(epochs):
  for m in range(steps_per_epoch):
    step += 1
    train_step(image)
    print(".", end='', flush=True)
  display.clear_output(wait=True)
  display.display(tensor_to_image(image))
  print("Train step: {}".format(step))

end = time.time()
print("Total time: {:.1f}".format(end-start))

Para finalizar, guarde el resultado:

file_name = 'stylized-image.png'
tensor_to_image(image).save(file_name)

try:
  from google.colab import files
except ImportError:
   pass
else:
  files.download(file_name)

Más información

En este tutorial se explica el algoritmo original de transferencia de estilo. Para ver una aplicación simple de la transferencia de estilo, échele un vistazo a este tutorial y encuentre más información sobre cómo usar el modelo arbitrario de transferencia de estilo de TensorFlow Hub.