Spaces:
Paused
Paused
| # -*- coding: utf-8 -*- | |
| """ | |
| # File name: poisson_blending.py | |
| # Time : 2021/11/5 15:22 | |
| # Author: [email protected] | |
| # Description: | |
| """ | |
| import numpy as np | |
| import scipy.sparse | |
| from scipy.sparse.linalg import spsolve | |
| def laplacian_matrix(n, m): | |
| mat_D = scipy.sparse.lil_matrix((m, m)) | |
| mat_D.setdiag(-1, -1) | |
| mat_D.setdiag(4) | |
| mat_D.setdiag(-1, 1) | |
| mat_A = scipy.sparse.block_diag([mat_D] * n).tolil() | |
| mat_A.setdiag(-1, 1 * m) | |
| mat_A.setdiag(-1, -1 * m) | |
| return mat_A | |
| def poisson_blending(source, target, mask, with_gamma=True): | |
| """ | |
| source: H * W * 3, cv2 image | |
| target: H * W * 3, cv2 image | |
| mask: H * W * 1 | |
| """ | |
| if with_gamma: | |
| gamma_value = 2.2 | |
| else: | |
| gamma_value = 1 | |
| source = source.astype('float') | |
| target = target.astype('float') | |
| source = np.power(source, 1 / gamma_value) | |
| target = np.power(target, 1 / gamma_value) | |
| res = target.copy() | |
| y_range, x_range = source.shape[:2] | |
| mat_A = laplacian_matrix(y_range, x_range) | |
| laplacian = mat_A.tocsc() | |
| mask[mask != 0] = 1 | |
| for y in range(1, y_range - 1): | |
| for x in range(1, x_range - 1): | |
| if mask[y, x] == 0: | |
| k = x + y * x_range | |
| mat_A[k, k] = 1 | |
| mat_A[k, k + 1] = 0 | |
| mat_A[k, k - 1] = 0 | |
| mat_A[k, k + x_range] = 0 | |
| mat_A[k, k - x_range] = 0 | |
| mat_A = mat_A.tocsc() | |
| y_min, y_max = 0, y_range | |
| x_min, x_max = 0, x_range | |
| mask_flat = mask.flatten() | |
| for channel in range(source.shape[2]): | |
| source_flat = source[y_min:y_max, x_min:x_max, channel].flatten() | |
| target_flat = target[y_min:y_max, x_min:x_max, channel].flatten() | |
| # inside the mask: | |
| # \Delta f = div v = \Delta g | |
| alpha = 1 | |
| mat_b = laplacian.dot(source_flat) * alpha | |
| # outside the mask: | |
| # f = t | |
| mat_b[mask_flat == 0] = target_flat[mask_flat == 0] | |
| x = spsolve(mat_A, mat_b) | |
| x = x.reshape((y_range, x_range)) | |
| res[:, :, channel] = x | |
| res = np.power(res, gamma_value) | |
| res[res > 255] = 255 | |
| res[res < 0] = 0 | |
| res = res.astype('uint8') | |
| return res | |