Modifications to random walker segmentation algorithm:

* returning the probability to belong to a label instead of only the most
  likely label is now possible

* fixing some type issues

* handling non-consecutive label values
This commit is contained in:
Emmanuelle Gouillart
2012-06-24 19:35:27 +02:00
parent 8ece051c9c
commit 4ab7d0a4fa
2 changed files with 93 additions and 23 deletions
@@ -27,6 +27,8 @@ try:
except ImportError:
amg_loaded = False
from scipy.sparse.linalg import cg
from ..util import img_as_float
from ..filter import rank_order
#-----------Laplacian--------------------
@@ -96,7 +98,10 @@ def _make_laplacian_sparse(edges, weights):
return lap.tocsr()
def _clean_labels_ar(X, labels):
def _clean_labels_ar(X, labels, copy=False):
X = X.astype(labels.dtype)
if copy:
labels = np.copy(labels)
labels = np.ravel(labels)
labels[labels == 0] = X
return labels
@@ -157,7 +162,8 @@ def _build_laplacian(data, mask=None, beta=50):
#----------- Random walker algorithm --------------------------------
def random_walker(data, labels, beta=130, mode='bf', tol=1.e-3, copy=True):
def random_walker(data, labels, beta=130, mode='bf', tol=1.e-3, copy=True,
return_full_prob=False, reorder_labels=False):
"""
Random walker algorithm for segmentation from markers.
@@ -172,7 +178,9 @@ def random_walker(data, labels, beta=130, mode='bf', tol=1.e-3, copy=True):
Array of seed markers labeled with different positive integers
for different phases. Zero-labeled pixels are unlabeled pixels.
Negative labels correspond to inactive pixels that are not taken
into account (they are removed from the graph).
into account (they are removed from the graph). If labels are not
consecutive integers and `reorder_labels` is True, the labels array
will be transformed so that labels are consecutive.
beta : float
Penalization coefficient for the random walker motion
@@ -208,12 +216,24 @@ def random_walker(data, labels, beta=130, mode='bf', tol=1.e-3, copy=True):
the result of the segmentation. Use copy=False if you want to
save on memory.
return_full_prob : bool, default False
If True, the probability that a pixel belongs to each of the labels
will be returned, instead of only the most likely label.
reorder_labels : bool, default False
If True, labels is transformed so that its values are consecutive
integers.
Returns
-------
output : ndarray of ints
Array in which each pixel has been labeled according to the marker
that reached the pixel first by anisotropic diffusion.
output : ndarray
If `return_full_prob` is False, array of ints of same shape as `data`,
in which each pixel has been labeled according to the marker that
reached the pixel first by anisotropic diffusion.
If `return_full_prob` is True, array of floats of shape
`(nlabels, data.shape)`. `output[label_nb, i, j]` is the probability
that label `label_nb` reaches the pixel `(i, j)` first.
See also
--------
@@ -247,7 +267,8 @@ def random_walker(data, labels, beta=130, mode='bf', tol=1.e-3, copy=True):
The weight w_ij is a decreasing function of the norm of the local gradient.
This ensures that diffusion is easier between pixels of similar values.
When the Laplacian is decomposed into blocks of marked and unmarked pixels::
When the Laplacian is decomposed into blocks of marked and unmarked
pixels::
L = M B.T
B A
@@ -257,7 +278,7 @@ def random_walker(data, labels, beta=130, mode='bf', tol=1.e-3, copy=True):
A x = - B x_m
where x_m=1 on markers of the given phase, and 0 on other markers.
where x_m = 1 on markers of the given phase, and 0 on other markers.
This linear system is solved in the algorithm using a direct method for
small images, and an iterative method for larger images.
@@ -282,11 +303,15 @@ def random_walker(data, labels, beta=130, mode='bf', tol=1.e-3, copy=True):
[ 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]])
"""
# We work with 3-D arrays
# We work with 3-D arrays of floats
data = img_as_float(data)
data = np.atleast_3d(data)
if copy:
labels = np.copy(labels)
labels = labels.astype(np.intp)
if reorder_labels:
mask = labels >= 0
labels[mask] = rank_order(labels[mask])[0].astype(labels.dtype)
labels = labels.astype(np.int32)
# If the array has pruned zones, be sure that no isolated pixels
# exist between pruned zones (they could not be determined)
if np.any(labels < 0):
@@ -304,7 +329,8 @@ def random_walker(data, labels, beta=130, mode='bf', tol=1.e-3, copy=True):
# where X[i, j] is the probability that a marker of label i arrives
# first at pixel j by anisotropic diffusion.
if mode == 'cg':
X = _solve_cg(lap_sparse, B, tol=tol)
X = _solve_cg(lap_sparse, B, tol=tol,
return_full_prob=return_full_prob)
if mode == 'cg_mg':
if not amg_loaded:
warnings.warn(
@@ -313,15 +339,23 @@ def random_walker(data, labels, beta=130, mode='bf', tol=1.e-3, copy=True):
instead.""")
X = _solve_cg(lap_sparse, B, tol=tol)
else:
X = _solve_cg_mg(lap_sparse, B, tol=tol)
X = _solve_cg_mg(lap_sparse, B, tol=tol,
return_full_prob=return_full_prob)
if mode == 'bf':
X = _solve_bf(lap_sparse, B)
X = _clean_labels_ar(X + 1, labels)
X = _solve_bf(lap_sparse, B,
return_full_prob=return_full_prob)
# Clean up results
data = np.squeeze(data)
return X.reshape(data.shape)
if return_full_prob:
labels = labels.astype(np.float)
X = np.array([_clean_labels_ar(Xline, labels,
copy=True).reshape(data.shape) for Xline in X])
else:
X = _clean_labels_ar(X + 1, labels).reshape(data.shape)
return X
def _solve_bf(lap_sparse, B):
def _solve_bf(lap_sparse, B, return_full_prob=False):
"""
solves lap_sparse X_i = B_i for each phase i. An LU decomposition
of lap_sparse is computed first. For each pixel, the label i
@@ -331,11 +365,12 @@ def _solve_bf(lap_sparse, B):
solver = sparse.linalg.factorized(lap_sparse.astype(np.double))
X = np.array([solver(np.array((-B[i]).todense()).ravel())\
for i in range(len(B))])
X = np.argmax(X, axis=0)
if not return_full_prob:
X = np.argmax(X, axis=0)
return X
def _solve_cg(lap_sparse, B, tol):
def _solve_cg(lap_sparse, B, tol, return_full_prob=False):
"""
solves lap_sparse X_i = B_i for each phase i, using the conjugate
gradient method. For each pixel, the label i corresponding to the
@@ -346,12 +381,13 @@ def _solve_cg(lap_sparse, B, tol):
for i in range(len(B)):
x0 = cg(lap_sparse, -B[i].todense(), tol=tol)[0]
X.append(x0)
X = np.array(X)
X = np.argmax(X, axis=0)
if not return_full_prob:
X = np.array(X)
X = np.argmax(X, axis=0)
return X
def _solve_cg_mg(lap_sparse, B, tol):
def _solve_cg_mg(lap_sparse, B, tol, return_full_prob=False):
"""
solves lap_sparse X_i = B_i for each phase i, using the conjugate
gradient method with a multigrid preconditioner (ruge-stuben from
@@ -364,6 +400,7 @@ def _solve_cg_mg(lap_sparse, B, tol):
for i in range(len(B)):
x0 = cg(lap_sparse, -B[i].todense(), tol=tol, M=M, maxiter=30)[0]
X.append(x0)
X = np.array(X)
X = np.argmax(X, axis=0)
if not return_full_prob:
X = np.array(X)
X = np.argmax(X, axis=0)
return X
@@ -54,6 +54,10 @@ def test_2d_bf():
data, labels = make_2d_syntheticdata(lx, ly)
labels_bf = random_walker(data, labels, beta=90, mode='bf')
assert (labels_bf[25:45, 40:60] == 2).all()
full_prob_bf = random_walker(data, labels, beta=90, mode='bf',
return_full_prob=True)
assert (full_prob_bf[1, 25:45, 40:60] >=
full_prob_bf[0, 25:45, 40:60]).all()
return data, labels_bf
@@ -63,6 +67,10 @@ def test_2d_cg():
data, labels = make_2d_syntheticdata(lx, ly)
labels_cg = random_walker(data, labels, beta=90, mode='cg')
assert (labels_cg[25:45, 40:60] == 2).all()
full_prob = random_walker(data, labels, beta=90, mode='cg',
return_full_prob=True)
assert (full_prob[1, 25:45, 40:60] >=
full_prob[0, 25:45, 40:60]).all()
return data, labels_cg
@@ -72,8 +80,33 @@ def test_2d_cg_mg():
data, labels = make_2d_syntheticdata(lx, ly)
labels_cg_mg = random_walker(data, labels, beta=90, mode='cg_mg')
assert (labels_cg_mg[25:45, 40:60] == 2).all()
full_prob = random_walker(data, labels, beta=90, mode='cg_mg',
return_full_prob=True)
assert (full_prob[1, 25:45, 40:60] >=
full_prob[0, 25:45, 40:60]).all()
return data, labels_cg_mg
def test_types():
lx = 70
ly = 100
data, labels = make_2d_syntheticdata(lx, ly)
data = 255 * (data - data.min()) / (data.max() - data.min())
data = data.astype(np.uint8)
labels_cg_mg = random_walker(data, labels, beta=90, mode='cg_mg')
assert (labels_cg_mg[25:45, 40:60] == 2).all()
return data, labels_cg_mg
def test_reorder_labels():
lx = 70
ly = 100
data, labels = make_2d_syntheticdata(lx, ly)
labels[labels == 2] == 4
labels_bf = random_walker(data, labels, beta=90, mode='bf',
reorder_labels=True)
assert (labels_bf[25:45, 40:60] == 2).all()
return data, labels_bf
def test_2d_inactive():
lx = 70