|
import gradio as gr |
|
import tensorflow as tf |
|
|
|
encoder = tf.keras.models.load_model("nst-encoder.h5", compile=False) |
|
decoder = tf.keras.models.load_model("nst-decoder.h5", compile=False) |
|
|
|
def get_mean_std(tensor, epsilon=1e-5): |
|
axes = [1, 2] |
|
tensor_mean, tensor_var = tf.nn.moments(tensor, axes=axes, keepdims=True) |
|
tensor_std = tf.sqrt(tensor_var + epsilon) |
|
return tensor_mean, tensor_std |
|
|
|
def ada_in(style, content, epsilon=1e-5): |
|
c_mean, c_std = get_mean_std(content) |
|
s_mean, s_std = get_mean_std(style) |
|
t = s_std * (content - c_mean) / c_std + s_mean |
|
return t |
|
|
|
def load_resize(image): |
|
image = tf.image.convert_image_dtype(image, dtype="float32") |
|
image = tf.image.resize(image, (224, 224)) |
|
return image |
|
|
|
def infer(style, content): |
|
style = load_resize(style) |
|
style = style[tf.newaxis, ...] |
|
content = load_resize(content) |
|
content = content[tf.newaxis, ...] |
|
|
|
style_enc = encoder(style) |
|
content_enc = encoder(content) |
|
|
|
t = ada_in(style=style_enc, content=content_enc) |
|
|
|
recons_image = decoder(t) |
|
return recons_image[0].numpy() |
|
|
|
iface = gr.Interface( |
|
fn=infer, |
|
inputs=[gr.inputs.Image(label="style"), |
|
gr.inputs.Image(label="content")], |
|
outputs="image").launch() |
|
|