import functools
import os

from matplotlib import gridspec
import matplotlib.pylab as plt
import numpy as np
import tensorflow as tf
import tensorflow_hub as hub
import gradio as gr


def crop_center(image):
  """Returns a cropped square image."""
  shape = image.shape
  new_shape = min(shape[1], shape[2])
  offset_y = max(shape[1] - shape[2], 0) // 2
  offset_x = max(shape[2] - shape[1], 0) // 2
  image = tf.image.crop_to_bounding_box(
      image, offset_y, offset_x, new_shape, new_shape)
  return image

@functools.lru_cache(maxsize=None)
def load_image(img, image_size=(256, 256), preserve_aspect_ratio=True):
  """Loads and preprocesses images."""
  img = plt.imread(img.name).astype(np.float32)[np.newaxis, ...]
  # Load and convert to float32 numpy array, add batch dimension, and normalize to range [0, 1].
  if img.max() > 1.0:
    img = img / 255.
  if len(img.shape) == 3:
    img = tf.stack([img, img, img], axis=-1)
  img = crop_center(img)
  img = tf.image.resize(img, image_size, preserve_aspect_ratio=True)
  return img

def show_n(images, titles=('',)):
  n = len(images)
  image_sizes = [image.shape[1] for image in images]
  w = (image_sizes[0] * 6) // 320
  plt.figure(figsize=(w  * n, w))
  gs = gridspec.GridSpec(1, n, width_ratios=image_sizes)
  for i in range(n):
    plt.subplot(gs[i])
    plt.imshow(images[i][0], aspect='equal')
    plt.axis('tight')
    plt.axis('off')
    plt.tight_layout(pad=0)
    plt.title(titles[i] if len(titles) > i else '')
    return plt

output_image_size = 384  

content_img_size = (output_image_size, output_image_size)

style_img_size = (256, 256) 


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

def style(content_image, style_image):
  content_image = load_image(content_image, content_img_size)
  style_image = load_image(style_image, style_img_size)
  style_image = tf.nn.avg_pool(style_image, ksize=[3,3], strides=[1,1], padding='SAME')
  outputs = hub_module(tf.constant(content_image), tf.constant(style_image))
  stylized_image = outputs[0]
  return show_n([stylized_image])

inputs = [
  gr.inputs.Image(type='file', label="Content Image"),
  gr.inputs.Image(type='file', label="Style Image")
]
outputs =  gr.outputs.Image(type='plot', label="Output Image")

title = "Fast Style Transfer for Arbitrary Styles"
description = "demo for Fast Style Transfer for Arbitrary Styles. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1705.06830'>Exploring the structure of a real-time, arbitrary neural artistic stylization network</a> | <a href='https://github.com/tensorflow/hub'>Github Repo</a></p>"

examples = [
    ['bridge.jpeg', 'wave.jpeg']
]

gr.Interface(style, inputs, outputs, title=title, description=description, article=article, examples=examples).launch()
import functools
import os

from matplotlib import gridspec
import matplotlib.pylab as plt
import numpy as np
import tensorflow as tf
import tensorflow_hub as hub
import gradio as gr


def crop_center(image):
  """Returns a cropped square image."""
  shape = image.shape
  new_shape = min(shape[1], shape[2])
  offset_y = max(shape[1] - shape[2], 0) // 2
  offset_x = max(shape[2] - shape[1], 0) // 2
  image = tf.image.crop_to_bounding_box(
      image, offset_y, offset_x, new_shape, new_shape)
  return image

@functools.lru_cache(maxsize=None)
def load_image(img, image_size=(256, 256), preserve_aspect_ratio=True):
  """Loads and preprocesses images."""
  img = plt.imread(img.name).astype(np.float32)[np.newaxis, ...]
  # Load and convert to float32 numpy array, add batch dimension, and normalize to range [0, 1].
  if img.max() > 1.0:
    img = img / 255.
  if len(img.shape) == 3:
    img = tf.stack([img, img, img], axis=-1)
  img = crop_center(img)
  img = tf.image.resize(img, image_size, preserve_aspect_ratio=True)
  return img

def show_n(images, titles=('',)):
  n = len(images)
  image_sizes = [image.shape[1] for image in images]
  w = (image_sizes[0] * 6) // 320
  plt.figure(figsize=(w  * n, w))
  gs = gridspec.GridSpec(1, n, width_ratios=image_sizes)
  for i in range(n):
    plt.subplot(gs[i])
    plt.imshow(images[i][0], aspect='equal')
    plt.axis('tight')
    plt.axis('off')
    plt.tight_layout(pad=0)
    plt.title(titles[i] if len(titles) > i else '')
    return plt

output_image_size = 384  

content_img_size = (output_image_size, output_image_size)

style_img_size = (256, 256) 


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

def style(content_image, style_image):
  content_image = load_image(content_image, content_img_size)
  style_image = load_image(style_image, style_img_size)
  style_image = tf.nn.avg_pool(style_image, ksize=[3,3], strides=[1,1], padding='SAME')
  outputs = hub_module(tf.constant(content_image), tf.constant(style_image))
  stylized_image = outputs[0]
  return show_n([stylized_image])

inputs = [
  gr.inputs.Image(type='file', label="Content Image"),
  gr.inputs.Image(type='file', label="Style Image")
]
outputs =  gr.outputs.Image(type='plot', label="Output Image")

title = "Fast Style Transfer for Arbitrary Styles"
description = "demo for Fast Style Transfer for Arbitrary Styles. To use it, simply upload your image, or click one of the examples to load them. Read more at the links below."
article = "<p style='text-align: center'><a href='https://arxiv.org/abs/1705.06830'>Exploring the structure of a real-time, arbitrary neural artistic stylization network</a> | <a href='https://github.com/tensorflow/hub'>Github Repo</a></p>"

examples = [
    ['bridge.jpeg', 'wave.jpeg']
]

gr.Interface(style, inputs, outputs, title=title, description=description, article=article, examples=examples).launch()
Returns a cropped square image. Loads and preprocesses images. # Load and convert to float32 numpy array, add batch dimension, and normalize to range [0, 1].