Add missing support for num_peaks parameter

This commit is contained in:
Johannes Schönberger
2012-10-10 09:28:47 +02:00
parent d9a7bac999
commit ede5df6c70
2 changed files with 20 additions and 1 deletions
+11 -1
View File
@@ -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
@@ -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()