Update app.py

app.py

@@ -1,219 +1,217 @@
 import cv2
 import numpy as np
-import scipy as sp
-import scipy.sparse.linalg
+import scipy.sparse as sp
+import scipy.sparse.linalg as splin
+from numba import jit
 import gradio as gr
-import os
 
-def get_image(img, mask=False):
+def get_image(img_path: str, mask: bool=False, scale: bool=True) -> np.array:
     if mask:
-        if isinstance(img, str):
-            img = cv2.imread(img, cv2.IMREAD_GRAYSCALE)
-        elif img.ndim == 3:
-            img = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY)
-        return np.where(img > 127, 1, 0)
-    else:
-        if isinstance(img, str):
-            img = cv2.imread(img)
-            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-        elif img.ndim == 2:
-            img = np.stack((img,)*3, axis=-1)
-        return img.astype('double') / 255.0
-
-def neighbours(i, j, max_i, max_j):
-    pairs = []
-    for n in [-1, 1]:
-        if 0 <= i+n <= max_i:
-            pairs.append((i+n, j))
-        if 0 <= j+n <= max_j:
-            pairs.append((i, j+n))
-    return pairs
+        img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)
+        _, binary_mask = cv2.threshold(img, 127, 255, cv2.THRESH_BINARY)
+        return np.where(binary_mask == 255, 1, 0)
+    
+    if scale:
+        return cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB).astype('double') / 255.0
+    
+    return cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)
 
-def poisson_blend(img_s, mask, img_t):
+@jit(nopython=True)
+def build_poisson_sparse_matrix(ys, xs, im2var, img_s, img_t, mask):
+    nnz = len(ys)
     img_s_h, img_s_w = img_s.shape
+    A_data = np.zeros(16 * nnz, dtype=np.float64)
+    A_rows = np.zeros(16 * nnz, dtype=np.int32)
+    A_cols = np.zeros(16 * nnz, dtype=np.int32)
+    b = np.zeros(4 * nnz, dtype=np.float64)
+    
+    offsets = np.array([(0, 1), (0, -1), (1, 0), (-1, 0)])
+    
+    idx = 0
+    for n in range(nnz):
+        y, x = ys[n], xs[n]
+        for i in range(4):
+            dy, dx = offsets[i]
+            n_y, n_x = y + dy, x + dx
+            e = 4 * n + i
+            
+            if 0 <= n_y < img_s_h and 0 <= n_x < img_s_w:
+                A_data[idx] = 1
+                A_rows[idx] = e
+                A_cols[idx] = im2var[y, x]
+                idx += 1
+                
+                b[e] = img_s[y, x] - img_s[n_y, n_x]
+                
+                if im2var[n_y, n_x] != -1:
+                    A_data[idx] = -1
+                    A_rows[idx] = e
+                    A_cols[idx] = im2var[n_y, n_x]
+                    idx += 1
+                else:
+                    b[e] += img_t[n_y, n_x]
+    
+    return A_data[:idx], A_rows[:idx], A_cols[:idx], b
+
+def poisson_blend_fast_jit(img_s: np.ndarray, mask: np.ndarray, img_t: np.ndarray) -> np.ndarray:
     nnz = np.sum(mask > 0)
-    im2var = np.full(mask.shape, -1, dtype='int32')
+    im2var = np.full(mask.shape, -1, dtype=np.int32)
     im2var[mask > 0] = np.arange(nnz)
     
-    ys, xs = np.where(mask == 1)
-    
-    # Precompute neighbor indices
-    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)
+    ys, xs = np.nonzero(mask)
     
-    # Compute differences
-    d = img_s[ys, xs][:, np.newaxis] - img_s[y_n, x_n]
+    A_data, A_rows, A_cols, b = build_poisson_sparse_matrix(ys, xs, im2var, img_s, img_t, mask)
     
-    # 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.csr_matrix((A_data, (A_rows, A_cols)), shape=(4*nnz, nnz))
+    v = splin.lsqr(A, b)[0]
     
