Fix pad

app.py

@@ -80,45 +80,105 @@ def make_example(image_path: Path, mask_path: Path) -> EditorValue:
     }
 
 
-def remove_padding(image, original_size):
-    original_width, original_height = original_size
-    left = max((image.width - original_width) // 2, 0)
-    top = max((image.height - original_height) // 2, 0)
-    right = left + original_width
-    bottom = top + original_height
-    return image.crop((left, top, right, bottom))
+def pad(
+    image: Image.Image,
+    size: tuple[int, int],
+    method: int = Image.Resampling.BICUBIC,
+    color: str | int | tuple[int, ...] | None = None,
+    centering: tuple[float, float] = (1, 1),
+) -> tuple[Image.Image, tuple[int, int]]:
+    resized = ImageOps.contain(image, size, method)
+    resized_size = resized.size
+    if resized_size == size:
+        out = resized
+    else:
+        out = Image.new(image.mode, size, color)
+        if resized.palette:
+            palette = resized.getpalette()
+            if palette is not None:
+                out.putpalette(palette)
+        if resized.width != size[0]:
+            x = round((size[0] - resized.width) * max(0, max(centering[0], 1)))
+            out.paste(resized, (x, 0))
+        else:
+            y = round((size[1] - resized.height) * max(0, max(centering[1], 1)))
+            out.paste(resized, (0, y))
+    return out, resized_size
+
+
+def unpad(
+    padded_image: Image.Image,
+    padded_size: tuple[int, int],
+    original_size: tuple[int, int],
+    centering: tuple[float, float] = (1, 1),
+    method: int = Image.Resampling.BICUBIC,
+) -> Image.Image:
+    """
+    Remove the padding added by the `pad` function to recover the original resized image.
+
+    Args:
+        padded_image (Image.Image): The padded image.
+        padded_size (tuple[int, int]): The original size of the resized image before padding.
+        centering (tuple[float, float]): The centering used during padding (x, y), defaults to (1, 1).
+
+    Returns:
+        Image.Image: The cropped image matching the original resized dimensions.
+
+    """
+    width, height = padded_image.size
+    padded_width, padded_height = padded_size
+
+    # Calculate the cropping box based on centering
+    left = round((width - padded_width) * centering[0])
+    top = round((height - padded_height) * centering[1])
+    right = left + padded_width
+    bottom = top + padded_height
+
+    # Crop the image to remove the padding
+    cropped_image = padded_image.crop((left, top, right, bottom))
+
+    # Resize the cropped image to match the original size
+    resized_image = cropped_image.resize(original_size, method)
+    return resized_image
 
 
 def adjust_bbox_to_divisible_16(
-    x_min, y_min, x_max, y_max, width, height, padding=MASK_CONTEXT_PADDING
-):
-    # Add padding
+    x_min: int,
+    y_min: int,
+    x_max: int,
+    y_max: int,
+    width: int,
+    height: int,
+    padding: int = MASK_CONTEXT_PADDING,
+) -> tuple[int, int, int, int]:
+    # Add context padding
     x_min = max(x_min - padding, 0)
     y_min = max(y_min - padding, 0)
     x_max = min(x_max + padding, width)
     y_max = min(y_max + padding, height)
 
-    # Calculate current bbox width and height
-    bbox_width = x_max - x_min
-    bbox_height = y_max - y_min
-
     # Ensure bbox dimensions are divisible by 16
-    if bbox_width % 16 != 0:
-        adjustment = 16 - (bbox_width % 16)
-        x_min = max(x_min - adjustment // 2, 0)
-        x_max = min(x_max + adjustment // 2, width)
-
-    if bbox_height % 16 != 0:
-        adjustment = 16 - (bbox_height % 16)
-        y_min = max(y_min - adjustment // 2, 0)
-        y_max = min(y_max + adjustment // 2, height)
-
-    # Ensure bbox is still within bounds
+    def make_divisible_16(val_min, val_max, max_limit):
+        size = val_max - val_min
+        if size % 16 != 0:
+            adjustment = 16 - (size % 16)
+            val_min = max(val_min - adjustment // 2, 0)
+            val_max = min(val_max + adjustment // 2, max_limit)
+        return val_min, val_max
+
+    x_min, x_max = make_divisible_16(x_min, x_max, width)
+    y_min, y_max = make_divisible_16(y_min, y_max, height)
+
+    # Re-check divisibility after bounds adjustment
     x_min = max(x_min, 0)
     y_min = max(y_min, 0)
     x_max = min(x_max, width)
     y_max = min(y_max, height)
 
+    # Final divisibility check (in case constraints pushed it off again)
+    x_min, x_max = make_divisible_16(x_min, x_max, width)
+    y_min, y_max = make_divisible_16(y_min, y_max, height)
+
     return x_min, y_min, x_max, y_max
 
 
@@ -176,18 +236,9 @@ def infer(
         mask_bbox_x_max,
         mask_bbox_y_max,
     ))
-    room_image_cropped = ImageOps.pad(
-        room_image_cropped,
-        (bbox_longest_side, bbox_longest_side),
-        # White padding
-        color=(255, 255, 255),
-        centering=(0.5, 0.5),
-    )
-    room_image_cropped = ImageOps.fit(
+    room_image_cropped, room_image_cropped_size = pad(
         room_image_cropped,
         (max_dimension, max_dimension),
-        method=Image.Resampling.BICUBIC,
-        centering=(0.5, 0.5),
     )
 
     room_mask_cropped = room_mask.crop((
@@ -196,30 +247,17 @@ def infer(
         mask_bbox_x_max,
         mask_bbox_y_max,
     ))
-    room_mask_cropped.save("room_mask_croppedv1.png")
-    room_mask_cropped = ImageOps.pad(
+    room_mask_cropped, _ = pad(
         room_mask_cropped,
         (max_dimension, max_dimension),
-        # White padding
-        color=(255, 255, 255),
-        centering=(0.5, 0.5),
-    )
-    room_mask_cropped = ImageOps.fit(
-        room_mask_cropped,
-        (max_dimension, max_dimension),
-        method=Image.Resampling.BICUBIC,
-        centering=(0.5, 0.5),
     )
 
     room_image_cropped.save("room_image_cropped.png")
     room_mask_cropped.save("room_mask_cropped.png")
 
-    furniture_image = ImageOps.pad(
+    furniture_image, _ = pad(
         furniture_image_input,
         (max_dimension, max_dimension),
-        # White padding
-        color=(255, 255, 255),
-        centering=(0.5, 0.5),
     )
 
     furniture_mask = Image.new("RGB", (max_dimension, max_dimension), (255, 255, 255))
@@ -271,19 +309,9 @@ def infer(
     for image in results_images:
         final_image = room_image.copy()
 
-        # Downscale back to the bbox_longest_side
-        image_generated = image.crop((
-            max_dimension,
-            0,
-            max_dimension * 2,
-            max_dimension,
-        ))
-        image_generated = image_generated.resize(
-            (bbox_longest_side, bbox_longest_side), Image.Resampling.BICUBIC
-        )
-        # Crop back to the bbox (remove the padding)
-        image_generated = remove_padding(
-            image_generated,
+        image_generated = unpad(
+            image,
+            room_image_cropped_size,
             (
                 mask_bbox_x_max - mask_bbox_x_min,
                 mask_bbox_y_max - mask_bbox_y_min,