adjust_bbox_to_divisible_16

app.py

@@ -89,6 +89,39 @@ def remove_padding(image, original_size):
     return image.crop((left, top, right, bottom))
 
 
+def adjust_bbox_to_divisible_16(
+    x_min, y_min, x_max, y_max, width, height, padding=MASK_CONTEXT_PADDING
+):
+    # Add 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
+    x_min = max(x_min, 0)
+    y_min = max(y_min, 0)
+    x_max = min(x_max, width)
+    y_max = min(y_max, height)
+
+    return x_min, y_min, x_max, y_max
+
+
 @spaces.GPU(duration=150)
 def infer(
     furniture_image_input: Image.Image,
@@ -120,11 +153,17 @@ def infer(
     mask_bbox_x_min, mask_bbox_y_min, mask_bbox_x_max, mask_bbox_y_max = (
         room_mask.getbbox(alpha_only=False)
     )
-    # Add MASK_CONTEXT_PADDING (16 pixels) for the context
-    mask_bbox_x_min = max(mask_bbox_x_min - MASK_CONTEXT_PADDING, 0)
-    mask_bbox_y_min = max(mask_bbox_y_min - MASK_CONTEXT_PADDING, 0)
-    mask_bbox_x_max = min(mask_bbox_x_max + MASK_CONTEXT_PADDING, room_mask.width)
-    mask_bbox_y_max = min(mask_bbox_y_max + MASK_CONTEXT_PADDING, room_mask.height)
+    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,
+            padding=MASK_CONTEXT_PADDING,
+        )
+    )
 
     bbox_longest_side = max(
         mask_bbox_x_max - mask_bbox_x_min,
@@ -175,19 +214,6 @@ def infer(
     room_image_cropped.save("room_image_cropped.png")
     room_mask_cropped.save("room_mask_cropped.png")
 
-    # _room_image = ImageOps.fit(
-    #     _room_image,
-    #     (max_dimension, max_dimension),
-    #     method=Image.Resampling.LANCZOS,
-    #     centering=(0.5, 0.5),
-    # )
-    # _room_image.save("room_image.png")
-    # _room_mask_with_white_background = Image.new(
-    #     "RGB", _room_mask.size, (255, 255, 255)
-    # )
-    # _room_mask_with_white_background.paste(_room_mask, (0, 0), _room_mask)
-    # _room_mask_with_white_background.save("room_mask.png")
-
     furniture_image = ImageOps.pad(
         furniture_image_input,
         (max_dimension, max_dimension),
@@ -196,8 +222,6 @@ def infer(
         centering=(0.5, 0.5),
     )
 
-    # _furniture_image.save("furniture_image.png")
-
     furniture_mask = Image.new("RGB", (max_dimension, max_dimension), (255, 255, 255))
 
     image = Image.new(
@@ -285,7 +309,7 @@ intro_markdown = r"""
     </div>
     <br>
     <div style="display: flex; text-align: center; font-size: 14px; padding-right: 300px; padding-left: 300px;">
-        AnyFurnish is a tool that allows you to generate furniture images using Flux.1 Fill Dev. 
+        AnyFurnish is a tool that allows you to generate furniture images using Flux.1 Fill Dev.
         You can upload a furniture image and a room image, and the tool will generate a new image with the furniture in the room.
     </div>
   </div>
@@ -358,7 +382,6 @@ with gr.Blocks(css=css) as demo:
                 sources=["upload"],
                 image_mode="RGBA",
                 layers=False,
-                # crop_size="1:1",
                 brush=gr.Brush(colors=["#FFFFFF"], color_mode="fixed"),
                 height=500,
             )