Spaces:
Runtime error
Runtime error
| # -*- coding: utf-8 -*- | |
| """cyclegan_inference.ipynb | |
| Automatically generated by Colaboratory. | |
| Original file is located at | |
| https://colab.research.google.com/drive/12lelsBZXqNOe7xaXI724rEHAbppRt07y | |
| """ | |
| 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 | |
| class ConvBlock(torch.nn.Module): | |
| def __init__(self,input_size,output_size,kernel_size=4,stride=2,padding=1,activation='relu',batch_norm=True): | |
| super(ConvBlock,self).__init__() | |
| self.conv = torch.nn.Conv2d(input_size,output_size,kernel_size,stride,padding) | |
| self.batch_norm = batch_norm | |
| self.bn = torch.nn.InstanceNorm2d(output_size) | |
| self.activation = activation | |
| self.relu = torch.nn.ReLU(True) | |
| self.lrelu = torch.nn.LeakyReLU(0.2,True) | |
| self.tanh = torch.nn.Tanh() | |
| self.sigmoid = torch.nn.Sigmoid() | |
| def forward(self,x): | |
| if self.batch_norm: | |
| out = self.bn(self.conv(x)) | |
| else: | |
| out = self.conv(x) | |
| if self.activation == 'relu': | |
| return self.relu(out) | |
| elif self.activation == 'lrelu': | |
| return self.lrelu(out) | |
| elif self.activation == 'tanh': | |
| return self.tanh(out) | |
| elif self.activation == 'no_act': | |
| return out | |
| elif self.activation =='sigmoid': | |
| return self.sigmoid(out) | |
| class ResnetBlock(torch.nn.Module): | |
| def __init__(self,num_filter,kernel_size=3,stride=1,padding=0): | |
| super(ResnetBlock,self).__init__() | |
| conv1 = torch.nn.Conv2d(num_filter,num_filter,kernel_size,stride,padding) | |
| conv2 = torch.nn.Conv2d(num_filter,num_filter,kernel_size,stride,padding) | |
| bn = torch.nn.InstanceNorm2d(num_filter) | |
| relu = torch.nn.ReLU(True) | |
| pad = torch.nn.ReflectionPad2d(1) | |
| self.resnet_block = torch.nn.Sequential( | |
| pad, | |
| conv1, | |
| bn, | |
| relu, | |
| pad, | |
| conv2, | |
| bn | |
| ) | |
| def forward(self,x): | |
| out = self.resnet_block(x) | |
| return out | |
| 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) | |
| class DeconvBlock(torch.nn.Module): | |
| def __init__(self,input_size,output_size,kernel_size=4,stride=2,padding=1,activation='relu',batch_norm=True): | |
| super(DeconvBlock,self).__init__() | |
| self.deconv = torch.nn.ConvTranspose2d(input_size,output_size,kernel_size,stride,padding) | |
| self.batch_norm = batch_norm | |
| self.bn = torch.nn.InstanceNorm2d(output_size) | |
| self.activation = activation | |
| self.relu = torch.nn.ReLU(True) | |
| self.tanh = torch.nn.Tanh() | |
| def forward(self,x): | |
| if self.batch_norm: | |
| out = self.bn(self.deconv(x)) | |
| else: | |
| out = self.deconv(x) | |
| if self.activation == 'relu': | |
| return self.relu(out) | |
| elif self.activation == 'lrelu': | |
| return self.lrelu(out) | |
| elif self.activation == 'tanh': | |
| return self.tanh(out) | |
| elif self.activation == 'no_act': | |
| return out | |
| class Generator(torch.nn.Module): | |
| def __init__(self,input_dim,num_filter,output_dim,num_resnet): | |
| super(Generator,self).__init__() | |
| #Reflection padding | |
| #self.pad = torch.nn.ReflectionPad2d(3) | |
| #Encoder | |
| self.conv1 = ConvBlock(input_dim,num_filter,kernel_size=4,stride=2,padding=1) | |
| self.conv2 = ConvBlock(num_filter,num_filter*2) | |
| #self.conv3 = ConvBlock(num_filter*2,num_filter*4) | |
| #Resnet blocks | |
| self.resnet_blocks = [] | |
| for i in range(num_resnet): | |
| self.resnet_blocks.append(ResnetBlock(num_filter*2)) | |
| self.resnet_blocks = torch.nn.Sequential(*self.resnet_blocks) | |
| #Decoder | |
| self.deconv1 = DeconvBlock(num_filter*2,num_filter) | |
| self.deconv2 = DeconvBlock(num_filter,output_dim,activation='tanh') | |
| #self.deconv3 = ConvBlock(num_filter,output_dim,kernel_size=7,stride=1,padding=0,activation='tanh',batch_norm=False) | |
| def forward(self,x): | |
| #Encoder | |
| enc1 = self.conv1(x) | |
| enc2 = self.conv2(enc1) | |
| #enc3 = self.conv3(enc2) | |
| #Resnet blocks | |
| res = self.resnet_blocks(enc2) | |
| #Decoder | |
| dec1 = self.deconv1(res) | |
| dec2 = self.deconv2(dec1) | |
| #out = self.deconv3(self.pad(dec2)) | |
| return dec2 | |
| def normal_weight_init(self,mean=0.0,std=0.02): | |
| for m in self.children(): | |
| if isinstance(m,ConvBlock): | |
| torch.nn.init.normal_(m.conv.weight,mean,std) | |
| if isinstance(m,DeconvBlock): | |
| torch.nn.init.normal_(m.deconv.weight,mean,std) | |
| if isinstance(m,ResnetBlock): | |
| torch.nn.init.normal_(m.conv.weight,mean,std) | |
| torch.nn.init.constant_(m.conv.bias,0) | |
| model = G_A = Generator(3, 32, 3, 4).cuda() # input_dim, num_filter, output_dim, num_resnet | |
| model.load_state_dict(torch.load('G_A_HW4_SAVE.pt',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) |