Spaces:
Sleeping
Sleeping
| #!/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() | |