Update app.py

app.py

@@ -5,17 +5,6 @@ import scipy.sparse.linalg as splin
 from numba import jit
 import gradio as gr
 
-def get_image(img_path: str, mask: bool=False, scale: bool=True) -> np.array:
-    if mask:
-        img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
-        _, binary_mask = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY)
-        return np.where(binary_mask == 255, 1, 0)
-    
-    if scale:
-        return cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB).astype('double') / 255.0
-    
-    return cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)
-
 @jit(nopython=True)
 def build_poisson_sparse_matrix(ys, xs, im2var, img_s, img_t, mask):
     nnz = len(ys)
@@ -195,11 +184,30 @@ def laplacian_blend(img1: np.ndarray, img2: np.ndarray, mask: np.ndarray, depth:
         
     return np.clip(imgs[-1], 0, 1)
 
+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:
+        img = cv2.imread(img_path)
+    
+    if mask:
+        if len(img.shape) == 3:
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        _, binary_mask = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY)
+        return np.where(binary_mask == 255, 1, 0)
+    
+    if scale:
+        return cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype('double') / 255.0
+    
+    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+
 def blend_images(bg_img, obj_img, mask_img, method):
-    # Convert images to the correct format
-    bg_img = cv2.cvtColor(bg_img, cv2.COLOR_RGB2BGR) / 255.0
-    obj_img = cv2.cvtColor(obj_img, cv2.COLOR_RGB2BGR) / 255.0
-    mask_img = cv2.cvtColor(mask_img, cv2.COLOR_RGB2GRAY) / 255.0
+    bg_img = get_image(bg_img)
+    obj_img = get_image(obj_img)
+    mask_img = get_image(mask_img, mask=True)
 
     if method == "Poisson":
         blend_img = np.zeros_like(bg_img)
@@ -210,7 +218,7 @@ def blend_images(bg_img, obj_img, mask_img, method):
         for b in range(3):
             blend_img[:,:,b] = mixed_blend_fast_jit(obj_img[:,:,b], mask_img, bg_img[:,:,b].copy())
     elif method == "Laplacian":
-        mask_stack = np.stack((mask_img,) * 3, axis=-1)
+        mask_stack = np.stack((mask_img.astype(float),) * 3, axis=-1)
         blend_img = laplacian_blend(obj_img, bg_img, mask_stack, 5, 25.0)
     
     return (blend_img * 255).astype(np.uint8)