Update app.py

app.py

@@ -38,28 +38,29 @@ def poisson_blend(img_s, mask, img_t):
     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)
+    y_n = np.array([ys-1, ys+1, ys, ys])
+    x_n = np.array([xs, xs, xs-1, xs+1])
+    
+    # Clip indices to image boundaries
+    y_n = np.clip(y_n, 0, img_s_h-1)
+    x_n = np.clip(x_n, 0, img_s_w-1)
     
     # Compute differences
-    d = img_s[ys, xs][:, np.newaxis] - img_s[y_n.T, x_n.T].T
+    d = img_s[ys, xs][:, np.newaxis] - img_s[y_n, x_n]
     
     # Construct sparse matrix A and vector b
-    rows = np.arange(4*nnz)
-    cols = np.repeat(im2var[ys, xs], 4)
-    data = np.ones(4*nnz)
-    
-    A = sp.sparse.csr_matrix((data, (rows, cols)), shape=(4*nnz, nnz))
-    
-    mask_n = (im2var[y_n, x_n] != -1)
-    cols_n = im2var[y_n, x_n][mask_n]
-    rows_n = np.arange(4*nnz)[mask_n.ravel()]
-    data_n = -np.ones(cols_n.size)
+    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.tile([1, -1], 4*nnz)
     
-    A += sp.sparse.csr_matrix((data_n, (rows_n, cols_n)), shape=(4*nnz, nnz))
+    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.ravel()] += img_t[y_n, x_n][~mask_n]
+    b[~mask_n] += img_t[y_n.ravel()[~mask_n], x_n.ravel()[~mask_n]]
     
     # Solve the system
     v = sp.sparse.linalg.lsqr(A, b)[0]
@@ -79,8 +80,12 @@ def mixed_blend(img_s, mask, img_t):
     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)
+    y_n = np.array([ys-1, ys+1, ys, ys])
+    x_n = np.array([xs, xs, xs-1, xs+1])
+    
+    # Clip indices to image boundaries
+    y_n = np.clip(y_n, 0, img_s_h-1)
+    x_n = np.clip(x_n, 0, img_s_w-1)
     
     # Compute differences
     ds = img_s[ys, xs][:, np.newaxis] - img_s[y_n, x_n]
@@ -92,7 +97,7 @@ def mixed_blend(img_s, mask, img_t):
     # 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()
+    data = np.tile([1, -1], 4*nnz)
     
     mask_n = (im2var[y_n, x_n] != -1).ravel()
     rows = rows[mask_n]
@@ -101,7 +106,7 @@ def mixed_blend(img_s, mask, img_t):
     
     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]
+    b[~mask_n] += img_t[y_n.ravel()[~mask_n], x_n.ravel()[~mask_n]]
     
     # Solve the system
     v = sp.sparse.linalg.lsqr(A, b)[0]