# -*- coding: utf-8 -*-
import IPython.display as display
import tensorflow as tf
import matplotlib.pyplot as plt
import matplotlib as mpl
import tensorflow_hub as hub
mpl.rcParams['figure.figsize'] = (12,12)
mpl.rcParams['axes.grid'] = False

import numpy as np
import PIL.Image
import time
import functools

def imshow(image, title=None):
  if len(image.shape) > 3:
    image = tf.squeeze(image, axis=0)

  plt.imshow(image)
  if title:
    plt.title(title)

def load_img(path_to_img):
  max_dim = 512
  img = tf.io.read_file(path_to_img)
  print(img)
  img = tf.image.decode_image(img, channels=3)
  img = tf.image.convert_image_dtype(img, tf.float32)

  shape = tf.cast(tf.shape(img)[:-1], tf.float32)
  long_dim = max(shape)
  scale = max_dim / long_dim

  new_shape = tf.cast(shape * scale, tf.int32)

  img = tf.image.resize(img, new_shape)
  img = img[tf.newaxis, :]
  return img

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)

content_path = tf.keras.utils.get_file('YellowLabradorLooking_new.jpg', 'https://storage.googleapis.com/download.tensorflow.org/example_images/YellowLabradorLooking_new.jpg')
style_path = tf.keras.utils.get_file('kandinsky5.jpg','https://storage.googleapis.com/download.tensorflow.org/example_images/Vassily_Kandinsky%2C_1913_-_Composition_7.jpg')
content_image = load_img(content_path)
style_image = load_img(style_path)

plt.subplot(1, 2, 1)
imshow(content_image, 'Content Image')

plt.subplot(1, 2, 2)
imshow(style_image, 'Style Image')

style_path

hub_model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
#hub_model = hub.load('https://hub.tensorflow.google.cn/sayakpaul/lite-model/arbitrary-image-stylization-inceptionv3-dynamic-shapes/dr/predict/1')

stylized_image = hub_model(tf.constant(content_image), tf.constant(style_image))[0]
tensor_to_image(stylized_image)


import numpy as np

import gradio as gr

def sepia(input_path,input_path2):
  input_img = load_img(input_path)
  input_img2 = load_img(input_path2)

  hub_model = hub.load('https://tfhub.dev/google/magenta/arbitrary-image-stylization-v1-256/2')
  stylized_image = hub_model(tf.constant(input_img), tf.constant(input_img2))[0]
  pre_img = tensor_to_image(stylized_image)
  return pre_img

demo_path="1.jpg"
demo_path2="2.jpg"
demo = gr.Interface(
    sepia,examples=[[demo_path,demo_path2],],
    title="上传两个图片，将图片转成另一图片风格",
    inputs=[gr.inputs.Image(label="待转换图片",type="filepath"),gr.inputs.Image(label="风格图片",type="filepath")],
    outputs=gr.outputs.Image(type="pil", label="输出"),theme="dark-seafoam"
          )

if __name__ == "__main__":
    demo.launch()