Update app.py

app.py

@@ -160,15 +160,14 @@ 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 both images have the same number of channels (3 for RGB)
+    if len(img1.shape) == 2:
+        img1 = cv2.cvtColor(img1, cv2.COLOR_GRAY2RGB)
+    if 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):
+    # Ensure mask has 3 channels
+    if len(mask.shape) == 2:
         mask = np.stack((mask,) * 3, axis=-1)
 
     # Resize all images to the same size (use the size of img1)
@@ -185,11 +184,18 @@ def laplacian_blend(img1: np.ndarray, img2: np.ndarray, mask: np.ndarray, depth:
     
     for d in range(depth-1):
         gaussian_img = _low_pass_filter(blended_img, sigma)
-        gaussian_img = cv2.resize(gaussian_img, blended[d].shape[:2])  # Ensure sizes match
+        gaussian_img = cv2.resize(gaussian_img, blended[d].shape[:2])
+        
+        # Ensure both images have 3 channels
+        if len(blended[d].shape) == 2:
+            blended[d] = cv2.cvtColor(blended[d], cv2.COLOR_GRAY2RGB)
+        if len(gaussian_img.shape) == 2:
+            gaussian_img = cv2.cvtColor(gaussian_img, cv2.COLOR_GRAY2RGB)
+        
         reconstructed_img = cv2.add(blended[d], gaussian_img)
         
         blended_img = cv2.pyrUp(reconstructed_img)
-        blended_img = cv2.resize(blended_img, (w, h))  # Resize to original size
+        blended_img = cv2.resize(blended_img, (w, h))
     
     return np.clip(blended_img, 0, 1)