File size: 3,404 Bytes
67953fe |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
import torch
from torchvision import transforms
from PIL import Image
import numpy as np
import gradio as gr
from torch import optim
import torchvision
device = "cuda" if torch.cuda.is_available() else "cpu"
def create_vgg_model():
model_weights = torchvision.models.VGG19_Weights.DEFAULT
model = torchvision.models.vgg19(weights=model_weights)
for param in model.parameters():
param.requires_grad = False
model = model.features
return model
def preprocess(img):
image = Image.fromarray(img).convert('RGB')
imsize = 196
transform = transforms.Compose([
transforms.Resize((imsize, imsize)),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])
image = transform(image)
image = image.unsqueeze(dim=0)
return image
def deprocess(image):
image = image.clone()
image = image.squeeze(0)
image = image.permute(1, 2, 0)
image = image.cpu().detach().numpy()
image = image * np.array([0.229, 0.224, 0.225]) + np.array([0.485, 0.456, 0.406])
image = image.clip(0, 1)
return image
def get_features(image, model):
features = {}
layers = {
'0': 'layer_1',
'5': 'layer_2',
'10': 'layer_3',
'19': 'layer_4',
'28': 'layer_5'
}
x = image
for name, layer in model._modules.items():
x = layer(x)
if name in layers:
features[layers[name]] = x
return features
def gram_matrix(image):
b, c, h, w = image.size()
image = image.view(c, h * w)
gram = torch.mm(image, image.t())
return gram
def content_loss(target, content):
return torch.mean((target - content) ** 2)
def style_loss(target_features, style_grams):
loss = 0
for layer in target_features:
target_gram = gram_matrix(target_features[layer])
style_gram = style_grams[layer]
layer_style_loss = torch.mean((target_gram - style_gram) ** 2)
loss += layer_style_loss
return loss
def total_loss(content_loss, style_loss, alpha, beta):
return alpha * content_loss + beta * style_loss
def predict(content_image, style_image):
model = create_vgg_model().to(device).eval()
content_img = preprocess(content_image).to(device)
style_img = preprocess(style_image).to(device)
target_img = content_img.clone().requires_grad_(True)
content_features = get_features(content_img, model)
style_features = get_features(style_img, model)
style_gram = {layer: gram_matrix(style_features[layer]) for layer in style_features}
optimizer = optim.Adam([target_img], lr=0.06)
alpha_param = 1
beta_param = 1e2
epochs = 60
for i in range(epochs):
target_features = get_features(target_img, model)
c_loss = content_loss(target_features['layer_4'], content_features['layer_4'])
s_loss = style_loss(target_features, style_gram)
t_loss = total_loss(c_loss, s_loss, alpha_param, beta_param)
optimizer.zero_grad()
t_loss.backward()
optimizer.step()
results = deprocess(target_img)
return Image.fromarray((results * 255).astype(np.uint8))
title = "Neural Style Transfer 🎨"
demo = gr.Interface(fn=predict,
inputs=['image', 'image'],
outputs=gr.Image(),
title=title)
demo.launch(debug=False, share=False)
|