Update app.py

app.py

@@ -160,6 +160,17 @@ def _blend(img1: np.ndarray, img2: np.ndarray, mask: np.ndarray) -> np.ndarray:
     return img1 * mask + img2 * (1.0 - mask)
 
 def laplacian_blend(img1: np.ndarray, img2: np.ndarray, mask: np.ndarray, depth: int, sigma: int) -> np.ndarray:
+    # Ensure both images have the same number of channels
+    if len(img1.shape) != len(img2.shape):
+        if len(img1.shape) == 2:
+            img1 = cv2.cvtColor(img1, cv2.COLOR_GRAY2RGB)
+        elif len(img2.shape) == 2:
+            img2 = cv2.cvtColor(img2, cv2.COLOR_GRAY2RGB)
+    
+    # Ensure mask has the same number of channels as the images
+    if len(mask.shape) != len(img1.shape):
+        mask = np.stack((mask,) * 3, axis=-1)
+
     mask_gaus_pyramid = _gaus_pyramid(mask, depth, sigma)
     img1_lap_pyramid, img2_lap_pyramid = _lap_pyramid(img1, depth, sigma), _lap_pyramid(img2, depth, sigma)