import tensorflow as tf
import numpy as np
from imageio.v2 import imread
import os, glob, cv2, shutil
from super_image import EdsrModel, ImageLoader
from PIL import Image
import gradio as gr

pb = 'dmt.pb'
style_dim = 8
img_size=256

model_scale = EdsrModel.from_pretrained('eugenesiow/edsr-base', scale=2)

def preprocess(img):
  return (img / 255. - 0.5) * 2

def deprocess(img):
  return (img + 1) / 2

def load_image(path):
  img = cv2.resize(imread(path), (img_size, img_size))
  img_ = np.expand_dims(preprocess(img), 0)
  return img / 255., img_

def inference(A,B):
  with tf.Graph().as_default():
      output_graph_def = tf.compat.v1.GraphDef()

      with open(pb, 'rb') as fr:
          output_graph_def.ParseFromString(fr.read())
          tf.import_graph_def(output_graph_def, name='')

      sess = tf.compat.v1.Session()
      sess.run(tf.compat.v1.global_variables_initializer())
      graph = tf.compat.v1.get_default_graph()
      Xs = graph.get_tensor_by_name('decoder_1/g:0')
      X = graph.get_tensor_by_name('X:0')
      Y = graph.get_tensor_by_name('Y:0')
  A_img, A_img_ = load_image(A)
  B_img, B_img_ = load_image(B)
  Xs_ = sess.run(Xs, feed_dict={X: A_img_, Y: B_img_})
  output = deprocess(Xs_)[0]
  output = np.array(np.array(output)*255,dtype=np.uint8)
  # output = cv2.cvtColor(output, cv2.COLOR_RGB2BGR)

  image = Image.fromarray(output)
  inputs = ImageLoader.load_image(image)
  preds = model_scale(inputs)
  print(preds.shape)
  ImageLoader.save_image(preds, 'output/scaled_2x.png')



def makeupTransfer(arr1,arr2):
  print("-"*8)
  shutil.rmtree("input/")
  os.makedirs("input/")
  output1 = cv2.cvtColor(arr1, cv2.COLOR_BGR2RGB)
  output2 = cv2.cvtColor(arr2, cv2.COLOR_BGR2RGB)
  cv2.imwrite("input/original.png",output1)
  cv2.imwrite("input/ref.png",output2)
  no_makeup = "input/original.png"
  makeup = "input/ref.png"
  inference(no_makeup, makeup)
  return cv2.cvtColor(cv2.imread("output/scaled_2x.png"), cv2.COLOR_BGR2RGB)
examples = [
            [,
                'faces/no_makeup/xfsy_0521.png',
                'faces/makeup/vFG756.png'],
            [,
                'faces/no_makeup/xfsy_0068.png',
                'faces/makeup/XMY-136.png']
]

app = gr.Interface(fn=makeupTransfer,
             inputs=[gr.Image(label="Reference Image",type='numpy'),
                     gr.Image(label="Makeup Image",type='numpy')],
             outputs=gr.Image(label="Makeup Transfer Image",type='numpy'),
             title="MakeUp Transfer APP",
             examples=examples
             )

app.launch()