Update app.py

app.py

@@ -160,44 +160,29 @@ 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 (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 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)
-    h, w = img1.shape[:2]
-    img2 = cv2.resize(img2, (w, h))
-    mask = cv2.resize(mask, (w, h))
-
     mask_gaus_pyramid = _gaus_pyramid(mask, depth, sigma)
     img1_lap_pyramid, img2_lap_pyramid = _lap_pyramid(img1, depth, sigma), _lap_pyramid(img2, depth, sigma)
 
     blended = [_blend(obj, bg, mask) for obj, bg, mask in zip(img1_lap_pyramid, img2_lap_pyramid, mask_gaus_pyramid)][::-1]
     
-    blended_img = _blend(img1, img2, mask)
+    h, w = blended[0].shape[:2]
     
-    for d in range(depth-1):
-        gaussian_img = _low_pass_filter(blended_img, sigma)
-        gaussian_img = cv2.resize(gaussian_img, (blended[d].shape[1], blended[d].shape[0]))
-        
-        # 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)
-        
+    img1 = cv2.resize(img1, (w, h))
+    img2 = cv2.resize(img2, (w, h))
+    mask = cv2.resize(mask, (w, h))
+
+    blanded_img = _blend(img1, img2, mask)
+    blanded_img = cv2.resize(blanded_img, blended[0].shape[:2])
+    
+    imgs = []
+    for d in range(0, depth-1):
+        gaussian_img = _low_pass_filter(blanded_img.copy(), sigma)
         reconstructed_img = cv2.add(blended[d], gaussian_img)
         
-        blended_img = cv2.pyrUp(reconstructed_img)
-        blended_img = cv2.resize(blended_img, (w, h))
-    
-    return np.clip(blended_img, 0, 1)
+        imgs.append(reconstructed_img)
+        blanded_img = cv2.pyrUp(reconstructed_img)
+        
+    return np.clip(imgs[-1], 0, 1)
 
 def get_image(img_path: str, mask: bool=False, scale: bool=True) -> np.array:
     """
@@ -251,7 +236,7 @@ with gr.Blocks(theme='bethecloud/storj_theme') as iface:
             mask_img = gr.Image(label="Mask Image", type="numpy", height=300)
     with gr.Row():
         with gr.Column():
-            method = gr.Radio(["Laplacian", "Mixed Gradient"], label="Blending Method", value="Laplacian")
+            method = gr.Radio(["Laplacian", "Poisson", "Mixed Gradient"], label="Blending Method", value="Laplacian")
         with gr.Column():
             blend_button = gr.Button("Blend Images")
     
@@ -265,7 +250,7 @@ with gr.Blocks(theme='bethecloud/storj_theme') as iface:
     
     gr.Examples(
         examples=[
-            ["img1.jpg", "img2.jpg", "mask1.jpg", "Mixed Gradient"],
+            ["img1.jpg", "img2.jpg", "mask1.jpg", "Poisson"],
             ["img3.jpg", "img4.jpg", "mask2.jpg", "Mixed Gradient"],
             ["img6.jpg", "img9.jpg", "mask3.jpg", "Laplacian"]
         ],