Update app.py

app.py

@@ -3,6 +3,7 @@ from PIL import Image
 import torch
 import torchvision.transforms as transforms
 import numpy as np
+import torch.nn.functional as F
 
 from archs.model import UNet
 
@@ -28,14 +29,26 @@ def load_img (filename):
     img_tensor = pil_to_tensor(img)
     return img_tensor
 
+
+def check_image_size(x):
+    _, _, h, w = x.size()
+    mod_pad_h = (32 - h % 32) % 32
+    mod_pad_w = (32 - w % 32) % 32
+    x = F.pad(x, (0, mod_pad_w, 0, mod_pad_h), value = 0)
+    return x
+
+
 def process_img(image):
     img = np.array(image)
     img = img / 255.
     img = img.astype(np.float32)
     y = torch.tensor(img).permute(2,0,1).unsqueeze(0).to(device)
+    _, _, H, W = y.shape
+    y= check_image_size(y)
 
     with torch.no_grad():
         x_hat = model(y)
+        x_hat = x_hat[:, :, :H, :W]
 
     restored_img = x_hat.squeeze().permute(1,2,0).clamp_(0, 1).cpu().detach().numpy()
     restored_img = np.clip(restored_img, 0. , 1.)