Merge pull request #486 from yangzetian/peak-local-max-indices-nd

Fix peak_local_max's output for ndarray when indices is set to False

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

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()