From aa92d5f0bd3d68144acecdaa1475acd74488a95e Mon Sep 17 00:00:00 2001 From: Emmanuelle Gouillart Date: Tue, 24 Jul 2012 23:01:15 +0200 Subject: [PATCH] Better handling of labeled pixels in random walker segmentation when we return the whole probability. --- skimage/segmentation/random_walker_segmentation.py | 5 +++++ skimage/segmentation/tests/test_random_walker.py | 2 +- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/skimage/segmentation/random_walker_segmentation.py b/skimage/segmentation/random_walker_segmentation.py index a5bed48b..eaae60c9 100644 --- a/skimage/segmentation/random_walker_segmentation.py +++ b/skimage/segmentation/random_walker_segmentation.py @@ -350,6 +350,11 @@ def random_walker(data, labels, beta=130, mode='bf', tol=1.e-3, copy=True, labels = labels.astype(np.float) X = np.array([_clean_labels_ar(Xline, labels, copy=True).reshape(data.shape) for Xline in X]) + for i in range(1, int(labels.max()) + 1): + mask_i = np.squeeze(labels == i) + X[i - 1, mask_i] = 1 + X[np.setdiff1d(np.arange(0, labels.max(), dtype=np.int), + [i - 1]), mask_i] = 0 else: X = _clean_labels_ar(X + 1, labels).reshape(data.shape) return X diff --git a/skimage/segmentation/tests/test_random_walker.py b/skimage/segmentation/tests/test_random_walker.py index aec9edea..6de312e7 100644 --- a/skimage/segmentation/tests/test_random_walker.py +++ b/skimage/segmentation/tests/test_random_walker.py @@ -58,7 +58,7 @@ def test_2d_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 + return data, labels_bf, full_prob_bf def test_2d_cg():