File size: 2,317 Bytes
f6ca457 d3d3291 f6ca457 d3d3291 f6ca457 |
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 |
from model import create_vgg_model
from data_setup import preprocess, deprocess, get_features, gram_matrix
from loss_functions import content_loss, style_loss, total_loss
from torch import optim
from pathlib import Path
import numpy as np
import gradio as gr
from PIL import Image
def predict(content_image, style_image):
# Create model
model = create_vgg_model()
# Transform images
content_img = preprocess(content_image)
style_img = preprocess(style_image)
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}
# Inference
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))
# Gradio Interface
example_list = [['content/content1.jpg',
'style/style1.jpg'],
['content/content2.jpg',
'style/style2.jpg'],
['content/content3.jpg',
'style/style3.jpg']]
title = "Neural Style Transfer 🎨"
description = "It will take about 1 minute for the result to be displayed. Since the algorithm runs a small number of epochs (to reduce the waiting time), the result will not always be as good as it should be."
article = "See the code on [GitHub](https://github.com/georgescutelnicu/Neural-Style-Transfer)."
demo = gr.Interface(fn=predict,
inputs=['image', 'image'],
outputs=gr.Image().style(width=256, height=256),
examples=example_list,
title=title,
description=description,
article=article)
demo.launch(debug=False,
share=False)
|