Spaces:
Running
on
T4
Running
on
T4
from __future__ import absolute_import, division, print_function, unicode_literals | |
import scipy.ndimage | |
from scipy.sparse.linalg import spsolve | |
from scipy import sparse | |
import scipy.io as sio | |
import numpy as np | |
from PIL import Image | |
import copy | |
import cv2 | |
import os | |
import argparse | |
def sub2ind(pi, pj, imgH, imgW): | |
return pj + pi * imgW | |
def Poisson_blend_img(imgTrg, imgSrc_gx, imgSrc_gy, holeMask, gradientMask=None, edge=None): | |
imgH, imgW, nCh = imgTrg.shape | |
if not isinstance(gradientMask, np.ndarray): | |
gradientMask = np.zeros((imgH, imgW), dtype=np.float32) | |
if not isinstance(edge, np.ndarray): | |
edge = np.zeros((imgH, imgW), dtype=np.float32) | |
# Initialize the reconstructed image | |
imgRecon = np.zeros((imgH, imgW, nCh), dtype=np.float32) | |
# prepare discrete Poisson equation | |
A, b, UnfilledMask = solvePoisson(holeMask, imgSrc_gx, imgSrc_gy, imgTrg, | |
gradientMask, edge) | |
# Independently process each channel | |
for ch in range(nCh): | |
# solve Poisson equation | |
x = scipy.sparse.linalg.lsqr(A, b[:, ch])[0] | |
imgRecon[:, :, ch] = x.reshape(imgH, imgW) | |
# Combined with the known region in the target | |
holeMaskC = np.tile(np.expand_dims(holeMask, axis=2), (1, 1, nCh)) | |
imgBlend = holeMaskC * imgRecon + (1 - holeMaskC) * imgTrg | |
# while((UnfilledMask * edge).sum() != 0): | |
# # Fill in edge pixel | |
# pi = np.expand_dims(np.where((UnfilledMask * edge) == 1)[0], axis=1) # y, i | |
# pj = np.expand_dims(np.where((UnfilledMask * edge) == 1)[1], axis=1) # x, j | |
# | |
# for k in range(len(pi)): | |
# if pi[k, 0] - 1 >= 0: | |
# if (UnfilledMask * edge)[pi[k, 0] - 1, pj[k, 0]] == 0: | |
# imgBlend[pi[k, 0], pj[k, 0], :] = imgBlend[pi[k, 0] - 1, pj[k, 0], :] | |
# UnfilledMask[pi[k, 0], pj[k, 0]] = 0 | |
# continue | |
# if pi[k, 0] + 1 <= imgH - 1: | |
# if (UnfilledMask * edge)[pi[k, 0] + 1, pj[k, 0]] == 0: | |
# imgBlend[pi[k, 0], pj[k, 0], :] = imgBlend[pi[k, 0] + 1, pj[k, 0], :] | |
# UnfilledMask[pi[k, 0], pj[k, 0]] = 0 | |
# continue | |
# if pj[k, 0] - 1 >= 0: | |
# if (UnfilledMask * edge)[pi[k, 0], pj[k, 0] - 1] == 0: | |
# imgBlend[pi[k, 0], pj[k, 0], :] = imgBlend[pi[k, 0], pj[k, 0] - 1, :] | |
# UnfilledMask[pi[k, 0], pj[k, 0]] = 0 | |
# continue | |
# if pj[k, 0] + 1 <= imgW - 1: | |
# if (UnfilledMask * edge)[pi[k, 0], pj[k, 0] + 1] == 0: | |
# imgBlend[pi[k, 0], pj[k, 0], :] = imgBlend[pi[k, 0], pj[k, 0] + 1, :] | |
# UnfilledMask[pi[k, 0], pj[k, 0]] = 0 | |
return imgBlend, UnfilledMask | |
def solvePoisson(holeMask, imgSrc_gx, imgSrc_gy, imgTrg, | |
gradientMask, edge): | |
# UnfilledMask indicates the region that is not completed | |
UnfilledMask_topleft = copy.deepcopy(holeMask) | |
UnfilledMask_bottomright = copy.deepcopy(holeMask) | |
# Prepare the linear system of equations for Poisson blending | |
imgH, imgW = holeMask.shape | |
N = imgH * imgW | |
# Number of unknown variables | |
numUnknownPix = holeMask.sum() | |
# 4-neighbors: dx and dy | |
dx = [1, 0, -1, 0] | |
dy = [0, 1, 0, -1] | |
# 3 | |
# | | |
# 2 -- * -- 0 | |
# | | |
# 1 | |
# | |
# Initialize (I, J, S), for sparse matrix A where A(I(k), J(k)) = S(k) | |
I = np.empty((0, 1), dtype=np.float32) | |
J = np.empty((0, 1), dtype=np.float32) | |
S = np.empty((0, 1), dtype=np.float32) | |
# Initialize b | |
b = np.empty((0, 3), dtype=np.float32) | |
# Precompute unkonwn pixel position | |
pi = np.expand_dims(np.where(holeMask == 1)[0], axis=1) # y, i | |
pj = np.expand_dims(np.where(holeMask == 1)[1], axis=1) # x, j | |
pind = sub2ind(pi, pj, imgH, imgW) | |
# |--------------------| | |
# | y (i) | | |
# | x (j) * | | |
# | | | |
# |--------------------| | |
# p[y, x] | |
qi = np.concatenate((pi + dy[0], | |
pi + dy[1], | |
pi + dy[2], | |
pi + dy[3]), axis=1) | |
qj = np.concatenate((pj + dx[0], | |
pj + dx[1], | |
pj + dx[2], | |
pj + dx[3]), axis=1) | |
# Handling cases at image borders | |
validN = (qi >= 0) & (qi <= imgH - 1) & (qj >= 0) & (qj <= imgW - 1) | |
qind = np.zeros((validN.shape), dtype=np.float32) | |
qind[validN] = sub2ind(qi[validN], qj[validN], imgH, imgW) | |
e_start = 0 # equation counter start | |
e_stop = 0 # equation stop | |
# 4 neighbors | |
I, J, S, b, e_start, e_stop = constructEquation(0, validN, holeMask, gradientMask, edge, imgSrc_gx, imgSrc_gy, imgTrg, pi, pj, pind, qi, qj, qind, I, J, S, b, e_start, e_stop) | |
I, J, S, b, e_start, e_stop = constructEquation(1, validN, holeMask, gradientMask, edge, imgSrc_gx, imgSrc_gy, imgTrg, pi, pj, pind, qi, qj, qind, I, J, S, b, e_start, e_stop) | |
I, J, S, b, e_start, e_stop = constructEquation(2, validN, holeMask, gradientMask, edge, imgSrc_gx, imgSrc_gy, imgTrg, pi, pj, pind, qi, qj, qind, I, J, S, b, e_start, e_stop) | |
I, J, S, b, e_start, e_stop = constructEquation(3, validN, holeMask, gradientMask, edge, imgSrc_gx, imgSrc_gy, imgTrg, pi, pj, pind, qi, qj, qind, I, J, S, b, e_start, e_stop) | |
nEqn = len(b) | |
# Construct the sparse matrix A | |
A = sparse.csr_matrix((S[:, 0], (I[:, 0], J[:, 0])), shape=(nEqn, N)) | |
# Check connected pixels | |
for ind in range(0, len(pi), 1): | |
ii = pi[ind, 0] | |
jj = pj[ind, 0] | |
# check up (3) | |
if ii - 1 >= 0: | |
if UnfilledMask_topleft[ii - 1, jj] == 0 and gradientMask[ii - 1, jj] == 0: | |
UnfilledMask_topleft[ii, jj] = 0 | |
# check left (2) | |
if jj - 1 >= 0: | |
if UnfilledMask_topleft[ii, jj - 1] == 0 and gradientMask[ii, jj - 1] == 0: | |
UnfilledMask_topleft[ii, jj] = 0 | |
for ind in range(len(pi) - 1, -1, -1): | |
ii = pi[ind, 0] | |
jj = pj[ind, 0] | |
# check bottom (1) | |
if ii + 1 <= imgH - 1: | |
if UnfilledMask_bottomright[ii + 1, jj] == 0 and gradientMask[ii, jj] == 0: | |
UnfilledMask_bottomright[ii, jj] = 0 | |
# check right (0) | |
if jj + 1 <= imgW - 1: | |
if UnfilledMask_bottomright[ii, jj + 1] == 0 and gradientMask[ii, jj] == 0: | |
UnfilledMask_bottomright[ii, jj] = 0 | |
UnfilledMask = UnfilledMask_topleft * UnfilledMask_bottomright | |
return A, b, UnfilledMask | |
def constructEquation(n, validN, holeMask, gradientMask, edge, imgSrc_gx, imgSrc_gy, imgTrg, pi, pj, pind, qi, qj, qind, I, J, S, b, e_start, e_stop): | |
# Pixel that has valid neighbors | |
validNeighbor = validN[:, n] | |
# Change the out-of-boundary value to 0, in order to run edge[y,x] | |
# in the next line. It won't affect anything as validNeighbor is saved already | |
qi_tmp = copy.deepcopy(qi) | |
qj_tmp = copy.deepcopy(qj) | |
qi_tmp[np.invert(validNeighbor), n] = 0 | |
qj_tmp[np.invert(validNeighbor), n] = 0 | |
NotEdge = (edge[pi[:, 0], pj[:, 0]] == 0) * (edge[qi_tmp[:, n], qj_tmp[:, n]] == 0) | |
# Have gradient | |
if n == 0: | |
HaveGrad = gradientMask[pi[:, 0], pj[:, 0]] == 0 | |
elif n == 2: | |
HaveGrad = gradientMask[pi[:, 0], pj[:, 0] - 1] == 0 | |
elif n == 1: | |
HaveGrad = gradientMask[pi[:, 0], pj[:, 0]] == 0 | |
elif n == 3: | |
HaveGrad = gradientMask[pi[:, 0] - 1, pj[:, 0]] == 0 | |
# Boundary constraint | |
Boundary = holeMask[qi_tmp[:, n], qj_tmp[:, n]] == 0 | |
valid = validNeighbor * NotEdge * HaveGrad * Boundary | |
J_tmp = pind[valid, :] | |
# num of equations: len(J_tmp) | |
e_stop = e_start + len(J_tmp) | |
I_tmp = np.arange(e_start, e_stop, dtype=np.float32).reshape(-1, 1) | |
e_start = e_stop | |
S_tmp = np.ones(J_tmp.shape, dtype=np.float32) | |
if n == 0: | |
b_tmp = - imgSrc_gx[pi[valid, 0], pj[valid, 0], :] + imgTrg[qi[valid, n], qj[valid, n], :] | |
elif n == 2: | |
b_tmp = imgSrc_gx[pi[valid, 0], pj[valid, 0] - 1, :] + imgTrg[qi[valid, n], qj[valid, n], :] | |
elif n == 1: | |
b_tmp = - imgSrc_gy[pi[valid, 0], pj[valid, 0], :] + imgTrg[qi[valid, n], qj[valid, n], :] | |
elif n == 3: | |
b_tmp = imgSrc_gy[pi[valid, 0] - 1, pj[valid, 0], :] + imgTrg[qi[valid, n], qj[valid, n], :] | |
I = np.concatenate((I, I_tmp)) | |
J = np.concatenate((J, J_tmp)) | |
S = np.concatenate((S, S_tmp)) | |
b = np.concatenate((b, b_tmp)) | |
# Non-boundary constraint | |
NonBoundary = holeMask[qi_tmp[:, n], qj_tmp[:, n]] == 1 | |
valid = validNeighbor * NotEdge * HaveGrad * NonBoundary | |
J_tmp = pind[valid, :] | |
# num of equations: len(J_tmp) | |
e_stop = e_start + len(J_tmp) | |
I_tmp = np.arange(e_start, e_stop, dtype=np.float32).reshape(-1, 1) | |
e_start = e_stop | |
S_tmp = np.ones(J_tmp.shape, dtype=np.float32) | |
if n == 0: | |
b_tmp = - imgSrc_gx[pi[valid, 0], pj[valid, 0], :] | |
elif n == 2: | |
b_tmp = imgSrc_gx[pi[valid, 0], pj[valid, 0] - 1, :] | |
elif n == 1: | |
b_tmp = - imgSrc_gy[pi[valid, 0], pj[valid, 0], :] | |
elif n == 3: | |
b_tmp = imgSrc_gy[pi[valid, 0] - 1, pj[valid, 0], :] | |
I = np.concatenate((I, I_tmp)) | |
J = np.concatenate((J, J_tmp)) | |
S = np.concatenate((S, S_tmp)) | |
b = np.concatenate((b, b_tmp)) | |
S_tmp = - np.ones(J_tmp.shape, dtype=np.float32) | |
J_tmp = qind[valid, n, None] | |
I = np.concatenate((I, I_tmp)) | |
J = np.concatenate((J, J_tmp)) | |
S = np.concatenate((S, S_tmp)) | |
return I, J, S, b, e_start, e_stop | |
def getUnfilledMask(holeMask, gradientMask): | |
# UnfilledMask indicates the region that is not completed | |
UnfilledMask_topleft = copy.deepcopy(holeMask) | |
UnfilledMask_bottomright = copy.deepcopy(holeMask) | |
# Get the shape information of the mask | |
imgH, imgW = holeMask.shape | |
# Precompute the unknown pixel position | |
pi = np.expand_dims(np.where(holeMask == 1)[0], axis=1) | |
pj = np.expand_dims(np.where(holeMask == 1)[1], axis=1) | |
# Check connected pixels | |
for ind in range(0, len(pi), 1): | |
ii = pi[ind, 0] | |
jj = pj[ind, 0] | |
# check up (3) | |
if ii - 1 >= 0: | |
if UnfilledMask_topleft[ii - 1, jj] == 0 and gradientMask[ii - 1, jj] == 0: | |
UnfilledMask_topleft[ii, jj] = 0 | |
# check left (2) | |
if jj - 1 >= 0: | |
if UnfilledMask_topleft[ii, jj - 1] == 0 and gradientMask[ii, jj - 1] == 0: | |
UnfilledMask_topleft[ii, jj] = 0 | |
for ind in range(len(pi) - 1, -1, -1): | |
ii = pi[ind, 0] | |
jj = pj[ind, 0] | |
# check bottom (1) | |
if ii + 1 <= imgH - 1: | |
if UnfilledMask_bottomright[ii + 1, jj] == 0 and gradientMask[ii, jj] == 0: | |
UnfilledMask_bottomright[ii, jj] = 0 | |
# check right (0) | |
if jj + 1 <= imgW - 1: | |
if UnfilledMask_bottomright[ii, jj + 1] == 0 and gradientMask[ii, jj] == 0: | |
UnfilledMask_bottomright[ii, jj] = 0 | |
UnfilledMask = UnfilledMask_topleft * UnfilledMask_bottomright | |
return UnfilledMask | |