Domina Neural Style Transfer – Explicación y ejemplo práctico (Pytorch)

En este artículo, explicaremos cómo podemos implementar la técnica de Neural Style Transfer. Para ello nos hemos basado en el siguiente artículo de la documentación de Pytorch, ademós de en el libro Pytorch Computer Vision Cookbook.

¿ Qué es Neural Style Transfer?

La técnica de Transferencia Neuronal de Estilo de forma resumida, consiste en a partir de dos imágenes, donde una es el contenido y otra es el estilo, combinarlas para generar la misma imagen que en el contenido pero con el estilo de la otra.

Como se puede ver, la imagen A, es la que representa el contenido, es simplemente una imagen a la que se le quiere aplicar esta técnica para generar la misma imagen pero con un estilo diferente. Las imágenes B, C y D son la misma imagen pero con estilos diferentes.

A diferencia de este artículo donde entrenábamos una Red Neuronal Convolucional, en este caso, lo que se optimiza es la imagen que se quiere generar, donde se quiere minimizar las diferencias con el contenido y el estilo.

Función de coste de la transferencia de estilo

El objetivo a la hora de utilizar algoritmos de Deep Learning, siempre es minimizar el valor de una función de coste. En este caso, se tiene una función de coste compuesta, ya que se combina tanto una función de contenido como una función de estilo. Optimizar ambos costes es clave para conseguir un style transfer correcto.

J(G) = \alpha J_{content}(C, G) + \beta J_{style}(S, G)

Función de coste del contenido

La función de coste del contenido

J_{content}(C, G)

calcula la diferencia del contenido entre la imagen original y la imagen a generar. Para ello, se coge la función de activación de una de las capas intermedias, y simplemente se comparan las salidas de ambas.

Función de coste del estilo

La función de coste del estilo

J_{style}(S, G)

es algo más compleja, ya que se debe calcular una matriz de estilo, llamada comúnmente, Gram Matrix, para medir las correlaciones existentes entre las funciones de activación.

Se calcula una matriz para la imagen original y para la generada, para seguidamente, comparar ambas matrices y ver así la diferencia de estilo.

Implementación del algoritmo en Pytorch

Para esta implementación, partiremos de dos imágenes. Donde una representa el contenido y otra el estilo.

En primer lugar, cargaremos todas los paquetes necesarios y prepararemos el entorno para utilizar la GPU.

from PIL import Image
import torchvision.transforms as transforms
import torch
import torchvision.models as models
import torch.nn.functional as F
from torch import optim
from torchvision.transforms.functional import to_pil_image
import matplotlib.pyplot as plt

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"Using: {device}")

A continuación cargaremos las imágenes y las transformaremos en formato tensor.

h, w = 256, 256
mean_rgb = (0.485, 0.456, 0.406)
std_rgb = (0.229, 0.224, 0.225)
transformer = transforms.Compose([
    transforms.Resize((h, w)),
    transforms.ToTensor(),
    transforms.Normalize(mean_rgb, std_rgb)])
content = "data/content.jpeg"
style = "data/style.jpg"
content_img = Image.open(content)
style_img = Image.open(style)
content_tensor = transformer(content_img)
style_tensor = transformer(style_img)

Como hemos mencionado anteriormente, en esta técnica no se entrena un modelo. Se parte de un modelo ya preentrenado y se cogen las salidas de capas intermedias para ir comparándolas. Por lo que el siguiente paso consistirá en cargar la red VGG-19 preentrenada. Esta es la red elegida en la publicación, pero si quieres coger otra, adelante.

model_vgg = models.vgg19(pretrained=True).features.to(device).eval()
for param in model_vgg.parameters():
    param.requires_grad = False

Además, en esta técnica, no se cogen todas las capas, sino unas determinadas, pero al igual que antes si quieres probar con otras siéntete libre.

content_layer = 'conv4_2'
feature_layers = {
    '0': 'conv1_1',
    '5': 'conv2_1',
    '10': 'conv3_1',
    '19': 'conv4_1',
    '21': 'conv4_2',
    '28': 'conv5_1'
}

Se han creado varias funciones auxiliares para calcular la función de contenido y la de estilo, así como para coger los valores de las capas intermedias tanto para la imagen de contenido como la de estilo.

Finalmente, optimizamos nuestra imagen, para que vaya obteniendo el estilo deseado.

for epoch in range(num_epochs + 1):
    optimizer.zero_grad()
    input_features = get_features(input_tensor, model_vgg, feature_layers)
    content_loss = get_content_loss(input_features[content_layer], content_features[content_layer])
    style_loss = get_style_loss(input_features, style_features, feature_layers, content_layer)
    neural_loss = content_weight * content_loss + style_weight * style_loss
    neural_loss.backward()
    optimizer.step()
    if(epoch % 50 == 0):
        print('epoch {} | Content Loss: {:.2} | Style Loss {:.2}'
            .format(epoch, content_loss, style_loss))

El resultado es el siguiente.

Conclusiones del Neural Style Transfer

Nosotros hemos partido de la imagen original, y la hemos modificado, pero si lo prefieres puedes inicializarla con valores aleatorios a ver qué resultados obtienes.

Aunque este método proporciona buenos resultados, es bastante lento, por lo que está lejos de ser un sistema que pueda usarse en tiempo real. Por aquá dejamos un enlace a una publicación que propone como hacer Neural Style Transfer en tiempo real.

Si quieres acceder a todo el código, te lo dejamos en este enlace de Github.

Esperamos que hayas aprendido algo en este artículo. Si tienes cualquier duda, no dudes en dejarla en los comentarios.