Support background replacement

app.py

@@ -44,9 +44,49 @@ def update_trimap(foreground_mask: dict[str, np.ndarray], unknown_mask: dict[str
     return trimap
 
 
+def adjust_background_image(background_image: PIL.Image.Image, target_size: tuple[int, int]) -> PIL.Image.Image:
+    target_w, target_h = target_size
+    bg_w, bg_h = background_image.size
+
+    scale = max(target_w / bg_w, target_h / bg_h)
+    new_bg_w = int(bg_w * scale)
+    new_bg_h = int(bg_h * scale)
+    background_image = background_image.resize((new_bg_w, new_bg_h))
+    left = (new_bg_w - target_w) // 2
+    top = (new_bg_h - target_h) // 2
+    right = left + target_w
+    bottom = top + target_h
+    background_image = background_image.crop((left, top, right, bottom))
+    return background_image
+
+
+def replace_background(
+    image: PIL.Image.Image, alpha: np.ndarray, background_image: PIL.Image.Image | None
+) -> PIL.Image.Image | None:
+    if background_image is None:
+        return None
+
+    if image.mode != "RGB":
+        raise gr.Error("Image must be RGB.")
+
+    background_image = background_image.convert("RGB")
+    background_image = adjust_background_image(background_image, image.size)
+
+    image = np.array(image).astype(float) / 255
+    background_image = np.array(background_image).astype(float) / 255
+    result = image * alpha[:, :, None] + background_image * (1 - alpha[:, :, None])
+    result = (result * 255).astype(np.uint8)
+    return result
+
+
 @spaces.GPU
 @torch.inference_mode()
-def run(image: PIL.Image.Image, trimap: PIL.Image.Image) -> tuple[PIL.Image.Image, PIL.Image.Image]:
+def run(
+    image: PIL.Image.Image,
+    trimap: PIL.Image.Image,
+    apply_background_replacement: bool,
+    background_image: PIL.Image.Image | None,
+) -> tuple[np.ndarray, PIL.Image.Image, PIL.Image.Image | None]:
     if image.size != trimap.size:
         raise gr.Error("Image and trimap must have the same size.")
     if max(image.size) > MAX_IMAGE_SIZE:
@@ -67,7 +107,12 @@ def run(image: PIL.Image.Image, trimap: PIL.Image.Image) -> tuple[PIL.Image.Imag
     foreground = (foreground * 255).astype(np.uint8)
     foreground = PIL.Image.fromarray(foreground)
 
-    return alpha, foreground
+    if apply_background_replacement:
+        res_bg_replacement = replace_background(image, alpha, background_image)
+    else:
+        res_bg_replacement = None
+
+    return alpha, foreground, res_bg_replacement
 
 
 with gr.Blocks(css="style.css") as demo:
@@ -104,19 +149,33 @@ with gr.Blocks(css="style.css") as demo:
                             height=500,
                         )
                         set_trimap_button = gr.Button("Set trimap")
+                apply_background_replacement = gr.Checkbox(label="Apply background replacement", checked=False)
+                background_image = gr.Image(label="Background image", type="pil", height=500, visible=False)
                 run_button = gr.Button("Run")
         with gr.Column():
             with gr.Box():
                 out_alpha = gr.Image(label="Alpha", height=500)
                 out_foreground = gr.Image(label="Foreground", height=500)
-
+                out_background_replacement = gr.Image(label="Background replacement", height=500, visible=False)
+
+    inputs = [
+        image,
+        trimap,
+        apply_background_replacement,
+        background_image,
+    ]
+    outputs = [
+        out_alpha,
+        out_foreground,
+        out_background_replacement,
+    ]
     gr.Examples(
         examples=[
-            ["assets/bulb_rgb.png", "assets/bulb_trimap.png"],
-            ["assets/retriever_rgb.png", "assets/retriever_trimap.png"],
+            ["assets/retriever_rgb.png", "assets/retriever_trimap.png", False, None],
+            ["assets/bulb_rgb.png", "assets/bulb_trimap.png", True, "assets/new_bg.jpg"],
         ],
-        inputs=[image, trimap],
-        outputs=[out_alpha, out_foreground],
+        inputs=inputs,
+        outputs=outputs,
         fn=run,
         cache_examples=os.getenv("CACHE_EXAMPLES") == "1",
     )
@@ -141,10 +200,18 @@ with gr.Blocks(css="style.css") as demo:
         queue=False,
         api_name=False,
     )
+    apply_background_replacement.change(
+        fn=lambda checked: (gr.update(visible=checked), gr.update(visible=checked)),
+        inputs=apply_background_replacement,
+        outputs=[background_image, out_background_replacement],
+        queue=False,
+        api_name=False,
+    )
+
     run_button.click(
         fn=run,
-        inputs=[image, trimap],
-        outputs=[out_alpha, out_foreground],
+        inputs=inputs,
+        outputs=outputs,
         api_name="run",
     )