Update app.py

app.py

@@ -4,9 +4,6 @@ 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):
@@ -187,88 +184,24 @@ 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 img_input is None:
-        raise ValueError("Image input is None")
-    
-    if isinstance(img_input, dict):
-        img = img_input.get('composite')
-        if img is None:
-            img = img_input.get('background')
-    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)
+def get_image(img_path: str, mask: bool=False, scale: bool=True) -> np.array:
+    """
+    Gets image in appropriate format
+    """
+    if isinstance(img_path, np.ndarray):
+        img = img_path
     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
-
-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")
+        img = cv2.imread(img_path)
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)  # Convert BGR to RGB for file inputs
     
-    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
+        _, binary_mask = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY)
+        return np.where(binary_mask == 255, 1, 0)
     
-    if scale and img.dtype != np.float64:
-        return img.astype('float64') / 255.0
+    if scale:
+        return img.astype('double') / 255.0
     
     return img
 
@@ -277,11 +210,6 @@ def blend_images(bg_img, obj_img, mask_img, method):
     obj_img = get_image(obj_img)
     mask_img = get_image(mask_img, mask=True)
 
-    # Ensure mask and images have the same dimensions
-    h, w = bg_img.shape[:2]
-    obj_img = cv2.resize(obj_img, (w, h))
-    mask_img = cv2.resize(mask_img, (w, h))
-
     if method == "Poisson":
         blend_img = np.zeros_like(bg_img)
         for b in range(3):
@@ -296,52 +224,35 @@ def blend_images(bg_img, obj_img, mask_img, method):
     
     return (blend_img * 255).astype(np.uint8)
 
-def update_mask_preview(image):
-    mask = create_mask(image)
-    return mask
-
 with gr.Blocks(theme='bethecloud/storj_theme') as iface:
     gr.HTML("<h1>Image Blending with Multiple Methods</h1>")
     
     with gr.Row():
-        bg_img = gr.Image(label="Background Image", type="numpy")
-        obj_img = gr.Image(label="Object Image", type="numpy")
-    
+        with gr.Column():
+            bg_img = gr.Image(label="Background Image", type="numpy", height=300)
+        with gr.Column():
+            obj_img = gr.Image(label="Object Image", type="numpy", height=300)
+        with gr.Column():
+            gr.Image(label="Mask Image", type="numpy", height=300)
     with gr.Row():
-        mask_img = gr.ImageEditor(
-            label="Mask Image",
-            type="numpy",
-            brush=gr.Brush(colors=["#ffffff"], color_mode="fixed"),
-            eraser=gr.Eraser(),
-        )
-        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")
+        with gr.Column():
+            method = gr.Radio(["Poisson", "Mixed Gradient", "Laplacian"], label="Blending Method", value="Poisson")
+        with gr.Column():
+            blend_button = gr.Button("Blend Images")
     
     output_image = gr.Image(label="Blended Image")
     
-    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 {
-            "background": image_path,
-            "layers": [],
-            "composite": image_path
-        }
+    )
     
     gr.Examples(
         examples=[
-            ["img1.jpg", "img2.jpg", create_image_editor_input("mask1.jpg"), "Poisson"],
-            ["img3.jpg", "img4.jpg", create_image_editor_input("mask2.jpg"), "Mixed Gradient"],
-            ["img6.jpg", "img9.jpg", create_image_editor_input("mask3.jpg"), "Laplacian"]
+            ["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,