File size: 4,123 Bytes
75f2ed4 493509d 75f2ed4 493509d 75f2ed4 493509d 75f2ed4 493509d |
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 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 |
from PIL import Image
import torch
import torch.nn as nn
import torch.optim as optim
import torchvision.transforms as transforms
import torchvision.models as models
import torchvision.utils as utils
import gradio as gr
device = 'cpu'
if torch.backends.mps.is_available():
device = 'mps'
if torch.cuda.is_available():
device = 'cuda'
print('DEVICE:', device)
class VGG_19(nn.Module):
def __init__(self):
super(VGG_19, self).__init__()
self.model = models.vgg19(pretrained=True).features[:30]
for i, _ in enumerate(self.model):
if i in [4, 9, 18, 27]:
self.model[i] = nn.AvgPool2d(kernel_size=2, stride=2, padding=0)
def forward(self, x):
features = []
for i, layer in enumerate(self.model):
x = layer(x)
if i in [0, 5, 10, 19, 28]:
features.append(x)
return features
model = VGG_19().to(device)
for param in model.parameters():
param.requires_grad = False
def load_img(img: Image, img_size):
original_size = img.size
transform = transforms.Compose([
transforms.Resize((img_size, img_size)),
transforms.ToTensor()
])
img = transform(img).unsqueeze(0)
return img, original_size
def load_img_from_path(path_to_image, img_size):
img = Image.open(path_to_image)
original_size = img.size
transform = transforms.Compose([
transforms.Resize((img_size, img_size)),
transforms.ToTensor()
])
img = transform(img).unsqueeze(0)
return img, original_size
def save_img(img, original_size):
img = img.cpu().clone()
img = img.squeeze(0)
# address tensor value scaling and quantization
img = torch.clamp(img, 0, 1)
img = img.mul(255).byte()
unloader = transforms.ToPILImage()
img = unloader(img)
img = img.resize(original_size, Image.Resampling.LANCZOS)
return img
def transfer_style(content_image):
style_img_filename = 'StarryNight.jpg'
img_size = 512
content_img, original_size = load_img(content_image, img_size)
content_img = content_img.to(device)
style_img = load_img_from_path(f'./style_images/{style_img_filename}', img_size)[0].to(device)
iters = 100
lr = 1e-1
alpha = 1
beta = 1
generated_img = content_img.clone().requires_grad_(True)
optimizer = optim.Adam([generated_img], lr=lr)
for iter in range(iters+1):
generated_features = model(generated_img)
content_features = model(content_img)
style_features = model(style_img)
content_loss = 0
style_loss = 0
for generated_feature, content_feature, style_feature in zip(generated_features, content_features, style_features):
batch_size, n_feature_maps, height, width = generated_feature.size()
content_loss += (torch.mean((generated_feature - content_feature) ** 2))
G = torch.mm((generated_feature.view(batch_size * n_feature_maps, height * width)), (generated_feature.view(batch_size * n_feature_maps, height * width)).t())
A = torch.mm((style_feature.view(batch_size * n_feature_maps, height * width)), (style_feature.view(batch_size * n_feature_maps, height * width)).t())
E_l = ((G - A) ** 2)
w_l = 1/5
style_loss += torch.mean(w_l * E_l)
total_loss = alpha * content_loss + beta * style_loss
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
yield save_img(generated_img, original_size), str(round(iter/iters*100))+'%'
yield save_img(generated_img, original_size), str(round(iter/iters*100))+'%'
interface = gr.Interface(
fn=transfer_style,
inputs=[gr.Image(label='Content', type='pil')],
outputs=[
gr.Image(label='Output', show_download_button=True),
gr.Label(label='Progress')
],
title="Starry Night Style Transfer",
api_name='style',
allow_flagging='never',
).launch(inbrowser=True) |