Update app.py

app.py

@@ -7,7 +7,7 @@ import torch
 from gradio.themes.utils import sizes
 from matplotlib import pyplot as plt
 from matplotlib.patches import Patch
-from PIL import Image
+from PIL import Image, ImageOps
 from torchvision import transforms
 
 # ----------------- HELPER FUNCTIONS ----------------- #
@@ -161,7 +161,29 @@ transform_fn = transforms.Compose(
 
 # ----------------- CORE FUNCTION ----------------- #
 
+def resize_and_pad(image: Image.Image, target_size=(768, 1024)):
+    img_ratio = image.width / image.height
+    target_ratio = target_size[0] / target_size[1]
+
+    if img_ratio > target_ratio:
+        new_width = target_size[0]
+        new_height = int(target_size[0] / img_ratio)
+    else:
+        new_height = target_size[1]
+        new_width = int(target_size[1] * img_ratio)
+
+    resized_image = image.resize((new_width, new_height), Image.ANTIALIAS)
+
+    delta_w = target_size[0] - new_width
+    delta_h = target_size[1] - new_height
+    padding = (delta_w // 2, delta_h // 2, delta_w - (delta_w // 2), delta_h - (delta_h // 2))
+    padded_image = ImageOps.expand(resized_image, padding, fill="black")
+
+    return padded_image
+
 def segment(image: Image.Image) -> Image.Image:
+    image = resize_and_pad(image, target_size=(768, 1024))
+    
     input_tensor = transform_fn(image).unsqueeze(0)
     preds = run_model(input_tensor, height=image.height, width=image.width)
     mask = preds.squeeze(0).cpu().numpy()