[PYTHON] Implémenter le transfert de style avec Pytorch

Implémentons le transfert de style avec Pytorch pour générer une image composite qui applique le style de l'image à une certaine image. Utilisez Google Colaboratory (Colab). Depuis que j'étudie, il peut y avoir des erreurs d'interprétation.

Implémenter le transfert de style dans Pytorch pour générer des images composites

Ceci est le produit fini de cet article.

Ce qui suit est l'image originale et l'image de Goch qui était la source du style. L'image a été prise de unsplash. Le style de Goch est bien appliqué à l'image du torii. unsplash

Préparation

Implémentons le traitement avant d'apprendre.

Bibliothèques requises et montage Google Drive

Importez les bibliothèques requises. Cette fois, je vais monter Google Drive et en apprendre davantage sur les images du lecteur.

Installons Pillow pour afficher le pytorch et les images.

!pip install torch torchvision
!pip install Pillow==4.0.0

Importez la bibliothèque.

%matplotlib inline
import torch
import torch.optim as optim
from torchvision import transforms, models
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np

Montez Google Drive. Vous pouvez désormais accéder au lecteur depuis Colab. Téléchargeons l'image à partir de laquelle l'image composite est générée et l'image à laquelle le style est appliqué au lecteur.

from google.colab import drive
drive.mount("/content/drive")

Définition du modèle

Utilisez VGG. Définissez également gpu pour le périphérique.

vgg = models.vgg19(pretrained=True).features

for param in vgg.parameters():
  param.requires_grad_(False)
 
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
vgg.to(device)

Définir une fonction pour afficher une image

Définissez une fonction qui charge une image à partir du chemin et renvoie un Tensor. Vous pouvez spécifier la taille.

def load_image(img_path, max_size=400, shape=None):
  
  image = Image.open(img_path).convert('RGB')
  if max(image.size) > max_size:
    size = max_size
  else:
    size = max(image.size)
  
  if shape:
    size = shape
  in_transform = transforms.Compose([
                    transforms.Resize(size),
                    transforms.ToTensor(),
                                            ])
  
  image = in_transform(image).unsqueeze(0)
  
  return image

Charger et afficher des images

Définit le chemin de l'image. Veuillez définir un chemin de commodité ici. L'image qui est la source de l'image finale est définie comme content_path, et l'image qui est la source du style est définie comme style_path.

images_path ='drive/My Drive/'
content_path = images_path + 'content.jpg'
style_path = images_path + 'style.jpg'

Chargez-le. Les tailles des deux images correspondent.

content = load_image(content_path).to(device)
style = load_image(style_path, shape=content.shape[-2:]).to(device)

Définissez une fonction qui convertit Tensor en numpy et vous permet d'afficher une image

def im_convert(tensor):
  #Tensor np.Convertir en tableau
  image = tensor.to("cpu").clone().detach()
  image = image.numpy().squeeze()
  image = image.transpose(1,2,0)
  image = image * np.array((0.5, 0.5, 0.5) + np.array((0.5, 0.5, 0.5)))
  image = image.clip(0, 1)
  
  return image

Affichez l'image sur Colab.

fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 10))
ax1.imshow(im_convert(content))
ax1.axis("off")
ax2.imshow(im_convert(style))
ax2.axis("off")

Extraire les fonctionnalités

CNN est à peu près divisé en deux couches. La première moitié de la couche qui extrait les caractéristiques de l'image et la seconde moitié de la couche qui classe finalement l'image à partir des caractéristiques extraites.

Les caractéristiques des deux images sont extraites par CNN, et la différence entre les caractéristiques est utilisée comme une fonction de perte pour entraîner l'image finale. Alors mettons l'image dans le CNN et extrayons les caractéristiques dans une couche spécifique. Extrait dans les 0e, 5e, 10e, 19e, 21e et 28e couches de VGG.

La première est la définition de la fonction.

def get_features(image, model):
#   Feature Extraction
#Couche pour extraire des fonctionnalités
  layers = {'0': 'conv1_1',
            '5': 'conv2_1',
            '10': 'conv3_1',
            '19': 'conv4_1',
            '21': 'conv4_2',
            '28': 'conv5_1',}
  
  features = {}
  
  for name, layer in model._modules.items():
