Update app.py

app.py

@@ -31,38 +31,39 @@ def neighbours(i, j, max_i, max_j):
 
 def poisson_blend(img_s, mask, img_t):
     img_s_h, img_s_w = img_s.shape
+    nnz = np.sum(mask > 0)
+    im2var = np.full(mask.shape, -1, dtype='int32')
+    im2var[mask > 0] = np.arange(nnz)
+    
+    ys, xs = np.where(mask == 1)
+    
+    # Precompute neighbor indices
+    y_n = np.clip(np.stack([ys-1, ys+1, ys, ys]), 0, img_s_h-1)
+    x_n = np.clip(np.stack([xs, xs, xs-1, xs+1]), 0, img_s_w-1)
     
-    nnz = (mask>0).sum()
-    im2var = -np.ones(mask.shape[0:2], dtype='int32')
-    im2var[mask>0] = np.arange(nnz)
-    
-    ys, xs = np.where(mask==1) 
-        
-    A = sp.sparse.lil_matrix((4*nnz, nnz))
-    b = np.zeros(4*nnz)
-    
-    e = 0
-    for n in range(nnz):
-        y, x = ys[n], xs[n]  
-        
-        for n_y, n_x in neighbours(y, x, img_s_h-1, img_s_w-1):
-            A[e, im2var[y][x]] = 1
-            b[e] = img_s[y][x] - img_s[n_y][n_x]
-            
-            if im2var[n_y][n_x] != -1:
-                A[e, im2var[n_y][n_x]] = -1
-            else:
-                b[e] += img_t[n_y][n_x]
-            e += 1
-    
-    A = sp.sparse.csr_matrix(A)
+    # Compute differences
+    d = img_s[ys, xs][:, np.newaxis] - img_s[y_n, x_n]
+    
+    # Construct sparse matrix A and vector b
+    rows = np.repeat(np.arange(4*nnz), 2)
+    cols = np.column_stack([np.repeat(im2var[ys, xs], 4), im2var[y_n, x_n].ravel()])
+    data = np.column_stack([np.ones(4*nnz), -np.ones(4*nnz)]).ravel()
+    
+    mask_n = (im2var[y_n, x_n] != -1).ravel()
+    rows = rows[mask_n]
+    cols = cols[mask_n]
+    data = data[mask_n]
+    
+    A = sp.sparse.csr_matrix((data, (rows, cols)), shape=(4*nnz, nnz))
+    b = d.ravel()
+    b[~mask_n] += img_t[y_n, x_n].ravel()[~mask_n]
+    
+    # Solve the system
     v = sp.sparse.linalg.lsqr(A, b)[0]
     
+    # Update the target image
     img_t_out = img_t.copy()
-    
-    for n in range(nnz):
-        y, x = ys[n], xs[n]
-        img_t_out[y][x] = v[im2var[y][x]]
+    img_t_out[ys, xs] = v
     
     return np.clip(img_t_out, 0, 1)
 
@@ -108,62 +109,61 @@ def mixed_blend(img_s, mask, img_t):
     
     return np.clip(img_t_out, 0, 1)
 
-def _2d_gaussian(sigma):
-    ksize = np.int(np.ceil(sigma)*6+1)
-    gaussian_1d = cv2.getGaussianKernel(ksize, sigma)
-    return gaussian_1d * np.transpose(gaussian_1d)
-
-def _low_pass_filter(img, sigma):
-    return cv2.filter2D(img, -1, _2d_gaussian(sigma))
-
-def _high_pass_filter(img, sigma):
-    return img - _low_pass_filter(img, sigma)
-
-def _gaus_pyramid(img, depth, sigma):
-    _im = img.copy()
-    pyramid = []
-    for d in range(depth-1):
-        _im = _low_pass_filter(_im.copy(), sigma)
-        pyramid.append(_im)
-        _im = cv2.pyrDown(_im)
-    return pyramid 
-
-def _lap_pyramid(img, depth, sigma):
-    _im = img.copy()
-    pyramid = []
-    for d in range(depth-1):
-        lap = _high_pass_filter(_im.copy(), sigma)
-        pyramid.append(lap)
-        _im = cv2.pyrDown(_im)
-    return pyramid 
-
-def _blend(img1, img2, mask):
-    return img1 * mask + img2 * (1.0 - mask)
-
 def laplacian_blend(img1, img2, mask, depth=5, sigma=25):
+    def _2d_gaussian(sigma):
+        ksize = int(np.ceil(sigma) * 6 + 1)
+        gaussian_1d = cv2.getGaussianKernel(ksize, sigma)
+        return gaussian_1d @ gaussian_1d.T
+
+    def _low_pass_filter(img, sigma):
+        return cv2.filter2D(img, -1, _2d_gaussian(sigma))
+
+    def _high_pass_filter(img, sigma):
+        return img - _low_pass_filter(img, sigma)
+
+    def _gaus_pyramid(img, depth, sigma):
+        pyramid = [img]
+        for _ in range(depth - 1):
+            img = _low_pass_filter(cv2.pyrDown(img), sigma)
+            pyramid.append(img)
+        return pyramid
+
+    def _lap_pyramid(img, depth, sigma):
+        pyramid = []
+        for d in range(depth - 1):
+            next_img = cv2.pyrDown(img)
+            lap = img - cv2.pyrUp(next_img, dstsize=img.shape[:2])
+            pyramid.append(lap)
+            img = next_img
+        pyramid.append(img)
+        return pyramid
+
+    def _blend(img1, img2, mask):
+        return img1 * mask + img2 * (1.0 - mask)
+
+    # Ensure mask is 3D
+    if mask.ndim == 2:
+        mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2)
+
+    # Create Gaussian pyramid for mask
     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]
-    
-    h, w = blended[0].shape[:2]
-    
-    img1 = cv2.resize(img1, (w, h))
-    img2 = cv2.resize(img2, (w, h))
-    mask = cv2.resize(mask, (w, h))
+    # Create Laplacian pyramids for images
+    img1_lap_pyramid = _lap_pyramid(img1, depth, sigma)
+    img2_lap_pyramid = _lap_pyramid(img2, depth, sigma)
 
-    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)
-        
-        imgs.append(reconstructed_img)
-        blanded_img = cv2.pyrUp(reconstructed_img)
-        
-    return np.clip(imgs[-1], 0, 1)
+    # Blend pyramids
+    blended_pyramid = [_blend(img1_lap, img2_lap, mask_gaus) 
+                       for img1_lap, img2_lap, mask_gaus 
+                       in zip(img1_lap_pyramid, img2_lap_pyramid, mask_gaus_pyramid)]
+
+    # Reconstruct image
+    blended_img = blended_pyramid[-1]
+    for lap in reversed(blended_pyramid[:-1]):
+        blended_img = cv2.pyrUp(blended_img, dstsize=lap.shape[:2])
+        blended_img += lap
+
+    return np.clip(blended_img, 0, 1)
 
 def load_example_images(bg_path, obj_path, mask_path):
     bg_img = cv2.imread(bg_path)