File size: 1,235 Bytes
dc604d9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | 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() |