diff --git a/skimage/segmentation/random_walker_segmentation.py b/skimage/segmentation/random_walker_segmentation.py index 6df714d2..bd143ed6 100644 --- a/skimage/segmentation/random_walker_segmentation.py +++ b/skimage/segmentation/random_walker_segmentation.py @@ -69,7 +69,7 @@ def _compute_weights_3d(data, beta=130, eps=1.e-6, depth=1., gradients = 0 for channel in range(0, data.shape[-1]): gradients += _compute_gradients_3d(data[..., channel], - depth=depth) ** 2 + depth=depth) ** 2 # All channels considered together in this standard deviation beta /= 10 * data.std() if multichannel: @@ -334,17 +334,18 @@ def random_walker(data, labels, beta=130, mode='bf', tol=1.e-3, copy=True, # Parse input data if not multichannel: # We work with 4-D arrays of floats + assert data.ndim > 1 and data.ndim < 4, 'For non-multichannel input, \ + data must be of dimension 2 \ + or 3.' dims = data.shape - data = np.atleast_3d(img_as_float(data)) - data.shape += (1,) + data = np.atleast_3d(img_as_float(data))[..., np.newaxis] else: dims = data[..., 0].shape assert multichannel and data.ndim > 2, 'For multichannel input, data \ must have >= 3 dimensions.' data = img_as_float(data) if data.ndim == 3: - data.shape += (1,) - data = data.transpose((0, 1, 3, 2)) + data = data[..., np.newaxis].transpose((0, 1, 3, 2)) if copy: labels = np.copy(labels) @@ -378,9 +379,9 @@ def random_walker(data, labels, beta=130, mode='bf', tol=1.e-3, copy=True, if mode == 'cg_mg': if not amg_loaded: warnings.warn( - """pyamg (http://code.google.com/p/pyamg/)) is needed to use - this mode, but is not installed. The 'cg' mode will be used - instead.""") + """pyamg (http://code.google.com/p/pyamg/)) is needed to use + this mode, but is not installed. The 'cg' mode will be used + instead.""") X = _solve_cg(lap_sparse, B, tol=tol, return_full_prob=return_full_prob) else: @@ -411,7 +412,7 @@ def _solve_bf(lap_sparse, B, return_full_prob=False): """ lap_sparse = lap_sparse.tocsc() solver = sparse.linalg.factorized(lap_sparse.astype(np.double)) - X = np.array([solver(np.array((-B[i]).todense()).ravel())\ + X = np.array([solver(np.array((-B[i]).todense()).ravel()) for i in range(len(B))]) if not return_full_prob: X = np.argmax(X, axis=0) diff --git a/skimage/segmentation/tests/test_random_walker.py b/skimage/segmentation/tests/test_random_walker.py index 7df0241c..1cc0a1ee 100644 --- a/skimage/segmentation/tests/test_random_walker.py +++ b/skimage/segmentation/tests/test_random_walker.py @@ -1,15 +1,11 @@ import numpy as np from skimage.segmentation import random_walker -try: - import pyamg - amg_loaded = True -except ImportError: - amg_loaded = False def make_2d_syntheticdata(lx, ly=None): if ly is None: ly = lx + np.random.seed(1234) data = np.zeros((lx, ly)) + 0.1 * np.random.randn(lx, ly) small_l = int(lx / 5) data[lx / 2 - small_l:lx / 2 + small_l, @@ -29,6 +25,7 @@ def make_3d_syntheticdata(lx, ly=None, lz=None): ly = lx if lz is None: lz = lx + np.random.seed(1234) data = np.zeros((lx, ly, lz)) + 0.1 * np.random.randn(lx, ly, lz) small_l = int(lx / 5) data[lx / 2 - small_l:lx / 2 + small_l, @@ -40,8 +37,8 @@ def make_3d_syntheticdata(lx, ly=None, lz=None): # make a hole hole_size = np.max([1, small_l / 8]) data[lx / 2 - small_l, - ly / 2 - hole_size:ly / 2 + hole_size,\ - lz / 2 - hole_size:lz / 2 + hole_size] = 0 + ly / 2 - hole_size:ly / 2 + hole_size, + lz / 2 - hole_size:lz / 2 + hole_size] = 0 seeds = np.zeros_like(data) seeds[lx / 5, ly / 5, lz / 5] = 1 seeds[lx / 2 + small_l / 4, ly / 2 - small_l / 4, lz / 2 - small_l / 4] = 2 @@ -54,17 +51,21 @@ 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() + assert data.shape == labels.shape full_prob_bf = random_walker(data, labels, beta=90, mode='bf', - return_full_prob=True) + return_full_prob=True) assert (full_prob_bf[1, 25:45, 40:60] >= - full_prob_bf[0, 25:45, 40:60]).all() + full_prob_bf[0, 25:45, 40:60]).all() + assert data.shape == labels.shape # Now test with more than two labels labels[55, 80] = 3 full_prob_bf = random_walker(data, labels, beta=90, mode='bf', - return_full_prob=True) + return_full_prob=True) assert (full_prob_bf[1, 25:45, 40:60] >= - full_prob_bf[0, 25:45, 40:60]).all() + full_prob_bf[0, 25:45, 40:60]).all() assert len(full_prob_bf) == 3 + assert data.shape == labels.shape + def test_2d_cg(): lx = 70 @@ -72,10 +73,12 @@ 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() + assert data.shape == labels.shape full_prob = random_walker(data, labels, beta=90, mode='cg', - return_full_prob=True) + return_full_prob=True) assert (full_prob[1, 25:45, 40:60] >= - full_prob[0, 25:45, 40:60]).all() + full_prob[0, 25:45, 40:60]).all() + assert data.shape == labels.shape return data, labels_cg @@ -85,10 +88,12 @@ 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() + assert data.shape == labels.shape full_prob = random_walker(data, labels, beta=90, mode='cg_mg', - return_full_prob=True) + return_full_prob=True) assert (full_prob[1, 25:45, 40:60] >= - full_prob[0, 25:45, 40:60]).all() + full_prob[0, 25:45, 40:60]).all() + assert data.shape == labels.shape return data, labels_cg_mg @@ -100,6 +105,7 @@ def test_types(): 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() + assert data.shape == labels.shape return data, labels_cg_mg @@ -110,6 +116,7 @@ def test_reorder_labels(): labels[labels == 2] = 4 labels_bf = random_walker(data, labels, beta=90, mode='bf') assert (labels_bf[25:45, 40:60] == 2).all() + assert data.shape == labels.shape return data, labels_bf @@ -121,6 +128,7 @@ def test_2d_inactive(): labels[46:50, 33:38] = -2 labels = random_walker(data, labels, beta=90) assert (labels.reshape((lx, ly))[25:45, 40:60] == 2).all() + assert data.shape == labels.shape return data, labels @@ -130,6 +138,7 @@ def test_3d(): data, labels = make_3d_syntheticdata(lx, ly, lz) labels = random_walker(data, labels, mode='cg') assert (labels.reshape(data.shape)[13:17, 13:17, 13:17] == 2).all() + assert data.shape == labels.shape return data, labels @@ -142,18 +151,19 @@ def test_3d_inactive(): after_labels = np.copy(labels) labels = random_walker(data, labels, mode='cg') assert (labels.reshape(data.shape)[13:17, 13:17, 13:17] == 2).all() + assert data.shape == labels.shape return data, labels, old_labels, after_labels def test_multispectral_2d(): lx, ly = 70, 100 data, labels = make_2d_syntheticdata(lx, ly) - data2 = data.copy() - data.shape += (1,) - data = data.repeat(2, axis=2) # Result should be identical + data = data[..., np.newaxis].repeat(2, axis=-1) # Expect identical output multi_labels = random_walker(data, labels, mode='cg', multichannel=True) - single_labels = random_walker(data2, labels, mode='cg') + assert data[..., 0].shape == labels.shape + single_labels = random_walker(data[..., 0], labels, mode='cg') assert (multi_labels.reshape(labels.shape)[25:45, 40:60] == 2).all() + assert data[..., 0].shape == labels.shape return data, multi_labels, single_labels, labels @@ -161,14 +171,16 @@ def test_multispectral_3d(): n = 30 lx, ly, lz = n, n, n data, labels = make_3d_syntheticdata(lx, ly, lz) - data.shape += (1,) - data = data.repeat(2, axis=3) # Result should be identical + data = data[..., np.newaxis].repeat(2, axis=-1) # Expect identical output multi_labels = random_walker(data, labels, mode='cg', multichannel=True) + assert data[..., 0].shape == labels.shape single_labels = random_walker(data[..., 0], labels, mode='cg') assert (multi_labels.reshape(labels.shape)[13:17, 13:17, 13:17] == 2).all() assert (single_labels.reshape(labels.shape)[13:17, 13:17, 13:17] == 2).all() + assert data[..., 0].shape == labels.shape return data, multi_labels, single_labels, labels + if __name__ == '__main__': from numpy import testing testing.run_module_suite()