From ede5df6c7030cc5acca7f09c3030efc08ef4db38 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Johannes=20Sch=C3=B6nberger?= Date: Wed, 10 Oct 2012 09:28:47 +0200 Subject: [PATCH] Add missing support for num_peaks parameter --- skimage/transform/hough_transform.py | 12 +++++++++++- skimage/transform/tests/test_hough_transform.py | 9 +++++++++ 2 files changed, 20 insertions(+), 1 deletion(-) diff --git a/skimage/transform/hough_transform.py b/skimage/transform/hough_transform.py index 0d47a376..9a37cf2e 100644 --- a/skimage/transform/hough_transform.py +++ b/skimage/transform/hough_transform.py @@ -248,4 +248,14 @@ def hough_peaks(hspace, angles, dists, min_distance=10, min_angle=10, dist_peaks.append(dists[dist_idx]) angle_peaks.append(angles[angle_idx]) - return np.array(hspace_peaks), np.array(dist_peaks), np.array(angle_peaks) + hspace_peaks = np.array(hspace_peaks) + dist_peaks = np.array(dist_peaks) + angle_peaks = np.array(angle_peaks) + + if num_peaks < len(hspace_peaks): + idx_maxsort = np.argsort(hspace_peaks)[::-1][:num_peaks] + hspace_peaks = hspace_peaks[idx_maxsort] + dist_peaks = dist_peaks[idx_maxsort] + angle_peaks = angle_peaks[idx_maxsort] + + return hspace_peaks, dist_peaks, angle_peaks diff --git a/skimage/transform/tests/test_hough_transform.py b/skimage/transform/tests/test_hough_transform.py index 3fb0a6f3..a38e6e76 100644 --- a/skimage/transform/tests/test_hough_transform.py +++ b/skimage/transform/tests/test_hough_transform.py @@ -101,5 +101,14 @@ def test_hough_peaks_angle(): assert len(tf.hough_peaks(hspace, angles, dists, min_angle=90)[0]) == 1 +def test_hough_peaks_num(): + img = np.zeros((100, 100), dtype=np.bool_) + img[:, 30] = True + img[:, 40] = True + hspace, angles, dists = tf.hough(img) + assert len(tf.hough_peaks(hspace, angles, dists, min_distance=0, + min_angle=0, num_peaks=1)[0]) == 1 + + if __name__ == "__main__": run_module_suite()