import cv2
import numpy as np
import paddle
import paddle.nn.functional as F
from Generater import Generater
import gradio as gr


generator = Generater()
# # oslist =os.listdir("纹理")
# # print(oslist)
G_path ='Gmodel_state33003.pdparams'
layer_state_dictg = paddle.load(G_path)
generator.set_state_dict(layer_state_dictg)#导入训练好的参数文件

def style_transfer(content_img,style_img):
    g_input = content_img.astype('float32') / 127.5 - 1  # 归一化
    g_input = g_input[np.newaxis, ...].transpose(0, 3, 1, 2)  # NHWC -> NCHW
    g_input = paddle.to_tensor(g_input)  # numpy -> tensor
    h,w = g_input.shape[-2:]
    p = max([h,w])
    g_input = F.interpolate(g_input,scale_factor=(512/p))

    g_input_s = style_img.astype('float32') / 127.5 - 1  # 归一化
    g_input_s = g_input_s[np.newaxis, ...].transpose(0, 3, 1, 2)  # NHWC -> NCHW
    g_input_s = paddle.to_tensor(g_input_s)  # numpy -> tensor
    h,w = g_input_s.shape[-2:]
    p = max([h,w])
    g_input_s = F.interpolate(g_input_s,scale_factor=(512/p))

    i = paddle.to_tensor([1.])
    g_output = generator(g_input,g_input_s,i)
    g_output = g_output.detach().numpy()                      # tensor -> numpy
    g_output = g_output.transpose(0, 2, 3, 1)[0]             # NCHW -> NHWC
    g_output = (g_output+1) *127.5                        # 反归一化
    g_output = g_output.astype(np.uint8)
    output = g_output
    # cv2.imwrite(os.path.join("./test", str(i.numpy()[0])+'qt.png'), g_output)#保存图片到本地
    return output

interface = gr.Interface(fn=style_transfer, inputs=["image","image"], outputs="image")
interface.launch(share=True)