| 13 | |
| 14 | |
| 15 | class Propagator(nn.Module): |
| 16 | def __init__(self, beta=0.9999): |
| 17 | super(Propagator, self).__init__() |
| 18 | self.beta = beta |
| 19 | |
| 20 | def process(self, initImg, contentImg): |
| 21 | |
| 22 | if type(contentImg) == str: |
| 23 | content = scipy.misc.imread(contentImg, mode='RGB') |
| 24 | else: |
| 25 | content = contentImg.copy() |
| 26 | # content = scipy.misc.imread(contentImg, mode='RGB') |
| 27 | |
| 28 | if type(initImg) == str: |
| 29 | B = scipy.misc.imread(initImg, mode='RGB').astype(np.float64) / 255 |
| 30 | else: |
| 31 | B = scipy.asarray(initImg).astype(np.float64) / 255 |
| 32 | # B = self. |
| 33 | # B = scipy.misc.imread(initImg, mode='RGB').astype(np.float64)/255 |
| 34 | h1,w1,k = B.shape |
| 35 | h = h1 - 4 |
| 36 | w = w1 - 4 |
| 37 | B = B[int((h1-h)/2):int((h1-h)/2+h),int((w1-w)/2):int((w1-w)/2+w),:] |
| 38 | content = scipy.misc.imresize(content,(h,w)) |
| 39 | B = self.__replication_padding(B,2) |
| 40 | content = self.__replication_padding(content,2) |
| 41 | content = content.astype(np.float64)/255 |
| 42 | B = np.reshape(B,(h1*w1,k)) |
| 43 | W = self.__compute_laplacian(content) |
| 44 | W = W.tocsc() |
| 45 | dd = W.sum(0) |
| 46 | dd = np.sqrt(np.power(dd,-1)) |
| 47 | dd = dd.A.squeeze() |
| 48 | D = scipy.sparse.csc_matrix((dd, (np.arange(0,w1*h1), np.arange(0,w1*h1)))) # 0.026 |
| 49 | S = D.dot(W).dot(D) |
| 50 | A = scipy.sparse.identity(w1*h1) - self.beta*S |
| 51 | A = A.tocsc() |
| 52 | solver = scipy.sparse.linalg.factorized(A) |
| 53 | V = np.zeros((h1*w1,k)) |
| 54 | V[:,0] = solver(B[:,0]) |
| 55 | V[:,1] = solver(B[:,1]) |
| 56 | V[:,2] = solver(B[:,2]) |
| 57 | V = V*(1-self.beta) |
| 58 | V = V.reshape(h1,w1,k) |
| 59 | V = V[2:2+h,2:2+w,:] |
| 60 | |
| 61 | img = Image.fromarray(np.uint8(np.clip(V * 255., 0, 255.))) |
| 62 | return img |
| 63 | |
| 64 | # Returns sparse matting laplacian |
| 65 | # The implementation of the function is heavily borrowed from |
| 66 | # https://github.com/MarcoForte/closed-form-matting/blob/master/closed_form_matting.py |
| 67 | # We thank Marco Forte for sharing his code. |
| 68 | def __compute_laplacian(self, img, eps=10**(-7), win_rad=1): |
| 69 | win_size = (win_rad*2+1)**2 |
| 70 | h, w, d = img.shape |
| 71 | c_h, c_w = h - 2*win_rad, w - 2*win_rad |
| 72 | win_diam = win_rad*2+1 |
no outgoing calls
no test coverage detected