Update app.py

app.py

@@ -42,21 +42,24 @@ def poisson_blend(img_s, mask, img_t):
     x_n = np.clip(np.stack([xs, xs, xs-1, xs+1]), 0, img_s_w-1)
     
     # Compute differences
-    d = img_s[ys, xs][:, np.newaxis] - img_s[y_n, x_n]
+    d = img_s[ys, xs][:, np.newaxis] - img_s[y_n.T, x_n.T].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()
-    
-    mask_n = (im2var[y_n, x_n] != -1).ravel()
-    rows = rows[mask_n]
-    cols = cols[mask_n]
-    data = data[mask_n]
+    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)
+    
+    A += sp.sparse.csr_matrix((data_n, (rows_n, cols_n)), shape=(4*nnz, nnz))
+    
     b = d.ravel()
-    b[~mask_n] += img_t[y_n, x_n].ravel()[~mask_n]
+    b[~mask_n.ravel()] += img_t[y_n, x_n][~mask_n]
     
     # Solve the system
     v = sp.sparse.linalg.lsqr(A, b)[0]