Spaces:
Running
Running
File size: 2,132 Bytes
75bdb59 fcb1411 75bdb59 dcf98e8 fcb1411 dcf98e8 fcb1411 dcf98e8 fcb1411 75bdb59 fcb1411 75bdb59 fcb1411 75bdb59 fcb1411 75bdb59 |
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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 |
#!/usr/bin/env python
import gradio as gr
import torch
from torch import Tensor
from torchvision import transforms
import numpy as np
import sys
model = None
def load_model():
global model
model = torch.jit.load("v4_gated_unet-rn18-rn18_mse_bn32_5e-3_1e-5.pt")
def denormalize_and_clip_as_tensor (im: Tensor) -> Tensor:
return torch.clip(im * 0.5 + 0.5, 0.0, 1.0)
def denormalize_and_clip_as_numpy (im: Tensor) -> np.ndarray:
im = im.squeeze(0)
return np.ascontiguousarray(denormalize_and_clip_as_tensor(im).permute(1,2,0).detach().cpu().numpy())
def pad_width (size: int, multiple: int):
return 0 if size % multiple == 0 else multiple - (size%multiple)
# pad image to a multiple of 64
def pad_image(im, multiple=64):
# B,C,H,W
rows = im.shape[2]
cols = im.shape[3]
rows_to_pad = pad_width(rows, multiple)
cols_to_pad = pad_width(cols, multiple)
if rows_to_pad == 0 and cols_to_pad == 0:
return im
return transforms.Pad(padding=(0, 0, cols_to_pad, rows_to_pad), padding_mode='reflect')(im)
def undo_antialiasing(im):
im_torch = torch.from_numpy (im).permute(2,0,1).unsqueeze(0).float() / 255.0
im_torch = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))(im_torch)
initial_rows = im.shape[0]
initial_cols = im.shape[1]
im_torch = pad_image(im_torch)
with torch.no_grad():
output_torch = model(im_torch)
output = denormalize_and_clip_as_numpy(output_torch.rgb)
output = (output*255.99).astype(np.uint8)
output = output[:initial_rows, :initial_cols, :]
return output
load_model()
iface = gr.Interface(fn=undo_antialiasing,
inputs=gr.inputs.Image(),
outputs=gr.outputs.Image(),
examples=[
['examples/opencv.png'],
['examples/matplotlib.png'],
['examples/coco_beach.png'],
['examples/plot-bowling.png']])
# FIXME: add OpenCV-background
# FIXME: add some arXiv test image
iface.launch()
|