'''NEURAL STYLE TRANSFER '''

import numpy as np
import tensorflow as tf
import tensorflow_hub as hub
import gradio as gr
from PIL import Image

np.set_printoptions(suppress=True)
model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')

def tensor_to_image(tensor):
    tensor *= 255
    tensor = np.array(tensor, dtype=np.uint8)
    if tensor.ndim > 3:
        tensor = tensor[0]
    return Image.fromarray(tensor)

def transform_my_model(content_image, style_image):
    content_image = content_image.astype(np.float32)[np.newaxis, ...] / 255.0
    style_image = style_image.astype(np.float32)[np.newaxis, ...] / 255.0
    stylized_image = model(tf.constant(content_image), tf.constant(style_image))[0]
    return tensor_to_image(stylized_image)

demo = gr.Interface(
    fn=transform_my_model,
    inputs=[gr.Image(label="Content Image"), gr.Image(label="Style Image")],
    outputs=gr.Image(label="Result"),
    title="Style Transfer",
    examples=[
        ["Content_Images/contnt12.jpg", "VG516.jpg"],
        ["Content_Images/contnt2.jpg", "Content_Images/styl9.jpg"],
        ["Content_Images/contnt.jpg", "Content_Images/styl22.jpg"]
    ],
    article="References-\n\nExploring the structure of a real-time, arbitrary neural artistic stylization network. Golnaz Ghiasi, Honglak Lee, Manjunath Kudlur, Vincent Dumoulin."
)

demo.launch(share=True)