def test_preconditioners():
atol = 1e-6
index = 1
scale = 0.2
name = f"GT{index:02d}"
# print(name)
image_dir = "data"
image = load_image(f"{image_dir}/input_training_lowres/{name}.png", "rgb",
scale, "bilinear")
trimap = load_image(
f"{image_dir}/trimap_training_lowres/Trimap1/{name}.png",
"gray",
scale,
"nearest",
)
A, b = make_linear_system(cf_laplacian(image), trimap)
preconditioners = [
("no", lambda A: None),
("jacobi", lambda A: jacobi(A)),
("icholt", lambda A: ichol(A, max_nnz=500000)),
("vcycle", lambda A: vcycle(A, trimap.shape)),
]
expected_iterations = {
"no": 532,
"jacobi": 250,
"icholt": 3,
"vcycle": 88,
}
for preconditioner_name, preconditioner in preconditioners:
callback = CounterCallback()
M = preconditioner(A)
x = cg(A, b, M=M, atol=atol, rtol=0, maxiter=10000, callback=callback)
r = b - A.dot(x)
norm_r = np.linalg.norm(r)
assert norm_r <= atol
n_expected = expected_iterations[preconditioner_name]
if callback.n > n_expected:
print(
"WARNING: Unexpected number of iterations. Expected %d, but got %d"
% (n_expected, callback.n))
assert callback.n <= n_expected
def main():
print("loading images")
size = (34, 22)
size = (680, 440)
image = np.array(
Image.open("images/lemur.png").convert("RGB").resize(
size, Image.BOX)) / 255.0
trimap = np.array(
Image.open("images/lemur_trimap.png").convert("L").resize(
size, Image.NEAREST)) / 255.0
is_fg = trimap == 1.0
is_bg = trimap == 0.0
is_known = is_fg | is_bg
is_unknown = ~is_known
b = 100.0 * is_fg.flatten()
c = 100.0 * is_known.flatten()
shape = trimap.shape
h, w = shape
L = cf_laplacian(image)
cache = {}
C = scipy.sparse.diags(c)
build_cache_L(L, shape, cache)
# If the trimap changes, simply call this function again with the new vector c
build_cache_c(c, shape, cache)
M = lambda x: vcycle_step(x, shape, cache)
x = cg(lambda x: L.dot(x) + c * x, b, M=M, callback=ProgressCallback())
print("\nbaseline:")
print("iteration 69 - 2.690681e-03 (0.00269068076571873623)")
print()
print("slight error due to discarding of C-offdiagonals")
alpha = np.clip(x, 0, 1).reshape(h, w)
import matplotlib.pyplot as plt
for i, img in enumerate([image, trimap, alpha]):
plt.subplot(1, 3, 1 + i)
plt.imshow(img, cmap='gray')
plt.axis('off')
plt.show()
def test_lkm():
index = 1
scale = 0.2
epsilon = 1e-7
radius = 2
image_dir = "data"
name = f"GT{index:02d}"
image = load_image(f"{image_dir}/input_training_lowres/{name}.png", "rgb",
scale, "bilinear")
trimap = load_image(
f"{image_dir}/trimap_training_lowres/Trimap1/{name}.png",
"gray",
scale,
"nearest",
)
L_lkm, diag_L_lkm = lkm_laplacian(image, epsilon=epsilon, radius=radius)
L_cf = cf_laplacian(image, epsilon=epsilon, radius=radius)
A_cf, b, c = make_linear_system(L_cf, trimap, return_c=True)
def A_lkm(x):
return L_lkm(x) + c * x
inv_diag_A_lkm = 1.0 / (diag_L_lkm + c)
def jacobi_lkm(r):
return inv_diag_A_lkm * r
jacobi_cf = jacobi(A_cf)
lkm_callback = ProgressCallback()
cf_callback = ProgressCallback()
x_lkm = cg(A_lkm, b, M=jacobi_lkm)
x_cf = cg(A_cf, b, M=jacobi_cf)
difference = np.linalg.norm(x_lkm - x_cf)
assert difference < 2e-4
assert abs(lkm_callback.n - cf_callback.n) <= 2
def main():
print("loading images")
size = (34, 22)
size = (680, 440)
image = np.array(
Image.open("images/lemur.png").convert("RGB").resize(
size, Image.BOX)) / 255.0
trimap = np.array(
Image.open("images/lemur_trimap.png").convert("L").resize(
size, Image.NEAREST)) / 255.0
is_fg = trimap == 1.0
is_bg = trimap == 0.0
is_known = is_fg | is_bg
is_unknown = ~is_known
b = 100.0 * is_fg.flatten()
c = 100.0 * is_known.flatten()
shape = trimap.shape
h, w = shape
L = cf_laplacian(image)
C = scipy.sparse.diags(c)
A = L + C
M = vcycle(A, (h, w))
x = cg(A, b, M=M, callback=ProgressCallback())
print("\nbaseline:")
print("iteration 69 - 2.690681e-03 (0.00269068076571873623)")
alpha = np.clip(x, 0, 1).reshape(h, w)
import matplotlib.pyplot as plt
for i, img in enumerate([image, trimap, alpha]):
plt.subplot(1, 3, 1 + i)
plt.imshow(img, cmap='gray')
plt.axis('off')
plt.show()
def run_solver_single_image(solver_name, scale, index):
# Load images
name = f"GT{index:02d}.png"
image_path = os.path.join(IMAGE_DIR, "input_training_lowres", name)
trimap_path = os.path.join(IMAGE_DIR, "trimap_training_lowres/Trimap1",
name)
image = load_image(image_path, "rgb", scale, "bilinear")
trimap = load_image(trimap_path, "gray", scale, "nearest")
# Create linear system
L = cf_laplacian(image)
A, b = make_linear_system(L, trimap)
is_fg, is_bg, is_known, is_unknown = trimap_split(trimap)
atol = ATOL * np.sum(is_known)
rtol = atol / np.linalg.norm(b)
# Compute various matrix representations
Acsr = A.tocsr()
Acsc = A.tocsc()
Acoo = A.tocoo()
AL = scipy.sparse.tril(Acoo)
# Start memory usage measurement thread
memory_usage = [get_memory_usage()]
thread = threading.Thread(target=log_memory_usage, args=(memory_usage, ))
thread.is_running = True
# All threads should die if the solver thread crashes so we can at least
# carry on with the other solvers.
thread.daemon = True
thread.start()
# Measure solver build time
# Note that it is not easily possible to separate build time from solve time
# for every solver, which is why only the sum of build_time and solve_time
# should be compared for fairness.
start_time = time.perf_counter()
run_solver = build_solver(solver_name, A, Acsr, Acsc, Acoo, AL, b, atol,
rtol)
build_time = time.perf_counter() - start_time
# Measure actual solve time
start_time = time.perf_counter()
x = run_solver()
solve_time = time.perf_counter() - start_time
# Stop memory usage measuring thread
thread.is_running = False
thread.join()
# Compute relative error
r = b - A.dot(x)
norm_r = np.linalg.norm(r)
# Store results
h, w = trimap.shape
result = dict(
solver_name=str(solver_name),
image_name=str(name),
scale=float(scale),
index=int(index),
norm_r=float(norm_r),
build_time=float(build_time),
solve_time=float(solve_time),
atol=float(atol),
rtol=float(rtol),
width=int(w),
height=int(h),
n_fg=int(np.sum(is_fg)),
n_bg=int(np.sum(is_bg)),
n_known=int(np.sum(is_known)),
n_unknown=int(np.sum(is_unknown)),
memory_usage=memory_usage,
)
print(result)
# Ensure that everything worked as expected
assert norm_r <= atol
# Inspect alpha for debugging
if 0:
alpha = np.clip(x, 0, 1).reshape(h, w)
show_images([alpha])
return result
def main():
solvers = []
if 0:
from pymatting import cg, ichol
def solve_cg_icholt(A, b, atol):
M = ichol(A, discard_threshold=1e-3, shifts=[0.002])
return cg(A, b, M=M, atol=atol, rtol=0)
solvers.append(
("cg_icholt", lambda: solve_cg_icholt(Acsc, b, atol=atol)))
if 1:
import pyamg
from pymatting import cg
def solve_pyamg(A, b, atol):
M = pyamg.smoothed_aggregation_solver(A).aspreconditioner()
return cg(A, b, M=M, atol=atol, rtol=0)
solvers.append(("pyamg", lambda: solve_pyamg(Acsr, b, atol=atol)))
if 0:
from solve_mumps import solve_mumps_coo, init_mpi, finalize_mpi
init_mpi()
def solve_mumps(A, b):
return solve_mumps_coo(A.data, A.row, A.col, b, is_symmetric=True)
solvers.append(("mumps", lambda: solve_mumps(AL, b)))
if 1:
from solve_petsc import solve_petsc_coo, init_petsc, finalize_petsc
init_petsc()
def solve_petsc(A, b, atol):
return solve_petsc_coo(A.data,
A.row,
A.col,
b,
atol=atol,
gamg_threshold=0.1)
solvers.append(("petsc", lambda: solve_petsc(Acoo, b, atol=atol)))
if 1:
from solve_amgcl import solve_amgcl_csr
def solve_amgcl(A, b, atol):
return solve_amgcl_csr(A.data,
A.indices,
A.indptr,
b,
atol=atol,
rtol=0)
solvers.append(("amgcl", lambda: solve_amgcl(Acsr, b, atol=atol)))
if 0:
solvers.append(
("umfpack",
lambda: scipy.sparse.linalg.spsolve(Acsc, b, use_umfpack=True)))
if 0:
def solve_superLU(A, b):
# Usually the same as:
# scipy.sparse.linalg.spsolve(A, b, use_umfpack=False)
return scipy.sparse.linalg.splu(A).solve(b)
solvers.append(("superlu", lambda: solve_superLU(Acsc, b)))
if 1:
from solve_eigen import solve_eigen_cholesky_coo
def solve_eigen_cholesky(A, b):
return solve_eigen_cholesky_coo(A.data, A.row, A.col, b)
solvers.append(
("eigen_cholesky", lambda: solve_eigen_cholesky(Acoo, b)))
if 0:
from solve_eigen import solve_eigen_icholt_coo
def solve_eigen_icholt(A, b, rtol):
# just large enough to not fail for given images (might fail with larger/different images)
initial_shift = 2e-4
return solve_eigen_icholt_coo(A.data,
A.row,
A.col,
b,
rtol=rtol,
initial_shift=initial_shift)
solvers.append(
("eigen_icholt", lambda: solve_eigen_icholt(Acoo, b, rtol=rtol)))
indices = np.arange(27)
scales = np.sqrt(np.linspace(0.1, 1.0, 11))
# indices = np.int32([1, 2])
# scales = [0.1]
atol0 = 1e-7
image_dir = "data"
for solver_name, solver in solvers:
results = []
for scale in scales:
for index in indices + 1:
name = f"GT{index:02d}"
image = load_image(
f"{image_dir}/input_training_lowres/{name}.png",
"rgb",
scale,
"bilinear",
)
trimap = load_image(
f"{image_dir}/trimap_training_lowres/Trimap1/{name}.png",
"gray",
scale,
"bilinear",
)
L = cf_laplacian(image)
A, b = make_linear_system(L, trimap)
is_fg, is_bg, is_known, is_unknown = trimap_split(trimap)
atol = atol0 * np.sum(is_known)
rtol = atol / np.linalg.norm(b)
Acsr = A.tocsr()
Acsc = A.tocsc()
Acoo = A.tocoo()
AL = scipy.sparse.tril(Acoo)
def log_memory_usage(memory_usage):
while running:
memory_usage.append(get_memory_usage())
time.sleep(0.01)
memory_usage = [get_memory_usage()]
running = True
thread = threading.Thread(target=log_memory_usage,
args=(memory_usage, ))
thread.start()
start_time = time.perf_counter()
x = solver()
elapsed_time = time.perf_counter() - start_time
running = False
thread.join()
r = b - A.dot(x)
norm_r = np.linalg.norm(r)
h, w = trimap.shape
result = dict(
solver_name=str(solver_name),
scale=float(scale),
index=int(index),
norm_r=float(norm_r),
t=float(elapsed_time),
atol=float(atol),
rtol=float(rtol),
width=int(w),
height=int(h),
n_fg=int(np.sum(is_fg)),
n_bg=int(np.sum(is_bg)),
n_known=int(np.sum(is_known)),
n_unknown=int(np.sum(is_unknown)),
memory_usage=memory_usage,
)
print(result)
assert norm_r <= atol
if 0:
alpha = np.clip(x, 0, 1).reshape(h, w)
show_images([alpha])
results.append(result)
path = f"benchmarks/results/solver/{solver_name}/results.json"
dir_name = os.path.split(path)[0]
if len(dir_name) > 0:
os.makedirs(dir_name, exist_ok=True)
with open(path, "w") as f:
json.dump(results, f, indent=4)
def main():
print("loading images")
mses = []
for idx in range(1, 28):
image = np.array(
Image.open(f"{DATASET_DIRECTORY}/input_training_lowres/GT%02d.png"
% idx).convert("RGB")) / 255.0
trimap = np.array(
Image.open(
f"{DATASET_DIRECTORY}/trimap_training_lowres/Trimap1/GT%02d.png"
% idx).convert("L")) / 255.0
alpha_gt = np.array(
Image.open(f"{DATASET_DIRECTORY}/gt_training_lowres/GT%02d.png" %
idx).convert("L")) / 255.0
is_fg = trimap == 1.0
is_bg = trimap == 0.0
is_known = is_fg | is_bg
is_unknown = ~is_known
b = 100.0 * is_fg.flatten()
c = 100.0 * is_known.flatten()
shape = trimap.shape
h, w = shape
L = cf_laplacian(image)
cache = {}
C = scipy.sparse.diags(c)
build_cache_L(L, shape, cache)
# If the trimap changes, simply call this function again with the new vector c
build_cache_c(c, shape, cache)
M = lambda x: vcycle_step(x, shape, cache)
x = cg(lambda x: L.dot(x) + c * x, b, M=M, callback=ProgressCallback())
alpha = np.clip(x, 0, 1).reshape(h, w)
difference = np.abs(alpha - alpha_gt)
mse = np.mean(np.square(difference[is_unknown]))
print("MSE:", mse)
mses.append(mse)
continue
import matplotlib.pyplot as plt
for i, img in enumerate([image, trimap, alpha]):
plt.subplot(1, 3, 1 + i)
plt.imshow(img, cmap='gray')
plt.axis('off')
plt.show()
print("mean MSE:", np.mean(mses))
print()
print("Expected:")
print("mean MSE: 0.021856688900784647")
def test_lkm():
index = 1
scale = 0.2
atol = 1e-5
epsilon = 1e-7
radius = 2
image_dir = "data"
errors = []
name = f"GT{index:02d}"
# print(name)
image = load_image(
f"{image_dir}/input_training_lowres/{name}.png", "rgb", scale, "bilinear"
)
trimap = load_image(
f"{image_dir}/trimap_training_lowres/Trimap1/{name}.png",
"gray",
scale,
"bilinear",
)
true_alpha = load_image(
f"{image_dir}/gt_training_lowres/{name}.png", "gray", scale, "nearest"
)
L_lkm, diag_L_lkm = lkm_laplacian(image, epsilon=epsilon, radius=radius)
L_cf = cf_laplacian(image, epsilon=epsilon, radius=radius)
A_cf, b, c = make_linear_system(L_cf, trimap, return_c=True)
def A_lkm(x):
return L_lkm(x) + c * x
inv_diag_A_lkm = 1.0 / (diag_L_lkm + c)
def jacobi_lkm(r):
return inv_diag_A_lkm * r
jacobi_cf = jacobi(A_cf)
lkm_callback = ProgressCallback()
cf_callback = ProgressCallback()
x_lkm = cg(A_lkm, b, M=jacobi_lkm)
x_cf = cg(A_cf, b, M=jacobi_cf)
if 0:
# show result
show_images(
[x_lkm.reshape(trimap.shape), x_cf.reshape(trimap.shape),]
)
difference = np.linalg.norm(x_lkm - x_cf)
# print("norm(x_lkm - x_cf):")
# print(difference)
# print("iterations:")
# print("lkm:", lkm_callback.n)
# print("cf:", cf_callback.n)
assert difference < 1e-5
assert abs(lkm_callback.n - cf_callback.n) <= 2