Spaces:
Runtime error
Runtime error
| import cv2 | |
| import numpy as np | |
| import scipy as sp | |
| import scipy.sparse.linalg | |
| import gradio as gr | |
| import os | |
| def get_image(img, mask=False): | |
| 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 | |
| def poisson_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) | |
| 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) | |
| # Compute differences | |
| 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.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.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 | |
| return np.clip(img_t_out, 0, 1) | |
| 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) | |
| 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) | |
| # 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.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() | |
| img_t_out[ys, xs] = v | |
| 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)] | |
| # 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 | |
| return np.clip(blended_img, 0, 1) | |
| 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) | |
| return bg_img, obj_img, mask_img | |
| # 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])) | |
| 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) | |
| 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()) | |
| return (blend_img * 255).astype(np.uint8) | |
| examples = [ | |
| ["img1.jpg", "img2.jpg", "mask1.jpg", "Poisson"], | |
| ["img3.jpg", "img4.jpg", "mask2.jpg", "Mixed Gradient"], | |
| ["img6.jpg", "img9.jpg", "mask3.jpg", "Laplacian"] | |
| ] | |
| iface = gr.Interface( | |
| fn=blend_images, | |
| inputs=[ | |
| gr.Image(label="Background Image", type="numpy"), | |
| gr.Image(label="Object Image", type="numpy"), | |
| gr.Image(label="Mask Image", type="numpy"), | |
| gr.Radio(["Poisson", "Mixed Gradient", "Laplacian"], label="Blending Method") | |
| ], | |
| outputs=gr.Image(label="Blended Image"), | |
| title="Image Blending with Examples", | |
| description="Choose from example images or upload your own to blend using different methods.", | |
| examples=examples, | |
| cache_examples=True | |
| ) | |
| iface.launch() |