-
Notifications
You must be signed in to change notification settings - Fork 107
/
solve_foreground_background.py
executable file
·155 lines (119 loc) · 5.82 KB
/
solve_foreground_background.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
#!/usr/bin/env python
"""Computes foreground and background images given source image and transparency map.
This module implements foreground and background reconstruction method described in Section 7 of:
Levin, Anat, Dani Lischinski, and Yair Weiss. "A closed-form solution to natural image
matting." IEEE Transactions on Pattern Analysis and Machine Intelligence 30.2 (2008): 228-242.
Please note, that the cost-function optimized by this code doesn't perfectly match Eq. 19 of the
paper, since our implementation mimics `solveFB.m` Matlab function provided by the authors of the
original paper (this implementation is
availale at http://people.csail.mit.edu/alevin/matting.tar.gz).
The code can be used in two ways:
1. By importing solve_foregound_background in your code:
```
from solve_foregound_background import solve_foregound_background
...
foreground, background = solve_foregound_background(image, alpha)
```
2. From command line:
```
./solve_foregound_background.py image.png alpha.png foreground.png background.png
```
Authors: Mikhail Erofeev, Yury Gitman.
"""
import numpy as np
import scipy.sparse
import scipy.sparse.linalg
CONST_ALPHA_MARGIN = 0.02
def __spdiagonal(diag):
"""Produces sparse matrix with given vector on its main diagonal."""
return scipy.sparse.spdiags(diag, (0,), len(diag), len(diag))
def get_grad_operator(mask):
"""Returns sparse matrix computing horizontal, vertical, and two diagonal gradients."""
horizontal_left = np.ravel_multi_index(np.nonzero(mask[:, :-1] | mask[:, 1:]), mask.shape)
horizontal_right = horizontal_left + 1
vertical_top = np.ravel_multi_index(np.nonzero(mask[:-1, :] | mask[1:, :]), mask.shape)
vertical_bottom = vertical_top + mask.shape[1]
diag_main_1 = np.ravel_multi_index(np.nonzero(mask[:-1, :-1] | mask[1:, 1:]), mask.shape)
diag_main_2 = diag_main_1 + mask.shape[1] + 1
diag_sub_1 = np.ravel_multi_index(np.nonzero(mask[:-1, 1:] | mask[1:, :-1]), mask.shape) + 1
diag_sub_2 = diag_sub_1 + mask.shape[1] - 1
indices = np.stack((
np.concatenate((horizontal_left, vertical_top, diag_main_1, diag_sub_1)),
np.concatenate((horizontal_right, vertical_bottom, diag_main_2, diag_sub_2))
), axis=-1)
return scipy.sparse.coo_matrix(
(np.tile([-1, 1], len(indices)), (np.arange(indices.size) // 2, indices.flatten())),
shape=(len(indices), mask.size))
def get_const_conditions(image, alpha):
"""Returns sparse diagonal matrix and vector encoding color prior conditions."""
falpha = alpha.flatten()
weights = (
(falpha < CONST_ALPHA_MARGIN) * 100.0 +
0.03 * (1.0 - falpha) * (falpha < 0.3) +
0.01 * (falpha > 1.0 - CONST_ALPHA_MARGIN)
)
conditions = __spdiagonal(weights)
mask = falpha < 1.0 - CONST_ALPHA_MARGIN
right_hand = (weights * mask)[:, np.newaxis] * image.reshape((alpha.size, -1))
return conditions, right_hand
def solve_foreground_background(image, alpha):
"""Compute foreground and background image given source image and transparency map."""
consts = (alpha < CONST_ALPHA_MARGIN) | (alpha > 1.0 - CONST_ALPHA_MARGIN)
grad = get_grad_operator(~consts)
grad_weights = np.power(np.abs(grad * alpha.flatten()), 0.5)
grad_only_positive = grad.maximum(0)
grad_weights_f = grad_weights + 0.003 * grad_only_positive * (1.0 - alpha.flatten())
grad_weights_b = grad_weights + 0.003 * grad_only_positive * alpha.flatten()
grad_pad = scipy.sparse.coo_matrix(grad.shape)
smoothness_conditions = scipy.sparse.vstack((
scipy.sparse.hstack((__spdiagonal(grad_weights_f) * grad, grad_pad)),
scipy.sparse.hstack((grad_pad, __spdiagonal(grad_weights_b) * grad))
))
composite_conditions = scipy.sparse.hstack((
__spdiagonal(alpha.flatten()),
__spdiagonal(1.0 - alpha.flatten())
))
const_conditions_f, b_const_f = get_const_conditions(image, 1.0 - alpha)
const_conditions_b, b_const_b = get_const_conditions(image, alpha)
non_zero_conditions = scipy.sparse.vstack((
composite_conditions,
scipy.sparse.hstack((
const_conditions_f,
scipy.sparse.coo_matrix(const_conditions_f.shape)
)),
scipy.sparse.hstack((
scipy.sparse.coo_matrix(const_conditions_b.shape),
const_conditions_b
))
))
b_composite = image.reshape(alpha.size, -1)
right_hand = non_zero_conditions.transpose() * np.concatenate((b_composite,
b_const_f,
b_const_b))
conditons = scipy.sparse.vstack((
non_zero_conditions,
smoothness_conditions
))
left_hand = conditons.transpose() * conditons
solution = scipy.sparse.linalg.spsolve(left_hand, right_hand).reshape(2, *image.shape)
foreground = solution[0, :, :, :].reshape(*image.shape)
background = solution[1, :, :, :].reshape(*image.shape)
return foreground, background
def main():
"""Parse command line arguments and apply solve_foregound_background."""
import argparse
import cv2
arg_parser = argparse.ArgumentParser(description=__doc__)
arg_parser.add_argument('image', type=str)
arg_parser.add_argument('alpha', type=str)
arg_parser.add_argument('foreground', type=str)
arg_parser.add_argument('background', type=str, default=None, nargs='?')
args = arg_parser.parse_args()
image = cv2.imread(args.image) / 255.0
alpha = cv2.imread(args.alpha, 0) / 255.0
foreground, background = solve_foreground_background(image, alpha)
cv2.imwrite(args.foreground, foreground * 255.0)
if args.background:
cv2.imwrite(args.background, background * 255.0)
if __name__ == "__main__":
main()