# -*- 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)