Update app.py

app.py

@@ -4,6 +4,9 @@ import scipy.sparse as sp
 import scipy.sparse.linalg as splin
 from numba import jit
 import gradio as gr
+from PIL import Image
+from typing import TypedDict
+
 
 @jit(nopython=True)
 def build_poisson_sparse_matrix(ys, xs, im2var, img_s, img_t, mask):
@@ -213,18 +216,65 @@ def get_image(img_input, mask=False, scale=True):
     
     return img
 
+class ImageData(TypedDict):
+    background: np.ndarray
+    layers: list[np.ndarray]
+
+def combine_masks(im_1: np.ndarray, im_2: np.ndarray):
+    if im_1.shape != im_2.shape:
+        raise ValueError("Images must have the same dimensions")
+    foreground_mask = im_2[:, :, 3] > 0
+    im_1[foreground_mask] = im_2[foreground_mask]
+    return im_1
+
+def make_grey(image: np.ndarray, grey_value: int = 128):
+    rgb = image[:, :, :3]
+    alpha = image[:, :, 3]
+    opaque_mask = alpha == 255
+    rgb[opaque_mask] = np.stack((grey_value, grey_value, grey_value), axis=-1)
+    return np.dstack((rgb, alpha))
+
+def create_mask(image: ImageData):
+    bg = image.get("background", None)
+    layers = image.get("layers", [])
+    if bg is None and not layers:
+        raise ValueError("No background or layers provided")
+    mask = layers[0] if layers else bg
+    for layer in layers[1:]:
+        mask = combine_masks(mask, layer)
+    mask = make_grey(mask)
+    return mask
+
+def get_image(img_input, mask=False, scale=True):
+    if img_input is None:
+        raise ValueError("Image input is None")
+    
+    if isinstance(img_input, dict):
+        img = create_mask(img_input)
+    elif isinstance(img_input, np.ndarray):
+        img = img_input
+    elif isinstance(img_input, str):
+        img = cv2.imread(img_input)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+    else:
+        raise ValueError(f"Unsupported image input type: {type(img_input)}")
+
+    if img is None:
+        raise ValueError("Failed to load image")
+
+    if mask:
+        if len(img.shape) == 3:
+            img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
+        return np.where(img > 127, 1, 0).astype(np.uint8)  # Threshold at 127 for the mask
+    
+    if scale and img.dtype != np.float64:
+        return img.astype('float64') / 255.0
+    
+    return img
+
 def blend_images(bg_img, obj_img, mask_img, method):
     bg_img = get_image(bg_img)
     obj_img = get_image(obj_img)
-    
-    # Handle different input types for mask_img
-    if isinstance(mask_img, dict):
-        mask_img = mask_img.get('composite')
-        if mask_img is None:
-            mask_img = mask_img.get('background')
-    elif isinstance(mask_img, str):
-        mask_img = cv2.imread(mask_img, cv2.IMREAD_GRAYSCALE)
-    
     mask_img = get_image(mask_img, mask=True)
 
     # Ensure mask and images have the same dimensions
@@ -246,14 +296,8 @@ def blend_images(bg_img, obj_img, mask_img, method):
     
     return (blend_img * 255).astype(np.uint8)
 
-def predict(im):
-    return im["composite"]
-
-def create_default_mask(img):
-    if img is None:
-        return None
-    gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
-    _, mask = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY)
+def update_mask_preview(image):
+    mask = create_mask(image)
     return mask
 
 with gr.Blocks(theme='bethecloud/storj_theme') as iface:
@@ -267,7 +311,6 @@ with gr.Blocks(theme='bethecloud/storj_theme') as iface:
         mask_img = gr.ImageEditor(
             label="Mask Image",
             type="numpy",
-            crop_size="1:1",
             brush=gr.Brush(colors=["#ffffff"], color_mode="fixed"),
             eraser=gr.Eraser(),
         )
@@ -279,25 +322,13 @@ with gr.Blocks(theme='bethecloud/storj_theme') as iface:
     
     output_image = gr.Image(label="Blended Image")
     
-    def update_mask(obj_image):
-        if obj_image is None:
-            return None
-        mask = create_default_mask(obj_image)
-        return {"background": mask, "layers": [], "composite": mask}
-    
-    obj_img.change(
-        update_mask,
-        inputs=[obj_img],
-        outputs=[mask_img]
-    )
-    
-    mask_img.change(predict, outputs=mask_preview, inputs=mask_img, show_progress="hidden")
+    mask_img.change(update_mask_preview, inputs=mask_img, outputs=mask_preview)
     
     blend_button.click(
         blend_images,
         inputs=[bg_img, obj_img, mask_img, method],
         outputs=output_image
-    )
+    )    
     
     def create_image_editor_input(image_path):
         return {