FGT_codes/tool/utils/Poisson_blend.py

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(imgTrg, imgSrc_gx, imgSrc_gy, holeMask, edge=None):

    imgH, imgW, nCh = imgTrg.shape

    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 = solvePoisson(holeMask, imgSrc_gx, imgSrc_gy, imgTrg, edge)

    # Independently process each channel
    for ch in range(nCh):

        # solve Poisson equation
        x = scipy.sparse.linalg.lsqr(A, b[:, ch, None])[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

    # Fill in edge pixel
    pi = np.expand_dims(np.where((holeMask * edge) == 1)[0], axis=1) # y, i
    pj = np.expand_dims(np.where((holeMask * edge) == 1)[1], axis=1) # x, j

    for k in range(len(pi)):
        if pi[k, 0] - 1 >= 0:
            if edge[pi[k, 0] - 1, pj[k, 0]] == 0:
                imgBlend[pi[k, 0], pj[k, 0], :] = imgBlend[pi[k, 0] - 1, pj[k, 0], :]
                continue
        if pi[k, 0] + 1 <= imgH - 1:
            if edge[pi[k, 0] + 1, pj[k, 0]] == 0:
                imgBlend[pi[k, 0], pj[k, 0], :] = imgBlend[pi[k, 0] + 1, pj[k, 0], :]
                continue
        if pj[k, 0] - 1 >= 0:
            if edge[pi[k, 0], pj[k, 0] - 1] == 0:
                imgBlend[pi[k, 0], pj[k, 0], :] = imgBlend[pi[k, 0], pj[k, 0] - 1, :]
                continue
        if pj[k, 0] + 1 <= imgW - 1:
            if edge[pi[k, 0], pj[k, 0] + 1] == 0:
                imgBlend[pi[k, 0], pj[k, 0], :] = imgBlend[pi[k, 0], pj[k, 0] + 1, :]

    return imgBlend

def solvePoisson(holeMask, imgSrc_gx, imgSrc_gy, imgTrg, edge):

    # 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, 2), 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)  *         |
    # |                    |
    # |--------------------|

    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, 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, 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, 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, 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))

    return A, b


def constructEquation(n, validN, holeMask, 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

    # Not edge
    NotEdge = (edge[pi[:, 0], pj[:, 0]] == 0) * (edge[qi_tmp[:, n], qj_tmp[:, n]] == 0)

    # Boundary constraint
    Boundary = holeMask[qi_tmp[:, n], qj_tmp[:, n]] == 0
    valid = validNeighbor * NotEdge * 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 * 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 gradient_mask(mask):  #产生梯度的mask

    gradient_mask = np.logical_or.reduce((mask,
        np.concatenate((mask[1:, :], np.zeros((1, mask.shape[1]), dtype=np.bool)), axis=0),
        np.concatenate((mask[:, 1:], np.zeros((mask.shape[0], 1), dtype=np.bool)), axis=1)))

    return gradient_mask

    
    
if __name__ == '__main__':
    import cvbase
    from skimage.feature import canny
    import argparse
    import imageio

    parser = argparse.ArgumentParser()
    parser.add_argument('--flow', type=str, default='../../test_blending/flow/00000.flo')
    parser.add_argument('--mask', type=str, default='../../test_blending/mask/00000.png')
    parser.add_argument('--width', type=int, default=432)
    parser.add_argument('--height', type=int, default=256)
    parser.add_argument('--output', type=str, default='../../test_blending/ret')
    args = parser.parse_args()
    flow, mask = args.flow, args.mask
    width, height = args.width, args.height
    output = args.output
    
    if not os.path.exists(output):
        os.makedirs(output)

    flow = cvbase.read_flow(flow)
    mask = cv2.imread(mask, 0)
    h, w, c = flow.shape
    flow_resized = np.zeros((height, width, 2))
    flow_resized[:, :, 0] = cv2.resize(flow[:, :, 0], (width, height), cv2.INTER_LINEAR) * width / w
    flow_resized[:, :, 1] = cv2.resize(flow[:, :, 1], (width, height), cv2.INTER_LINEAR) * height / h
    flow = flow_resized

    mask = cv2.resize(mask, (width, height), cv2.INTER_NEAREST)
    mask_gradient = gradient_mask(mask)

    flow_gray = (flow[:, :, 0] ** 2 + flow[:, :, 1] ** 2) ** 0.5
    flow_gray = flow_gray / flow_gray.max()
    edge = canny(flow_gray, sigma=1, low_threshold=0.1, high_threshold=0.2)
    edge = edge.astype(np.bool)
    masked_edge = edge * mask

    # gradients
    gradient_x = np.concatenate((np.diff(flow, axis=1), np.zeros((height, 1, 2), dtype=np.float32)), axis=1)
    gradient_y = np.concatenate((np.diff(flow, axis=0), np.zeros((1, width, 2), dtype=np.float32)), axis=0)
    gradient = np.concatenate((gradient_x, gradient_y), axis=2)
    gradient[mask_gradient, :] = 0  # 把中间的梯度设置成了0

    # complete flow
    imgSrc_gy = gradient[:, :, 2: 4]
    imgSrc_gy = imgSrc_gy[0: h - 1, :, :]
    imgSrc_gx = gradient[:, :, 0: 2]
    imgSrc_gx = imgSrc_gx[:, 0: w - 1, :]
    compFlow = Poisson_blend(flow, imgSrc_gx, imgSrc_gy, mask, masked_edge)  # todo: edge or masked_edge ?

    # save flow
    flow_n = cvbase.flow2rgb(flow)
    compFlow_n = cvbase.flow2rgb(compFlow)
    imageio.imwrite(os.path.join(output, 'flow.png'), flow_n)
    imageio.imwrite(os.path.join(output, 'compFlow.png'), compFlow_n)
    imageio.imwrite(os.path.join(output, 'edge.png'), masked_edge)
    # imageio.imwrite(os.path.join(output, 'gx.png'), imgSrc_gx)
    # imageio.imwrite(os.path.join(output, 'gy.png'), imgSrc_gy)
    imageio.imwrite(os.path.join(output, 'mask.png'), mask)
    imageio.imwrite(os.path.join(output, 'grad0.png'), gradient[:, :, 0])
    imageio.imwrite(os.path.join(output, 'grad1.png'), gradient[:, :, 1])