ENH: Add num_peaks to limit the number of peaks.

skimage/feature/peak.py

@@ -4,7 +4,7 @@ from scipy import ndimage
 
 
 def peak_local_max(image, min_distance=10, threshold='deprecated',
-                   threshold_abs=0, threshold_rel=0.1):
+                   threshold_abs=0, threshold_rel=0.1, num_peaks=np.inf):
     """Return coordinates of peaks in an image.
 
     Peaks are the local maxima in a region of `2 * min_distance + 1`
@@ -30,6 +30,10 @@ def peak_local_max(image, min_distance=10, threshold='deprecated',
     threshold_rel: float
         Minimum intensity of peaks calculated as `max(image) * threshold_rel`.
 
+    num_peaks : int
+        Maximum number of peaks. When the number of peaks exceeds `num_peaks`,
+        return `num_peaks` coordinates based on peak intensity.
+
     Returns
     -------
     coordinates : (N, 2) array
@@ -61,5 +65,10 @@ def peak_local_max(image, min_distance=10, threshold='deprecated',
     # get coordinates of peaks
     coordinates = np.transpose(image_t.nonzero())
 
+    if len(coordinates) > num_peaks:
+        intensities = image[tuple(coordinates.T)]
+        idx_maxsort = np.argsort(intensities)[::-1]
+        coordinates = coordinates[idx_maxsort][:2]
+
     return coordinates
 

skimage/feature/tests/test_peak.py

@@ -50,6 +50,18 @@ def test_flat_peak():
     assert len(peaks) == 4
 
 
+def test_num_peaks():
+    image = np.zeros((3, 7), dtype=np.uint8)
+    image[1, 1] = 10
+    image[1, 3] = 11
+    image[1, 5] = 12
+    assert len(peak.peak_local_max(image, min_distance=1)) == 3
+    peaks_limited = peak.peak_local_max(image, min_distance=1, num_peaks=2)
+    assert len(peaks_limited) == 2
+    assert (1, 3) in peaks_limited
+    assert (1, 5) in peaks_limited
+
+
 if __name__ == '__main__':
     from numpy import testing
     testing.run_module_suite()