#Tournez le CNN
    image = layer(image)
#Extraire des entités sur un calque spécifique
#Ici, 0e, 5e, 10e, 19e, 21e, 28e
    if name in layers:
      features[layers[name]] = image
      
  return features

Extrayons les fonctionnalités.

content_features = get_features(content, vgg)
style_features = get_features(style, vgg)

L'image qui est à la base du style de peinture ne compare pas simplement les caractéristiques, mais compare la matrice gramme des caractéristiques. Définissons une fonction qui calcule la matrice gramme.

def gram_matrix(tensor):
  #Calculer la matrice du gramme
  _, d, h, w = tensor.size()
  tensor = tensor.view(d, h * w)
#Matrice de transmutation et multiplication matricielle
  gram = torch.mm(tensor, tensor.t())
  
  return gram

Tenir une matrice gramme de styles

style_grams = {layer: gram_matrix(style_features[layer]) for layer in style_features}

Définit le poids des couches à comparer.

style_weights = {'conv1_1':1.,
                 'conv2_1':0.75,
                 'conv3_1':0.2,
                 'conv4_1':0.2,
                 'conv5_1':0.2}

content_weight = 1 #alpha
style_weight = 1e6 #blue

L'image finale sera formée en copiant l'image de contenu, alors copiez-la et définissez-la comme cible.

target = content.clone().requires_grad_(True).to(device)

Apprentissage

La préparation est terminée. Apprenons! Enregistrez la progression afin de pouvoir faire une vidéo plus tard. Définissons les hyper paramètres.

show_every = 300
optimizer = optim.Adam([target], lr=0.003)
steps = 10000
total_capture_frame_number = 500

height, width, channels = im_convert(target).shape
image_array = np.empty(shape=(total_capture_frame_number, height, width, channels))
capture_frame =steps/total_capture_frame_number
counter = 0

Maintenant, c'est apprendre. Veuillez vérifier les commentaires.

for ii in range(1, steps+1):
  target_features = get_features(target, vgg)
#Calcul de la fonction de perte avec le contenu
  content_loss = torch.mean((target_features['conv4_2'] - content_features['conv4_2'])**2)
#Calcul de la fonction de perte avec style
  style_loss = 0
  
  for layer in style_weights:
    target_feature = target_features[layer]
    target_gram = gram_matrix(target_feature)
    style_gram = style_grams[layer]
#poids*Erreur de somme des carrés
    layer_style_loss = style_weights[layer] * torch.mean((target_gram - style_gram)**2)
    _, d, h, w = target_feature.shape
    style_loss += layer_style_loss / (d * h * w)
  
  #Fonction de perte totale
  total_loss = content_weight * content_loss + style_weight * style_loss
  
  optimizer.zero_grad()
  total_loss.backward()
  optimizer.step()
  
  #suivre
  if ii % show_every == 0:
    print('Total loss: ', total_loss.item())
    print('Iteration: ', ii)
    plt.imshow(im_convert(target))
    plt.axis("off")
    plt.show()
    
#Stocker pour la vidéo
  if ii % capture_frame == 0:
    image_array[counter] = im_convert(target)
    counter = counter + 1

Je pense que cela prendra du temps. 600e

5100e

9900e

J'apprends correctement.

Exporter la vidéo

Créez une vidéo en utilisant OpenCV.

import cv2

frame_height, frame_width, _ = im_convert(target).shape
vid = cv2.VideoWriter('output.mp4', cv2.VideoWriter_fourcc(*'XVID'), 30, (frame_width, frame_height))

for i in range(total_capture_frame_number):
  img = image_array[i]
  img = img*255
  img = np.array(img, dtype = np.uint8)
  img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
  vid.write(img)
  
vid.release()

Téléchargeons.

from google.colab import files
files.download('output.mp4')

Vous pouvez regarder le processus d'apprentissage en lisant la vidéo.

en conclusion

Je vous remercie pour votre travail acharné! Cette fois, j'ai essayé d'implémenter Style Transfer en utilisant Pytorch. Vous pouvez facilement l'essayer en utilisant le GPU Google Colaboratory. Si vous l'enregistrez dans Google Drive, ce n'est pas grave si vous manquez d'instances Colab.

Veuillez essayer diverses autres images!