-    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.ravel()] += img_t[y_n, x_n][~mask_n]
-    
-    # Solve the system
-    v = sp.sparse.linalg.lsqr(A, b)[0]
-    
-    # Update the target image
     img_t_out = img_t.copy()
-    img_t_out[ys, xs] = v
+    img_t_out[mask > 0] = v[im2var[mask > 0]]
     
     return np.clip(img_t_out, 0, 1)
 
-def mixed_blend(img_s, mask, img_t):
+@jit(nopython=True)
+def neighbours(i: int, j: int, max_i: int, max_j: int):
+    pairs = []
+    for n in (-1, 1):
+        if 0 <= i+n <= max_i:
+            pairs.append((i+n, j))
+        if 0 <= j+n <= max_j:
+            pairs.append((i, j+n))
+    return pairs
+
+@jit(nopython=True)
+def build_mixed_blend_sparse_matrix(ys, xs, im2var, img_s, img_t, mask):
+    nnz = len(ys)
     img_s_h, img_s_w = img_s.shape
+    A_data = np.zeros(8 * nnz, dtype=np.float64)
+    A_rows = np.zeros(8 * nnz, dtype=np.int32)
+    A_cols = np.zeros(8 * nnz, dtype=np.int32)
+    b = np.zeros(4 * nnz, dtype=np.float64)
+    
+    idx = 0
+    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_data[idx] = 1
+            A_rows[idx] = e
+            A_cols[idx] = im2var[y, x]
+            idx += 1
+            
+            b[e] = d
+            
+            if im2var[n_y, n_x] != -1:
+                A_data[idx] = -1
+                A_rows[idx] = e
+                A_cols[idx] = im2var[n_y, n_x]
+                idx += 1
+            else:
+                b[e] += img_t[n_y, n_x]
+            e += 1
+    
+    return A_data[:idx], A_rows[:idx], A_cols[:idx], b[:e]
+
+def mixed_blend_fast_jit(img_s: np.ndarray, mask: np.ndarray, img_t: np.ndarray) -> np.ndarray:
     nnz = np.sum(mask > 0)
-    im2var = np.full(mask.shape, -1, dtype='int32')
+    im2var = np.full(mask.shape, -1, dtype=np.int32)
     im2var[mask > 0] = np.arange(nnz)
     
-    ys, xs = np.where(mask == 1)
-    
-    # Precompute neighbor indices
-    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]
-    dt = img_t[ys, xs][:, np.newaxis] - img_t[y_n, x_n]
-    
-    # Choose larger gradient
-    d = np.where(np.abs(ds) > np.abs(dt), ds, dt)
-    
-    # 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)
+    ys, xs = np.nonzero(mask)
     
-    A += sp.sparse.csr_matrix((data_n, (rows_n, cols_n)), shape=(4*nnz, nnz))
+    A_data, A_rows, A_cols, b = build_mixed_blend_sparse_matrix(ys, xs, im2var, img_s, img_t, mask)
     
-    b = d.ravel()
-    b[~mask_n.ravel()] += img_t[y_n, x_n][~mask_n]
+    A = sp.csr_matrix((A_data, (A_rows, A_cols)), shape=(len(b), nnz))
     
-    # Solve the system
-    v = sp.sparse.linalg.lsqr(A, b)[0]
+    v = splin.spsolve(A.T @ A, A.T @ b)
     
-    # Update the target image
     img_t_out = img_t.copy()
-    img_t_out[ys, xs] = v
+    img_t_out[mask > 0] = v[im2var[mask > 0]]
     
     return np.clip(img_t_out, 0, 1)
 
-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)
-
-    # Create Laplacian pyramids for images
-    img1_lap_pyramid = _lap_pyramid(img1, depth, sigma)
-    img2_lap_pyramid = _lap_pyramid(img2, depth, sigma)
-
-    # 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)]
+def _2d_gaussian(sigma: float) -> np.ndarray:
+    ksize = np.int64(np.ceil(sigma)*6+1)
+    gaussian_1d = cv2.getGaussianKernel(ksize, sigma)
+    return gaussian_1d * np.transpose(gaussian_1d)
 
