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

# Load model from TF-Hub
style_transfer_model = hub.load("https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2")

# function to Stylize the Image or to perform a style transfer
def do_style_transfer(content_image, style_image):
    # Convert to float32 numpy array, add batch dimension, and normalize to range [0, 1]. Example using numpy:
    content_image = tf.convert_to_tensor(content_image, np.float32)[tf.newaxis, ...] / 255.
    style_image = tf.convert_to_tensor(style_image, np.float32)[tf.newaxis, ...] / 255.
    # Stylize image
    output = style_transfer_model(content_image, style_image)
    stylized_image = output[0]
    
    return Image.fromarray(np.uint8(stylized_image[0] * 255))


content_image_input = gr.inputs.Image(label="Content Image")
style_image_input = gr.inputs.Image(shape=(256, 256), label="Style Image")

# Add image examples for users    
golden_gate = ["example_deadpool2.jpeg", "style_the_great_wave.jpeg"]
joshua_tree = ["example_deadpool.png", "style_starry_night.jpeg"]
avatar = ["example_avatar.jpeg", "style_the_scream.jpg"]
joker = ["example_joker.jpeg", "style_polasticot1.jpeg"]
einstein = ["example_einstein.jpeg", "style_polasticot2.jpeg"]
monalisa = ["example_mona1.jpeg", "style_polasticot3.jpeg"]
paris = ["example_paris.jpeg", "style_vangogh.jpeg"]

# Customize interface
title = "Fast Neural Style Transfer using TF-Hub"
description = "<p style='text-align: center'> Demo for Neural Style Transfer using a pretrained Arbitrary 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.</p>"
article = r"""<b>:: References ::</b>
<br>
<a href='https://www.tensorflow.org/hub/tutorials/tf2_arbitrary_image_stylization' target='_blank'>
Tutorial to implement Fast Neural Style Transfer using the pretrained model from TensorFlow Hub
</a>
<br>
<a href='https://arxiv.org/abs/1705.06830'>
Exploring the structure of a real-time, arbitrary neural artistic stylization network
</a>
<br>
<center><img src='https://visitor-badge.glitch.me/badge?page_id=dj_arbi_img_stylization' alt='visitor badge'></center>
"""

content_input = gr.inputs.Image(label="Content Image", source="upload")
style_input = gr.inputs.Image(label="Style Image", source="upload")

app_interface = gr.Interface(fn=do_style_transfer,
                             inputs=[content_image_input, style_image_input],
                             outputs="image",
                             title=title,
                             description=description,
                             examples=[golden_gate,joshua_tree,avatar,joker,einstein,monalisa,paris],
                             article=article
                            )
app_interface.launch()