different sizes compatibility

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import torch
 from PIL import Image
 from torchvision import transforms
-from utils import normalize_lab, denormalize_lab
+from utils import normalize_lab, denormalize_lab, pad_image
 from model import Generator
 import kornia.color as color
 
@@ -15,15 +15,13 @@ model = model.to(device)
 model.eval()
 
 
-# Define preprocessing transforms
-transform = transforms.Compose([
-    transforms.Resize((256, 256), Image.BICUBIC),
-    transforms.ToTensor(),
-])
-
-
 def preprocess(image):
     image = image.convert('RGB')
+    image = pad_image(image)
+    transform = transforms.Compose([
+        #transforms.Resize((height, width), Image.BICUBIC),
+        transforms.ToTensor(),
+    ])
     image = transform(image)
     image = image.to(device)
     image = color.rgb_to_lab(image)
@@ -33,8 +31,13 @@ def preprocess(image):
     print(L.shape)
     return L.unsqueeze(0)
 
+def crop_to_original_size(image, original_size):
+    width, height = original_size
+    return transforms.functional.crop(image, top=0, left=0, height=height, width=width)
+
 
 def predict(image):
+    original_size = image.size
     L = preprocess(image)
     with torch.no_grad():
         output = model(L)
@@ -42,6 +45,7 @@ def predict(image):
     L, ab = denormalize_lab(L, output)
     output = torch.cat([L, ab], dim=1)
     output = color.lab_to_rgb(output)
+    output = crop_to_original_size(output, original_size)
     image = transforms.ToPILImage()(output.squeeze().cpu())
 
     return image

utils.py

@@ -1,3 +1,5 @@
+from torchvision import transforms
+
 def normalize_lab(L, ab):
     """
     Normalize the L and ab channels of an image in Lab color space.
@@ -15,4 +17,33 @@ def denormalize_lab(L, ab):
     """
     L = (L + 1) * 50.
     ab = ab * 110.
-    return L, ab
+    return L, ab
+
+def decide_size(image):
+    height = image.size[1]
+    width = image.size[0]
+
+    new_height = 2
+    new_width = 2
+
+    while new_height < height:
+        new_height *= 2
+    while new_width < width:
+        new_width *= 2
+
+    return new_height, new_width
+
+def pad_image(image):
+    height = image.size[1]
+    width = image.size[0]
+
+    new_height, new_width = decide_size(image)
+
+    pad_height = new_height - height
+    pad_width = new_width - width
+
+    padding = (0, 0, pad_width, pad_height)
+
+    image = transforms.Pad(padding, padding_mode='reflect')(image)
+
+    return image