Update app.py

app.py

@@ -7,8 +7,18 @@ import os
 
 def get_image(img, mask=False):
     if mask:
+        if isinstance(img, str):
+            img = cv2.imread(img, cv2.IMREAD_GRAYSCALE)
+        elif img.ndim == 3:
+            img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
         return np.where(img > 127, 1, 0)
-    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype('double') / 255.0
+    else:
+        if isinstance(img, str):
+            img = cv2.imread(img)
+            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
+        elif img.ndim == 2:
+            img = np.stack((img,)*3, axis=-1)
+        return img.astype('double') / 255.0
 
 def neighbours(i, j, max_i, max_j):
     pairs = []
@@ -172,6 +182,10 @@ def blend_images(bg_img, obj_img, mask_img, blend_method):
     obj_img = get_image(obj_img)
     mask_img = get_image(mask_img, mask=True)
 
+    # Ensure mask is 2D
+    if mask_img.ndim == 3:
+        mask_img = mask_img[:,:,0]  # Take the first channel if it's 3D
+
     # Resize mask to match object image size
     mask_img = cv2.resize(mask_img, (obj_img.shape[1], obj_img.shape[0]))