luca-martial's picture
debug
0da3808
raw
history blame
1.51 kB
import gradio as gr
import tensorflow as tf
import tensorflow_hub as hub
import matplotlib.pyplot as plt
import numpy as np
import PIL.Image
hub_model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
def tensor_to_image(tensor):
tensor = tensor*255
tensor = np.array(tensor, dtype=np.uint8)
if np.ndim(tensor)>3:
assert tensor.shape[0] == 1
tensor = tensor[0]
return PIL.Image.fromarray(tensor)
def stylize(content_image, style_image):
content_image = content_image.astype(np.float32)[np.newaxis, ...] / 255.
style_image = style_image.astype(np.float32)[np.newaxis, ...] / 255.
stylized_image = hub_model(tf.constant(content_image), tf.constant(style_image))[0]
return tensor_to_image(stylized_image)
content_examples =[["example_paris.jpeg"], ["example_vangogh.jpeg"]]
style_examples = [["example_aristotle.jpeg"], ["example_dali.jpeg"]]
title = "Fast Neural Style Transfer using TF-Hub"
description = "Demo for neural style transfer using the pretrained Arbitrary Image Stylization model from TensorFlow Hub."
inputs = gr.inputs.Image(label="Content Image", source=["upload", "webcam"])
outputs = gr.outputs.Image(label="Style Image")
iface = gr.Interface(fn=stylize,
inputs=["image", "image"],
outputs="image",
title=title,
description=description,
examples=[content_examples, style_examples])
iface.launch()