Upload cyclegan_inference.py

cyclegan_inference.py

@@ -0,0 +1,168 @@
+# -*- 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)