[PYTHON] Implementieren Sie Style Transfer mit Pytorch
Implementieren wir die Stilübertragung mit Pytorch, um ein zusammengesetztes Bild zu generieren, das den Stil des Bildes auf ein bestimmtes Bild anwendet. Verwenden Sie Google Colaboratory (Colab). Da ich studiere, kann es zu Fehlinterpretationen kommen.
Implementieren Sie die Stilübertragung in Pytorch, um zusammengesetzte Bilder zu generieren
Dies ist das fertige Produkt dieses Artikels.
Das Folgende ist das Originalbild und Gochs Bild, das die Quelle des Stils war. Das Bild wurde von unsplash aufgenommen. Gochs Stil passt gut zum Bild der Torii. unsplash
Vorbereitung
Lassen Sie uns die Verarbeitung vor dem Lernen implementieren.
Erforderliche Bibliotheken und Google Drive-Mount
Importieren Sie die erforderlichen Bibliotheken. Dieses Mal werde ich Google Drive mounten und mehr über die Bilder auf dem Laufwerk erfahren.
Lassen Sie uns Pillow installieren, um Pytorch und Bilder anzuzeigen.
!pip install torch torchvision
!pip install Pillow==4.0.0
Importieren Sie die Bibliothek.
%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
Mounten Sie Google Drive. Sie können jetzt von Colab aus auf das Laufwerk zugreifen. Laden Sie das Image hoch, aus dem das zusammengesetzte Image generiert wird, und das Image, auf das der Stil auf das Laufwerk angewendet wird.
from google.colab import drive
drive.mount("/content/drive")
Modelldefinition
Verwenden Sie VGG. Definieren Sie auch die GPU für das Gerät.
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)
Definieren Sie eine Funktion zum Anzeigen eines Bildes
Definieren Sie eine Funktion, die ein Bild aus dem Pfad lädt und einen Tensor zurückgibt. Sie können die Größe angeben.
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
Laden und Anzeigen von Bildern
Definiert den Pfad des Bildes. Bitte legen Sie hier einen praktischen Pfad fest. Das Bild, das die Quelle des endgültigen Bildes ist, wird als content_path definiert, und das Bild, das die Quelle des Stils ist, wird als style_path definiert.
images_path ='drive/My Drive/'
content_path = images_path + 'content.jpg'
style_path = images_path + 'style.jpg'
Laden wir es. Die Größen der beiden Bilder stimmen überein.
content = load_image(content_path).to(device)
style = load_image(style_path, shape=content.shape[-2:]).to(device)
Definieren Sie eine Funktion, die Tensor in Numpy konvertiert und es Ihnen ermöglicht, ein Bild anzuzeigen
def im_convert(tensor):
#Tensor np.In Array konvertieren
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
Zeigen Sie das Bild auf Colab an.
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")
Features extrahieren
CNN ist grob in zwei Schichten unterteilt. Die erste Halbebene, die Features aus dem Bild extrahiert, und die zweite Halbebene, die das Bild schließlich aus den extrahierten Features klassifiziert.
Die Merkmale der beiden Bilder werden von CNN extrahiert, und die Differenz zwischen den Merkmalen wird als Verlustfunktion verwendet, um das endgültige Bild zu trainieren. Fügen wir also das Bild in das CNN ein und extrahieren die Features in einer bestimmten Ebene. Extrahieren Sie in der 0., 5., 10., 19., 21. und 28. Schicht von VGG.
Erstens ist die Definition der Funktion.
def get_features(image, model):
# Feature Extraction
#Ebene zum Extrahieren von Features
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():
#Schalten Sie den CNN
image = layer(image)
#Extrahieren Sie Features auf einer bestimmten Ebene
#Hier 0, 5, 10, 19, 21, 28
if name in layers:
features[layers[name]] = image
return features
Lassen Sie uns die Funktionen extrahieren.
content_features = get_features(content, vgg)
style_features = get_features(style, vgg)
Das Bild, das dem Malstil zugrunde liegt, vergleicht nicht nur Merkmale, sondern auch die Gramm-Matrix von Merkmalen. Definieren wir eine Funktion, die die Gramm-Matrix berechnet.
def gram_matrix(tensor):
#Berechnen Sie die Gramm-Matrix
_, d, h, w = tensor.size()
tensor = tensor.view(d, h * w)
#Transmutationsmatrix und Matrixmultiplikation
gram = torch.mm(tensor, tensor.t())
return gram
Halten Sie eine Gramm-Matrix von Stilen
style_grams = {layer: gram_matrix(style_features[layer]) for layer in style_features}
Legt das Gewicht der zu vergleichenden Ebenen fest.
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
Das endgültige Bild wird durch Kopieren des Inhaltsbilds trainiert. Kopieren Sie es also und definieren Sie es als Ziel.
target = content.clone().requires_grad_(True).to(device)
Lernen
Die Vorbereitung ist abgeschlossen. Lass uns lernen! Notieren Sie den Fortschritt, damit Sie später ein Video erstellen können. Lassen Sie uns die Hyperparameter einstellen.
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
Jetzt lernt es. Bitte überprüfen Sie die Kommentare.
for ii in range(1, steps+1):
target_features = get_features(target, vgg)
#Berechnung der Verlustfunktion mit Inhalt
content_loss = torch.mean((target_features['conv4_2'] - content_features['conv4_2'])**2)
#Berechnung der Verlustfunktion mit Stil
style_loss = 0
for layer in style_weights:
target_feature = target_features[layer]
target_gram = gram_matrix(target_feature)
style_gram = style_grams[layer]
#Gewicht*Summe der Fehlerquadrate
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)
#Totalverlustfunktion
total_loss = content_weight * content_loss + style_weight * style_loss
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
#nachverfolgen
if ii % show_every == 0:
print('Total loss: ', total_loss.item())
print('Iteration: ', ii)
plt.imshow(im_convert(target))
plt.axis("off")
plt.show()
#Für Video speichern
if ii % capture_frame == 0:
image_array[counter] = im_convert(target)
counter = counter + 1
Ich denke, es wird einige Zeit dauern. 600 ..
5100th
9900th
Ich lerne richtig.
Video exportieren
Erstellen Sie ein Video mit 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()
Lassen Sie uns herunterladen.
from google.colab import files
files.download('output.mp4')
Sie können den Lernprozess verfolgen, indem Sie das Video abspielen.
abschließend
Danke für deine harte Arbeit! Dieses Mal habe ich versucht, Style Transfer mit Pytorch zu implementieren. Sie können es ganz einfach mit der Google Colaboratory-GPU ausprobieren. Wenn Sie es in Google Drive speichern, ist es in Ordnung, wenn Ihnen die Colab-Instanzen ausgehen.
Bitte probieren Sie verschiedene andere Bilder!
Recommended Posts