import tensorflow.compat.v1 as tf
tf.disable_v2_behavior()
tf.reset_default_graph()
import numpy as np
from PIL import Image
import src.model
import src.util
import os
import sys

model_PATH='/home/jovyan/work/src/output/models/model2000.ckpt'

import time
def run_time(func):
    def inner(model, image, question):
        back = func(model, image, question)
        print("Runned time: {} s".format(round((time.time() - t)/10, 3)))
        return back
    t = time.time()
    return inner

def load_demo_image(in_PATH):
    img = np.array(Image.open(in_PATH).convert('RGB'))[np.newaxis] / 255.0
    img_p = util.preprocess_images_outpainting(img)
    return img_p

def image_to_path(img):
    resize_img = img
    path = uuid.uuid4().hex + '.png'
    resize_img.save(path)
    return path

@run_time
def inference(model_PATH, img_p):
    G_Z = tf.placeholder(tf.float32, shape=[None, IMAGE_SZ, IMAGE_SZ, 4], name='G_Z')
    G_sample = model.generator(G_Z)
    
    saver = tf.train.Saver()
    with tf.Session() as sess:
        saver.restore(sess, model_PATH)
        output, = sess.run([G_sample], feed_dict={G_Z: img_p})
        img_norm = (output[0] * 255.0).astype(np.uint8)
        img = Image.fromarray(img_norm, 'RGB')
        #util.save_image(output[0], out_PATH)
    return img

def handle(conf):
    """
    该方法是部署之后，其他人调用你的服务时候的处理方法。
    请按规范填写参数结构，这样我们就能替你自动生成配置文件，方便其他人的调用。
    范例：
    params['key'] = value # value_type: str # description: some description
    value_type 可以选择：img, video, audio, str, int, float, [int], [str], [float]
    参数请放到params字典中，我们会自动解析该变量。
    """
    base64_str = conf['Photo']
    image = load_demo_image(base64_str)
    res = inference(model_PATH, image)
    image_str = image_to_path(res)
    return {'Output': image_str}