Akash James
Initial commit
ef01fd7
"""
Copyright (C) 2018 NVIDIA Corporation. All rights reserved.
Licensed under the CC BY-NC-SA 4.0 license (https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode).
"""
from __future__ import division
import torch.nn as nn
import scipy.misc
import numpy as np
import scipy.sparse
import scipy.sparse.linalg
from numpy.lib.stride_tricks import as_strided
from PIL import Image
import cv2
class Propagator(nn.Module):
def __init__(self, beta=0.9999):
super(Propagator, self).__init__()
self.beta = beta
def process(self, initImg, contentImg):
if type(contentImg) == str:
content = cv2.imread(contentImg)
content = cv2.cvtColor(np.content, cv2.COLOR_BGR2RGB)
else:
content = np.array(contentImg, dtype=np.float64)
if type(initImg) == str:
B = cv2.imread(initImg)
B = cv2.cvtColor(B, cv2.COLOR_BGR2RGB)
B /= 255
else:
B = np.array(initImg, dtype=np.float64)
B /= 255
h1,w1,k = B.shape
h = h1 - 4
w = w1 - 4
B = B[int((h1-h)/2):int((h1-h)/2+h),int((w1-w)/2):int((w1-w)/2+w),:]
content = cv2.resize(content, (w, h))
B = self.__replication_padding(B,2)
content = self.__replication_padding(content,2)
content = content.astype(np.float64)/255
B = np.reshape(B,(h1*w1,k))
W = self.__compute_laplacian(content)
W = W.tocsc()
dd = W.sum(0)
dd = np.sqrt(np.power(dd,-1))
dd = dd.A.squeeze()
D = scipy.sparse.csc_matrix((dd, (np.arange(0,w1*h1), np.arange(0,w1*h1)))) # 0.026
S = D.dot(W).dot(D)
A = scipy.sparse.identity(w1*h1) - self.beta*S
A = A.tocsc()
solver = scipy.sparse.linalg.factorized(A)
V = np.zeros((h1*w1,k))
V[:,0] = solver(B[:,0])
V[:,1] = solver(B[:,1])
V[:,2] = solver(B[:,2])
V = V*(1-self.beta)
V = V.reshape(h1,w1,k)
V = V[2:2+h,2:2+w,:]
img = Image.fromarray(np.uint8(np.clip(V * 255., 0, 255.)))
return img
# Returns sparse matting laplacian
# The implementation of the function is heavily borrowed from
# https://github.com/MarcoForte/closed-form-matting/blob/master/closed_form_matting.py
# We thank Marco Forte for sharing his code.
def __compute_laplacian(self, img, eps=10**(-7), win_rad=1):
win_size = (win_rad*2+1)**2
h, w, d = img.shape
c_h, c_w = h - 2*win_rad, w - 2*win_rad
win_diam = win_rad*2+1
indsM = np.arange(h*w).reshape((h, w))
ravelImg = img.reshape(h*w, d)
win_inds = self.__rolling_block(indsM, block=(win_diam, win_diam))
win_inds = win_inds.reshape(c_h, c_w, win_size)
winI = ravelImg[win_inds]
win_mu = np.mean(winI, axis=2, keepdims=True)
win_var = np.einsum('...ji,...jk ->...ik', winI, winI)/win_size - np.einsum('...ji,...jk ->...ik', win_mu, win_mu)
inv = np.linalg.inv(win_var + (eps/win_size)*np.eye(3))
X = np.einsum('...ij,...jk->...ik', winI - win_mu, inv)
vals = (1/win_size)*(1 + np.einsum('...ij,...kj->...ik', X, winI - win_mu))
nz_indsCol = np.tile(win_inds, win_size).ravel()
nz_indsRow = np.repeat(win_inds, win_size).ravel()
nz_indsVal = vals.ravel()
L = scipy.sparse.coo_matrix((nz_indsVal, (nz_indsRow, nz_indsCol)), shape=(h*w, h*w))
return L
def __replication_padding(self, arr,pad):
h,w,c = arr.shape
ans = np.zeros((h+pad*2,w+pad*2,c))
for i in range(c):
ans[:,:,i] = np.pad(arr[:,:,i],pad_width=(pad,pad),mode='edge')
return ans
def __rolling_block(self, A, block=(3, 3)):
shape = (A.shape[0] - block[0] + 1, A.shape[1] - block[1] + 1) + block
strides = (A.strides[0], A.strides[1]) + A.strides
return as_strided(A, shape=shape, strides=strides)