Spaces:
Runtime error
Runtime error
import gradio as gr | |
import torch | |
import torchvision | |
from torch import nn | |
from typing import List | |
def ifnone(a, b): # a fastai-specific (fastcore) function used below, redefined so it's independent | |
"`b` if `a` is None else `a`" | |
return b if a is None else a | |
def convT_norm_relu(ch_in:int, ch_out:int, norm_layer:nn.Module, ks:int=3, stride:int=2, bias:bool=True): | |
return [nn.ConvTranspose2d(ch_in, ch_out, kernel_size=ks, stride=stride, padding=1, output_padding=1, bias=bias), | |
norm_layer(ch_out), nn.ReLU(True)] | |
def pad_conv_norm_relu(ch_in:int, ch_out:int, pad_mode:str, norm_layer:nn.Module, ks:int=3, bias:bool=True, | |
pad=1, stride:int=1, activ:bool=True, init=nn.init.kaiming_normal_, init_gain:int=0.02)->List[nn.Module]: | |
layers = [] | |
if pad_mode == 'reflection': layers.append(nn.ReflectionPad2d(pad)) | |
elif pad_mode == 'border': layers.append(nn.ReplicationPad2d(pad)) | |
p = pad if pad_mode == 'zeros' else 0 | |
conv = nn.Conv2d(ch_in, ch_out, kernel_size=ks, padding=p, stride=stride, bias=bias) | |
if init: | |
if init == nn.init.normal_: | |
init(conv.weight, 0.0, init_gain) | |
else: | |
init(conv.weight) | |
if hasattr(conv, 'bias') and hasattr(conv.bias, 'data'): conv.bias.data.fill_(0.) | |
layers += [conv, norm_layer(ch_out)] | |
if activ: layers.append(nn.ReLU(inplace=True)) | |
return layers | |
class ResnetBlock(nn.Module): | |
"nn.Module for the ResNet Block" | |
def __init__(self, dim:int, pad_mode:str='reflection', norm_layer:nn.Module=None, dropout:float=0., bias:bool=True): | |
super().__init__() | |
assert pad_mode in ['zeros', 'reflection', 'border'], f'padding {pad_mode} not implemented.' | |
norm_layer = ifnone(norm_layer, nn.InstanceNorm2d) | |
layers = pad_conv_norm_relu(dim, dim, pad_mode, norm_layer, bias=bias) | |
if dropout != 0: layers.append(nn.Dropout(dropout)) | |
layers += pad_conv_norm_relu(dim, dim, pad_mode, norm_layer, bias=bias, activ=False) | |
self.conv_block = nn.Sequential(*layers) | |
def forward(self, x): return x + self.conv_block(x) | |
def resnet_generator(ch_in:int, ch_out:int, n_ftrs:int=64, norm_layer:nn.Module=None, | |
dropout:float=0., n_blocks:int=9, pad_mode:str='reflection')->nn.Module: | |
norm_layer = ifnone(norm_layer, nn.InstanceNorm2d) | |
bias = (norm_layer == nn.InstanceNorm2d) | |
layers = pad_conv_norm_relu(ch_in, n_ftrs, 'reflection', norm_layer, pad=3, ks=7, bias=bias) | |
for i in range(2): | |
layers += pad_conv_norm_relu(n_ftrs, n_ftrs *2, 'zeros', norm_layer, stride=2, bias=bias) | |
n_ftrs *= 2 | |
layers += [ResnetBlock(n_ftrs, pad_mode, norm_layer, dropout, bias) for _ in range(n_blocks)] | |
for i in range(2): | |
layers += convT_norm_relu(n_ftrs, n_ftrs//2, norm_layer, bias=bias) | |
n_ftrs //= 2 | |
layers += [nn.ReflectionPad2d(3), nn.Conv2d(n_ftrs, ch_out, kernel_size=7, padding=0), nn.Tanh()] | |
return nn.Sequential(*layers) | |
model = resnet_generator(ch_in=3, ch_out=3, n_ftrs=64, norm_layer=None, dropout=0, n_blocks=9) | |
model.load_state_dict(torch.load('generator.pth',map_location=torch.device('cpu'))) | |
model.eval() | |
totensor = torchvision.transforms.ToTensor() | |
normalize_fn = torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)) | |
topilimage = torchvision.transforms.ToPILImage() | |
def predict(input): | |
im = normalize_fn(totensor(input)) | |
print(im.shape) | |
preds = model(im.unsqueeze(0))/2 + 0.5 | |
print(preds.shape) | |
return topilimage(preds.squeeze(0).detach()) | |
gr_interface = gr.Interface(fn=predict, inputs=gr.inputs.Image(shape=(256, 256)), outputs="image", title='Horse-to-Zebra CycleGAN') | |
gr_interface.launch(inline=False,share=False) | |