import numpy as np
import torch
import torch.nn as nn
import gradio as gr
from PIL import Image
import torchvision.transforms as transforms
from lle import SYELLENetS

kwargs = {'channels': 12}
model = SYELLENetS(**kwargs)
model.load_state_dict(torch.load('./model_best_slim.pkl', map_location='cpu'))
model.eval()
def predict(input_img, ver):
    input_img = Image.open(input_img)


    # transform = transforms.Compose([transforms.Resize((400,60), Image.BICUBIC)])

    input_img = np.array(input_img).transpose([2, 0, 1])
    input_img = input_img.astype(np.float32) / 255.0
    input_img = torch.from_numpy(input_img).unsqueeze(0)
    with torch.no_grad():
        outputs = model(input_img)
        out_img = (outputs.clip(0, 1)[0] * 255).permute([1, 2, 0]).cpu().numpy().astype(np.uint8)[..., ::-1]
    return transforms.ToPILImage()(out_img)

title="Image to Line Drawings - Complex and Simple Portraits and Landscapes"

examples=['./examples/1.png', './examples/22.png', './examples/23.png', './examples/55.png', './examples/79.png'
]

iface = gr.Interface(predict, inputs=gr.Image(type='filepath'),
                     outputs='image',
                     title=title,
                     examples=examples)

iface.launch()