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

model = hub.load("https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2")

def style(normal_img,style_img):
    normal_img=tf.convert_to_tensor(normal_img,np.float32)[tf.newaxis, ...] /255.
    style_img=tf.convert_to_tensor(style_img,np.float32)[tf.newaxis,...] /255.
    out=model(normal_img,style_img)
    print(out)
    res=out[0]
    print(res)
    return Image.fromarray(np.uint8(res[0]*255))
normal_image_input = gr.inputs.Image(label="Content Image")
style_image_input = gr.inputs.Image(shape=(256, 256), label="Style Image")
# Examples
golden_gate = ["samples/golden_gate_bridge.jpeg", "samples/the_great_wave.jpeg"]
joshua_tree = ["samples/joshua_tree.jpeg", "samples/starry_night.jpeg"]
glacier = ["samples/glacier_national_park.jpeg", "samples/the_scream.jpg"]

app_interface = gr.Interface(fn=style,
                             inputs=[normal_image_input, style_image_input],
                             outputs="image",
                             title="Fast Neural Style Transfer",
                             description="Gradio demo for Fast Neural Style Transfer using a pretrained Image Stylization model from TensorFlow Hub. To use it, simply upload a content image and style image, or click one of the examples to load them. To learn more about the project, please find the references listed below.",
                             examples=[glacier, golden_gate, joshua_tree],
                             article="**References**\n\n"
                             "<a href='https://www.tensorflow.org/hub/tutorials/tf2_arbitrary_image_stylization' target='_blank'>1. Tutorial to implement Fast Neural Style Transfer using the pretrained model from TensorFlow Hub</a>  \n"
                             "<a href='https://huggingface.co/spaces/luca-martial/neural-style-transfer' target='_blank'>2. The idea to build a neural style transfer application was inspired from this Hugging Face Space </a>")
app_interface.launch()