Add smooth transition

README.md

@@ -8,8 +8,8 @@ python_version: 3.12
 sdk_version: 5.12.0
 suggested_hardware: a100-large
 app_file: app.py
-fullWidth: true
-header: mini
+# fullWidth: true
+# header: mini
 # models: blanchon/anyfurnish
 # datasets: blanchon/anyfurnish-dataset
 tags:

app.py

@@ -1,3 +1,4 @@
+import math
 import secrets
 from pathlib import Path
 from typing import cast
@@ -25,7 +26,9 @@ MASK_CONTEXT_PADDING = 16 * 8
 if not torch.cuda.is_available():
 
     def _dummy_pipe(image: Image.Image, *args, **kwargs):  # noqa: ARG001
-        return {"images": [image]}
+        # return {"images": [image]}  # noqa: ERA001
+        blue_image = Image.new("RGB", image.size, (0, 0, 255))
+        return {"images": [blue_image]}
 
     pipe = _dummy_pipe
 else:
@@ -98,10 +101,10 @@ def pad(
             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)))
+            x = round((size[0] - resized.width) * max(0, min(centering[0], 1)))
             out.paste(resized, (x, 0))
         else:
-            y = round((size[1] - resized.height) * max(0, max(centering[1], 1)))
+            y = round((size[1] - resized.height) * max(0, min(centering[1], 1)))
             out.paste(resized, (0, y))
     return out, resized_size
 
@@ -113,18 +116,6 @@ def unpad(
     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
 
@@ -210,50 +201,54 @@ def infer(
         raise ValueError(msg)
     room_mask = cast("Image.Image", room_mask)
 
-    mask_bbox_x_min, mask_bbox_y_min, mask_bbox_x_max, mask_bbox_y_max = (
-        room_mask.getbbox(alpha_only=False)
-    )
     mask_bbox_x_min, mask_bbox_y_min, mask_bbox_x_max, mask_bbox_y_max = (
         adjust_bbox_to_divisible_16(
-            mask_bbox_x_min,
-            mask_bbox_y_min,
-            mask_bbox_x_max,
-            mask_bbox_y_max,
-            room_mask.width,
-            room_mask.height,
+            *room_mask.getbbox(alpha_only=False),
+            width=room_mask.width,
+            height=room_mask.height,
             padding=MASK_CONTEXT_PADDING,
         )
     )
 
-    bbox_longest_side = max(
-        mask_bbox_x_max - mask_bbox_x_min,
-        mask_bbox_y_max - mask_bbox_y_min,
-    )
-
     room_image_cropped = room_image.crop((
         mask_bbox_x_min,
         mask_bbox_y_min,
         mask_bbox_x_max,
         mask_bbox_y_max,
     ))
-    room_image_cropped, room_image_cropped_size = pad(
+    room_image_padded, room_image_padded_size = pad(
         room_image_cropped,
         (max_dimension, max_dimension),
     )
 
-    room_mask_cropped = room_mask.crop((
+    # grow_and_blur_mask
+    grow_pixels = 10
+    sigma_grow = grow_pixels / 4
+    kernel_size_grow = math.ceil(sigma_grow * 1.5 + 1)
+    room_mask_grow = room_mask.filter(
+        ImageFilter.MaxFilter(size=2 * kernel_size_grow + 1)
+    )
+
+    blur_pixels = 33
+    sigma_blur = blur_pixels / 4
+    kernel_size_blur = math.ceil(sigma_blur * 1.5 + 1)
+    room_mask_blurred = room_mask_grow.filter(
+        ImageFilter.GaussianBlur(radius=kernel_size_blur)
+    )
+
+    room_mask_cropped = room_mask_blurred.crop((
         mask_bbox_x_min,
         mask_bbox_y_min,
         mask_bbox_x_max,
         mask_bbox_y_max,
     ))
-    room_mask_cropped, _ = pad(
+    room_mask_padded, _ = pad(
         room_mask_cropped,
         (max_dimension, max_dimension),
     )
 
-    room_image_cropped.save("room_image_cropped.png")
-    room_mask_cropped.save("room_mask_cropped.png")
+    room_image_padded.save("room_image_padded.png")
+    room_mask_padded.save("room_mask_padded.png")
 
     furniture_image, _ = pad(
         furniture_image_input,
@@ -269,7 +264,7 @@ def infer(
     )
     # Paste on the center of the image
     image.paste(furniture_image, (0, 0))
-    image.paste(room_image_cropped, (max_dimension, 0))
+    image.paste(room_image_padded, (max_dimension, 0))
 
     mask = Image.new(
         "RGB",
@@ -277,7 +272,7 @@ def infer(
         (255, 255, 255),
     )
     mask.paste(furniture_mask, (0, 0))
-    mask.paste(room_mask_cropped, (max_dimension, 0), room_mask_cropped)
+    mask.paste(room_mask_padded, (max_dimension, 0), room_mask_padded)
     # Invert the mask
     mask = ImageOps.invert(mask)
     # Blur the mask
@@ -311,7 +306,7 @@ def infer(
 
         image_generated = unpad(
             image,
-            room_image_cropped_size,
+            room_image_padded_size,
             (
                 mask_bbox_x_max - mask_bbox_x_min,
                 mask_bbox_y_max - mask_bbox_y_min,
@@ -319,7 +314,11 @@ def infer(
         )
         # Paste the image on the room image as the crop was done
         # on the room image
-        final_image.paste(image_generated, (mask_bbox_x_min, mask_bbox_y_min))
+        final_image.paste(
+            image_generated,
+            (mask_bbox_x_min, mask_bbox_y_min),
+            room_mask_cropped,
+        )
         final_images.append(final_image)
 
     return final_images, seed
@@ -439,7 +438,6 @@ with gr.Blocks(css=css) as demo:
             )
             results = gr.Gallery(
                 label="Results",
-                format="png",
                 show_label=False,
                 columns=2,
                 height=500,