Spaces:
Running
Running
File size: 4,224 Bytes
f51ce08 148a566 f51ce08 148a566 f51ce08 148a566 f51ce08 148a566 f51ce08 |
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 |
# Install gradio
# Import necessary libs
import tensorflow as tf
import keras
import numpy as np
# Load the VGG19 model
from keras.applications.vgg19 import VGG19
from keras.applications.vgg19 import preprocess_input, decode_predictions
model = VGG19(include_top=False, weights='imagenet')
# Define style layers and content layers for VGG19
content_layers = 'block5_conv2'
style_layers = ['block1_conv1',
'block2_conv1',
'block3_conv1',
'block4_conv1',
'block5_conv1']
num_content_layers = len(content_layers)
num_style_layers = len(style_layers)
# Model for extracting style feature
style_model = keras.Model(inputs=model.input, outputs=[model.get_layer(name).output for name in style_layers], name='Style_model')
# Model for extracting content feature
content_model = keras.Model(inputs=model.input, outputs=model.get_layer(content_layers).output, name='Content_model')
# Style model summary
print(f'Style model: {style_model.name}')
style_model.summary()
# Content model summary
print(f'Content model: {content_model.name}')
content_model.summary()
# Compute content loss
def content_loss(gen, content_img):
return tf.reduce_sum(tf.square(content_img - gen))
# Gram matrix
def gram_matrix(image):
# Reshape image to 2-D array
channel = image.shape[-1] # Number of channels
a = tf.reshape(image, [-1, channel])
# Compute Gram matrix corresponding to the input image
dimension = a.shape[0] # Dimension M*N of an MxNxc image
return tf.matmul(a, a, transpose_a=True) / tf.cast(dimension, tf.float32)
# Style loss
def style_loss(gen, style_img):
total_loss = 0.
for i in range(len(gen)):
gen_gram = gram_matrix(gen[i])
style_gram = gram_matrix(style_img[i])
# Compute style loss
total_loss += tf.reduce_sum(tf.square(gen_gram-style_gram))/4
return total_loss
# Training parameter
lr = 10.
optimizer = keras.optimizers.Adam(lr)
alpha = 1e-7
beta = 1e-10
style_weight = len(style_layers)
# Train step
@tf.function(reduce_retracing=True)
def train_step(image, process_content_image, process_style_image):
with tf.GradientTape() as tape:
# Calculate content loss
gen_content_feature = content_model(image)
content_feature = content_model(process_content_image)
loss_content = content_loss(gen_content_feature, content_feature)
# Calculate style loss
gen_style_feature = style_model(image)
style_feature = style_model(process_style_image)
loss_style = style_loss(gen_style_feature, style_feature)
# Calculate total loss
total_loss = alpha*loss_content + beta*loss_style/style_weight
# Calculate gradient
grad = tape.gradient(total_loss, image)
# Apply gradient
optimizer.apply_gradients([(grad, image)])
# Return loss for visualize
return total_loss
# Deprocess image
def deprocess_image(x):
# Util function to convert a tensor into a valid image
x = x.reshape((256,256,3))
# Remove zero-center by mean pixel
x[:, :, 0] += 103.939
x[:, :, 1] += 116.779
x[:, :, 2] += 123.68
# 'BGR'->'RGB'
x = x[:, :, ::-1]
x = np.clip(x, 0, 255).astype("uint8")
return x
def load_image(image, RESO=256):
image = tf.image.convert_image_dtype(image, tf.float32) * 255
image = tf.image.resize(image,[RESO, RESO],method=tf.image.ResizeMethod.BICUBIC)
image = tf.expand_dims(preprocess_input(image), 0)
return image
def style_transfer(content_image, style_image, step=100):
process_content_image = load_image(content_image)
process_style_image = load_image(style_image)
generate_image = tf.Variable(process_content_image, trainable=True)
for step in range(1, step+1):
# Train step
loss = train_step(generate_image, process_content_image, process_style_image)
return deprocess_image(generate_image[0].numpy())
import gradio as gr
demo = gr.Interface(
fn=style_transfer,
inputs=[gr.Image(label='Input Image'), gr.Image(label='Style Image'), gr.Slider(10, 200, 50, step=10, label='Step', show_label=True)],
outputs=gr.Image(label='Style Transfer Image'),
)
demo.launch(debug=True, share=True) |