-    # 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
+def _low_pass_filter(img: np.ndarray, sigma: float) -> np.ndarray:
+    return cv2.filter2D(img, -1, _2d_gaussian(sigma))
 
-    return np.clip(blended_img, 0, 1)
+def _high_pass_filter(img: np.ndarray, sigma: float) -> np.ndarray:
+    return img - _low_pass_filter(img, sigma)
 
-def load_example_images(bg_path, obj_path, mask_path):
-    bg_img = cv2.imread(bg_path)
-    bg_img = cv2.cvtColor(bg_img, cv2.COLOR_BGR2RGB)
-    
-    obj_img = cv2.imread(obj_path)
-    obj_img = cv2.cvtColor(obj_img, cv2.COLOR_BGR2RGB)
-    
-    mask_img = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
-    mask_img = np.where(mask_img > 127, 255, 0).astype(np.uint8)
+def _gaus_pyramid(img: np.ndarray, depth: int, sigma: int):
+    _im = img.copy()
     
-    return bg_img, obj_img, mask_img
+    pyramid = []
+    for d in range(depth-1):
+        _im = _low_pass_filter(_im.copy(), sigma)
+        pyramid.append(_im)
+        _im = cv2.pyrDown(_im)
+        
+    return pyramid 
 
-# Modify the blend_images function to accept numpy arrays directly
-def blend_images(bg_img, obj_img, mask_img, blend_method):
-    bg_img = get_image(bg_img)
-    obj_img = get_image(obj_img)
-    mask_img = get_image(mask_img, mask=True)
-
-    # Ensure mask is 2D
-    if mask_img.ndim == 3:
-        mask_img = mask_img[:,:,0]  # Take the first channel if it's 3D
-
-    # Resize mask to match object image size
-    mask_img = cv2.resize(mask_img, (obj_img.shape[1], obj_img.shape[0]))
+def _lap_pyramid(img: np.ndarray, depth: int, sigma: int):
+    _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 
 
-    if blend_method == "Poisson":
-        blend_func = poisson_blend
-    elif blend_method == "Mixed Gradient":
-        blend_func = mixed_blend
-    else:  # Laplacian
-        return laplacian_blend(obj_img, bg_img, np.stack((mask_img,)*3, axis=-1), 5, 25)
+def _blend(img1: np.ndarray, img2: np.ndarray, mask: np.ndarray) -> np.ndarray:
+    return img1 * mask + img2 * (1.0 - mask)
 
-    blend_img = np.zeros(bg_img.shape)
-    for b in range(3):
-        blend_img[:,:,b] = blend_func(obj_img[:,:,b], mask_img, bg_img[:,:,b].copy())
+def laplacian_blend(img1: np.ndarray, img2: np.ndarray, mask: np.ndarray, depth: int, sigma: int) -> np.ndarray:
+    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))
+
+    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)
+
+def blend_images(bg_img, obj_img, mask_img, method):
+    # Convert images to the correct format
+    bg_img = cv2.cvtColor(bg_img, cv2.COLOR_RGB2BGR) / 255.0
+    obj_img = cv2.cvtColor(obj_img, cv2.COLOR_RGB2BGR) / 255.0
+    mask_img = cv2.cvtColor(mask_img, cv2.COLOR_RGB2GRAY) / 255.0
+
+    if method == "Poisson":
+        blend_img = np.zeros_like(bg_img)
+        for b in range(3):
+            blend_img[:,:,b] = poisson_blend_fast_jit(obj_img[:,:,b], mask_img, bg_img[:,:,b].copy())
+    elif method == "Mixed Gradient":
+        blend_img = np.zeros_like(bg_img)
+        for b in range(3):
+            blend_img[:,:,b] = mixed_blend_fast_jit(obj_img[:,:,b], mask_img, bg_img[:,:,b].copy())
+    elif method == "Laplacian":
+        mask_stack = np.stack((mask_img,) * 3, axis=-1)
+        blend_img = laplacian_blend(obj_img, bg_img, mask_stack, 5, 25.0)
     
     return (blend_img * 255).astype(np.uint8)