Update app.py

app.py

@@ -137,6 +137,41 @@ class MiniSRGAN(nn.Module):
         
         return x, feat
 
+class TinySRGAN(nn.Module):
+    
+    def __init__(self, img_feat = 3, n_feats = 32, kernel_size = 3, num_block = 6, act = nn.PReLU(), scale=4):
+        super(TinySRGAN, self).__init__()
+        
+        self.conv01 = conv(in_channel = img_feat, out_channel = n_feats, kernel_size = 9, BN = False, act = act)
+        
+        resblocks = [ResBlock(channels = n_feats, kernel_size = 3, act = act) for _ in range(num_block)]
+        self.body = nn.Sequential(*resblocks)
+        
+        self.conv02 = conv(in_channel = n_feats, out_channel = n_feats, kernel_size = 3, BN = True, act = None)
+        
+        if(scale == 4):
+            upsample_blocks = [Upsampler(channel = n_feats, kernel_size = 3, scale = 2, act = act) for _ in range(2)]
+        else:
+            upsample_blocks = [Upsampler(channel = n_feats, kernel_size = 3, scale = scale, act = act)]
+
+        self.tail = nn.Sequential(*upsample_blocks)
+        
+        self.last_conv = conv(in_channel = n_feats, out_channel = img_feat, kernel_size = 3, BN = False, act = nn.Tanh())
+        
+    def forward(self, x):
+        
+        x = self.conv01(x)
+        _skip_connection = x
+        
+        x = self.body(x)
+        x = self.conv02(x)
+        feat = x + _skip_connection
+        
+        x = self.tail(feat)
+        x = self.last_conv(x)
+        
+        return x, feat
+
 
 def build_generator():
     
@@ -262,8 +297,6 @@ def translate_image(image, sharpen, model_name, save):
         model=build_generator().to(device)
         model.load_state_dict(torch.load('./weights/mobile_sr.pt', map_location=torch.device('cpu')))
 
-
-
         low_res = transform(resized_image)
         low_res = low_res.unsqueeze(dim=0).to(device)
         model.eval()
@@ -274,15 +307,6 @@ def translate_image(image, sharpen, model_name, save):
         
         sr_img = Image.fromarray((((fake_imgs[0] + 1) / 2) * 255).astype(np.uint8))
 
-    elif(model_name=='EDSR'):
-        model = EdsrModel.from_pretrained('eugenesiow/edsr-base', scale=4)    
-        inputs = ImageLoader.load_image(resized_image)
-        with torch.no_grad():
-            preds = model(inputs)
-
-        sr_img = preds.data.cpu().numpy()
-        sr_img = sr_img[0].transpose((1, 2, 0)) * 255.0
-        sr_img = Image.fromarray(sr_img.astype(np.uint8))
     elif(model_name=='MiniSRGAN'):
         model = MiniSRGAN().to(device)
         model.load_state_dict(torch.load('./weights/miniSRGAN.pt', map_location=torch.device('cpu')))
@@ -300,9 +324,9 @@ def translate_image(image, sharpen, model_name, save):
         output = output.transpose(1, 2, 0)
         sr_img = Image.fromarray((output * 255.0).astype(np.uint8))
         
-    elif(model_name=='MiniSRResNET'):
-        model = MiniSRGAN().to(device)
-        model.load_state_dict(torch.load('./weights/miniSRResNET.pt', map_location=torch.device('cpu')))
+    elif(model_name=='TinySRGAN'):
+        model = TinySRGAN().to(device)
+        model.load_state_dict(torch.load('./weights/tinySRGAN.pt', map_location=torch.device('cpu')))
         
         inputs = np.array(resized_image)
         inputs = (inputs / 127.5) - 1.0   
@@ -335,13 +359,13 @@ def translate_image(image, sharpen, model_name, save):
             sr_img.save('super_resolved_image.png')
         
         return sr_img
-
+    
 interface = gr.Interface(
     fn=translate_image,
     inputs=[
         gr.Image(type="pil"),
         gr.Checkbox(label="Sharpen Image"),
-        gr.Radio(choices=["MobileSR", "MiniSRGAN", "MiniSRResNET"], label="Select Model", value="MobileSR"),
+        gr.Radio(choices=["MobileSR", "MiniSRGAN", "TinySRGAN"], label="Select Model", value="MobileSR"),
         gr.Radio(choices=["True", "False"], label="Save Output", value="False")
     ],
     outputs=gr.Image(type="pil", label="Translated Image"),
@@ -350,6 +374,5 @@ interface = gr.Interface(
     allow_flagging='never'
 )
 
-
 # Launch the Gradio app
 interface.launch()