Update app.py

app.py

@@ -68,42 +68,43 @@ def poisson_blend(img_s, mask, img_t):
 
 def mixed_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)
     
-    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)
     
-    ys, xs = np.where(mask==1) 
-        
-    A = sp.sparse.lil_matrix((4*nnz, nnz))
-    b = np.zeros(4*nnz)
+    # 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)
     
-    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):
-            ds = img_s[y][x] - img_s[n_y][n_x]
-            dt = img_t[y][x] - img_t[n_y][n_x]
-            d = ds if abs(ds) > abs(dt) else dt
-            
-            A[e, im2var[y][x]] = 1
-            b[e] = d
-            
-            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
+    # Compute differences
+    ds = img_s[ys, xs][:, np.newaxis] - img_s[y_n, x_n]
+    dt = img_t[ys, xs][:, np.newaxis] - img_t[y_n, x_n]
     
-    A = sp.sparse.csr_matrix(A)
+    # Choose larger gradient
+    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.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)