diff --git a/skimage/feature/peak.py b/skimage/feature/peak.py index 9c7a934f..3268831a 100644 --- a/skimage/feature/peak.py +++ b/skimage/feature/peak.py @@ -33,8 +33,8 @@ def peak_local_max(image, min_distance=10, threshold_abs=0, threshold_rel=0.1, If True, `min_distance` excludes peaks from the border of the image as well as from each other. indices : bool - If True, the output will be a matrix representing peak coordinates. - If False, the output will be a boolean matrix shaped as `image.shape` + If True, the output will be an array representing peak coordinates. + If False, the output will be a boolean array shaped as `image.shape` with peaks present at True elements. num_peaks : int Maximum number of peaks. When the number of peaks exceeds `num_peaks`, @@ -130,11 +130,12 @@ def peak_local_max(image, min_distance=10, threshold_abs=0, threshold_rel=0.1, image *= mask if exclude_border: - # Remove the image borders - image[:min_distance] = 0 - image[-min_distance:] = 0 - image[:, :min_distance] = 0 - image[:, -min_distance:] = 0 + # zero out the image borders + for i in range(image.ndim): + image = image.swapaxes(0, i) + image[:min_distance] = 0 + image[-min_distance:] = 0 + image = image.swapaxes(0, i) # find top peak candidates above a threshold peak_threshold = max(np.max(image.ravel()) * threshold_rel, threshold_abs) @@ -150,5 +151,6 @@ def peak_local_max(image, min_distance=10, threshold_abs=0, threshold_rel=0.1, if indices is True: return coordinates else: - out[coordinates[:, 0], coordinates[:, 1]] = True + nd_indices = tuple(coordinates.T) + out[nd_indices] = True return out diff --git a/skimage/feature/tests/test_peak.py b/skimage/feature/tests/test_peak.py index 3ef1f12d..39dc8d8f 100644 --- a/skimage/feature/tests/test_peak.py +++ b/skimage/feature/tests/test_peak.py @@ -117,6 +117,24 @@ def test_indices_with_labels(): assert (result == np.transpose(expected.nonzero())).all() +def test_ndarray_indices_false(): + nd_image = np.zeros((5,5,5)) + nd_image[2,2,2] = 1 + peaks = peak.peak_local_max(nd_image, min_distance=1, indices=False) + assert (peaks == nd_image.astype(np.bool)).all() + + +def test_ndarray_exclude_border(): + nd_image = np.zeros((5,5,5)) + nd_image[[1,0,0],[0,1,0],[0,0,1]] = 1 + nd_image[3,0,0] = 1 + nd_image[2,2,2] = 1 + expected = np.zeros_like(nd_image, dtype=np.bool) + expected[2,2,2] = True + result = peak.peak_local_max(nd_image, min_distance=2, indices=False) + assert (result == expected).all() + + if __name__ == '__main__': from numpy import testing testing.run_module_suite()