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='./src/output/models/model227000.ckpt'
out_PATH='./results/test_output.png'
IMAGE_SZ = 128

import time
def run_time(func):
    def inner(model_PATH, img_p):
        back = func(model_PATH, img_p)
        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).resize((128, 128), Image.ANTIALIAS).convert('RGB'))[np.newaxis] / 255.0
    img_p = src.util.preprocess_images_outpainting(img)
    return img_p

def image_to_path(img):
    resize_img = img
    path = out_PATH
    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 = src.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}