Update app.py

app.py

@@ -49,18 +49,21 @@ def poisson_blend(img_s, mask, img_t):
     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.tile([1, -1], 4*nnz)
-    
-    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.ravel()[~mask_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]
@@ -95,18 +98,21 @@ def mixed_blend(img_s, mask, img_t):
     d = np.where(np.abs(ds) > np.abs(dt), ds, dt)
     
     # 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.tile([1, -1], 4*nnz)
-    
-    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.ravel()[~mask_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]