Update app.py

app.py

@@ -184,11 +184,9 @@ def laplacian_blend(img1: np.ndarray, img2: np.ndarray, mask: np.ndarray, depth:
         
     return np.clip(imgs[-1], 0, 1)
 
-
-
 def get_image(img_input, mask=False, scale=True):
-    if isinstance(img_input, dict) and 'composite' in img_input:
-        img = img_input['composite']
+    if isinstance(img_input, dict):
+        img = img_input.get('composite') or img_input.get('background')
     elif isinstance(img_input, np.ndarray):
         img = img_input
     elif isinstance(img_input, str):
@@ -226,28 +224,32 @@ def blend_images(bg_img, obj_img, mask_img, method):
     
     return (blend_img * 255).astype(np.uint8)
 
+def predict(im):
+    return im["composite"]
+
 with gr.Blocks(theme='bethecloud/storj_theme') as iface:
     gr.HTML("<h1>Image Blending with Multiple Methods</h1>")
     
     with gr.Row():
-        with gr.Column():
-            bg_img = gr.Image(label="Background Image", type="numpy")
-            obj_img = gr.Image(label="Object Image", type="numpy")
-        with gr.Column():
-            mask_img = gr.ImageEditor(
-                label="Mask Image",
-                type="numpy",
-                brush=gr.Brush(colors=["#ffffff"], color_mode="fixed"),
-                eraser=gr.Eraser(),
-                height=300,
-                width=300,
-            )
-            method = gr.Radio(["Poisson", "Mixed Gradient", "Laplacian"], label="Blending Method", value="Poisson")
+        bg_img = gr.Image(label="Background Image", type="numpy")
+        obj_img = gr.Image(label="Object Image", type="numpy")
+    
+    with gr.Row():
+        mask_img = gr.ImageEditor(
+            label="Mask Image",
+            type="numpy",
+            crop_size="1:1",
+        )
+        mask_preview = gr.Image(label="Mask Preview")
+    
+    method = gr.Radio(["Poisson", "Mixed Gradient", "Laplacian"], label="Blending Method", value="Poisson")
     
     blend_button = gr.Button("Blend Images")
     
     output_image = gr.Image(label="Blended Image")
     
+    mask_img.change(predict, outputs=mask_preview, inputs=mask_img, show_progress="hidden")
+    
     blend_button.click(
         blend_images,
         inputs=[bg_img, obj_img, mask_img, method],
@@ -255,11 +257,11 @@ with gr.Blocks(theme='bethecloud/storj_theme') as iface:
     )
     
     gr.Examples(
-    examples=[
-        ["img1.jpg", "img2.jpg", {"composite": "mask1.jpg"}, "Poisson"],
-        ["img3.jpg", "img4.jpg", {"composite": "mask2.jpg"}, "Mixed Gradient"],
-        ["img6.jpg", "img9.jpg", {"composite": "mask3.jpg"}, "Laplacian"]
-    ],
+        examples=[
+            ["img1.jpg", "img2.jpg", "mask1.jpg", "Poisson"],
+            ["img3.jpg", "img4.jpg", "mask2.jpg", "Mixed Gradient"],
+            ["img6.jpg", "img9.jpg", "mask3.jpg", "Laplacian"]
+        ],
         inputs=[bg_img, obj_img, mask_img, method],
         outputs=output_image,
         fn=blend_images,