1use crate::definitions::Image;
6use crate::drawing::draw_line_segment_mut;
7use crate::suppress::suppress_non_maximum;
8use image::{GenericImage, GenericImageView, GrayImage, ImageBuffer, Luma, Pixel};
9use std::f32;
10
11#[derive(Copy, Clone, Debug, PartialEq)]
13pub struct PolarLine {
14 pub r: f32,
16 pub angle_in_degrees: u32,
19}
20
21#[derive(Copy, Clone, Debug, PartialEq, Eq)]
23pub struct LineDetectionOptions {
24 pub vote_threshold: u32,
26 pub suppression_radius: u32,
31}
32
33pub fn detect_lines(image: &GrayImage, options: LineDetectionOptions) -> Vec<PolarLine> {
40 let (width, height) = image.dimensions();
41
42 let rmax = ((width * width + height * height) as f64).sqrt() as i32;
44
45 let mut acc: ImageBuffer<Luma<u32>, Vec<u32>> = ImageBuffer::new(2 * rmax as u32 + 1, 180u32);
49
50 let lut: Vec<(f32, f32)> = (0..180u32)
52 .map(|deg| (deg as f32).to_radians())
53 .map(f32::sin_cos)
54 .collect();
55
56 for y in 0..height {
57 for x in 0..width {
58 let p = unsafe { image.unsafe_get_pixel(x, y)[0] };
59
60 if p > 0 {
61 for (m, (s, c)) in lut.iter().enumerate() {
62 let r = (x as f32) * c + (y as f32) * s;
63 let d = r as i32 + rmax;
64
65 if d <= 2 * rmax && d >= 0 {
66 unsafe {
67 let vote_incr = acc.unsafe_get_pixel(d as u32, m as u32)[0] + 1;
68 acc.unsafe_put_pixel(d as u32, m as u32, Luma([vote_incr]));
69 }
70 }
71 }
72 }
73 }
74 }
75
76 let acc_sup = suppress_non_maximum(&acc, options.suppression_radius);
77
78 let mut lines = Vec::new();
79
80 for m in 0..acc_sup.height() {
81 for r in 0..acc_sup.width() {
82 let votes = unsafe { acc_sup.unsafe_get_pixel(r, m)[0] };
83 if votes >= options.vote_threshold {
84 let line = PolarLine {
85 r: (r as i32 - rmax) as f32,
86 angle_in_degrees: m,
87 };
88 lines.push(line);
89 }
90 }
91 }
92
93 lines
94}
95
96pub fn draw_polar_lines<P>(image: &Image<P>, lines: &[PolarLine], color: P) -> Image<P>
100where
101 P: Pixel,
102{
103 let mut out = image.clone();
104 draw_polar_lines_mut(&mut out, lines, color);
105 out
106}
107
108pub fn draw_polar_lines_mut<P>(image: &mut Image<P>, lines: &[PolarLine], color: P)
112where
113 P: Pixel,
114{
115 for line in lines {
116 draw_polar_line(image, *line, color);
117 }
118}
119
120fn draw_polar_line<P>(image: &mut Image<P>, line: PolarLine, color: P)
121where
122 P: Pixel,
123{
124 if let Some((s, e)) = intersection_points(line, image.width(), image.height()) {
125 draw_line_segment_mut(image, s, e, color);
126 }
127}
128
129fn intersection_points(
134 line: PolarLine,
135 image_width: u32,
136 image_height: u32,
137) -> Option<((f32, f32), (f32, f32))> {
138 let r = line.r;
139 let m = line.angle_in_degrees;
140 let w = image_width as f32;
141 let h = image_height as f32;
142
143 if m == 0 {
145 return if r >= 0.0 && r <= w {
146 Some(((r, 0.0), (r, h)))
147 } else {
148 None
149 };
150 }
151
152 if m == 90 {
154 return if r >= 0.0 && r <= h {
155 Some(((0.0, r), (w, r)))
156 } else {
157 None
158 };
159 }
160
161 let theta = (m as f32).to_radians();
162 let (sin, cos) = theta.sin_cos();
163
164 let right_y = cos.mul_add(-w, r) / sin;
165 let left_y = r / sin;
166 let bottom_x = sin.mul_add(-h, r) / cos;
167 let top_x = r / cos;
168
169 let mut start = None;
170
171 if right_y >= 0.0 && right_y <= h {
172 let right_intersect = (w, right_y);
173 if let Some(s) = start {
174 return Some((s, right_intersect));
175 }
176 start = Some(right_intersect);
177 }
178
179 if left_y >= 0.0 && left_y <= h {
180 let left_intersect = (0.0, left_y);
181 if let Some(s) = start {
182 return Some((s, left_intersect));
183 }
184 start = Some(left_intersect);
185 }
186
187 if bottom_x >= 0.0 && bottom_x <= w {
188 let bottom_intersect = (bottom_x, h);
189 if let Some(s) = start {
190 return Some((s, bottom_intersect));
191 }
192 start = Some(bottom_intersect);
193 }
194
195 if top_x >= 0.0 && top_x <= w {
196 let top_intersect = (top_x, 0.0);
197 if let Some(s) = start {
198 return Some((s, top_intersect));
199 }
200 }
201
202 None
203}
204
205#[cfg(test)]
206mod tests {
207 use super::*;
208 use image::{GrayImage, ImageBuffer, Luma};
209 use test::{black_box, Bencher};
210
211 fn assert_points_eq(
212 actual: Option<((f32, f32), (f32, f32))>,
213 expected: Option<((f32, f32), (f32, f32))>,
214 ) {
215 match (actual, expected) {
216 (None, None) => {}
217 (Some(ps), Some(qs)) => {
218 let points_eq = |p: (f32, f32), q: (f32, f32)| {
219 (p.0 - q.0).abs() < 1.0e-6 && (p.1 - q.1).abs() < 1.0e-6
220 };
221
222 match (points_eq(ps.0, qs.0), points_eq(ps.1, qs.1)) {
223 (true, true) => {}
224 _ => {
225 panic!("Expected {:?}, got {:?}", expected, actual);
226 }
227 };
228 }
229 (Some(_), None) => {
230 panic!("Expected None, got {:?}", actual);
231 }
232 (None, Some(_)) => {
233 panic!("Expected {:?}, got None", expected);
234 }
235 }
236 }
237
238 #[test]
239 fn intersection_points_zero_signed_distance() {
240 assert_points_eq(
242 intersection_points(
243 PolarLine {
244 r: 0.0,
245 angle_in_degrees: 0,
246 },
247 10,
248 5,
249 ),
250 Some(((0.0, 0.0), (0.0, 5.0))),
251 );
252 assert_points_eq(
254 intersection_points(
255 PolarLine {
256 r: 0.0,
257 angle_in_degrees: 90,
258 },
259 10,
260 5,
261 ),
262 Some(((0.0, 0.0), (10.0, 0.0))),
263 );
264 assert_points_eq(
266 intersection_points(
267 PolarLine {
268 r: 0.0,
269 angle_in_degrees: 45,
270 },
271 10,
272 5,
273 ),
274 Some(((0.0, 0.0), (0.0, 0.0))),
275 );
276 assert_points_eq(
278 intersection_points(
279 PolarLine {
280 r: 0.0,
281 angle_in_degrees: 135,
282 },
283 10,
284 5,
285 ),
286 Some(((0.0, 0.0), (5.0, 5.0))),
287 );
288 assert_points_eq(
291 intersection_points(
292 PolarLine {
293 r: 0.0,
294 angle_in_degrees: 135,
295 },
296 10,
297 10,
298 ),
299 Some(((10.0, 10.0), (0.0, 0.0))),
300 );
301 }
302
303 #[test]
304 fn intersection_points_positive_signed_distance() {
305 assert_points_eq(
307 intersection_points(
308 PolarLine {
309 r: 9.0,
310 angle_in_degrees: 0,
311 },
312 10,
313 5,
314 ),
315 Some(((9.0, 0.0), (9.0, 5.0))),
316 );
317 assert_points_eq(
319 intersection_points(
320 PolarLine {
321 r: 8.0,
322 angle_in_degrees: 0,
323 },
324 5,
325 10,
326 ),
327 None,
328 );
329 assert_points_eq(
331 intersection_points(
332 PolarLine {
333 r: 9.0,
334 angle_in_degrees: 90,
335 },
336 5,
337 10,
338 ),
339 Some(((0.0, 9.0), (5.0, 9.0))),
340 );
341 assert_points_eq(
343 intersection_points(
344 PolarLine {
345 r: 8.0,
346 angle_in_degrees: 90,
347 },
348 10,
349 5,
350 ),
351 None,
352 );
353 assert_points_eq(
355 intersection_points(
356 PolarLine {
357 r: 5.0,
358 angle_in_degrees: 45,
359 },
360 10,
361 5,
362 ),
363 Some(((50f32.sqrt() - 5.0, 5.0), (50f32.sqrt(), 0.0))),
364 );
365 }
366
367 #[test]
368 fn intersection_points_negative_signed_distance() {
369 assert_points_eq(
371 intersection_points(
372 PolarLine {
373 r: -1.0,
374 angle_in_degrees: 0,
375 },
376 10,
377 5,
378 ),
379 None,
380 );
381 assert_points_eq(
383 intersection_points(
384 PolarLine {
385 r: -1.0,
386 angle_in_degrees: 90,
387 },
388 5,
389 10,
390 ),
391 None,
392 );
393 assert_points_eq(
395 intersection_points(
396 PolarLine {
397 r: -5.0,
398 angle_in_degrees: 135,
399 },
400 10,
401 5,
402 ),
403 Some(((10.0, 10.0 - 50f32.sqrt()), (50f32.sqrt(), 0.0))),
404 );
405 }
406
407 fn separated_horizontal_line_segment() -> GrayImage {
415 let white = Luma([255u8]);
416 let mut image = GrayImage::new(20, 5);
417 for i in 5..10 {
418 image.put_pixel(i, 2, white);
419 }
420 for i in 12..17 {
421 image.put_pixel(i, 2, white);
422 }
423 image
424 }
425
426 #[test]
427 fn detect_lines_horizontal_below_threshold() {
428 let image = separated_horizontal_line_segment();
429 let options = LineDetectionOptions {
430 vote_threshold: 11,
431 suppression_radius: 0,
432 };
433 let detected = detect_lines(&image, options);
434 assert_eq!(detected.len(), 0);
435 }
436
437 #[test]
438 fn detect_lines_horizontal_above_threshold() {
439 let image = separated_horizontal_line_segment();
440 let options = LineDetectionOptions {
441 vote_threshold: 10,
442 suppression_radius: 8,
443 };
444 let detected = detect_lines(&image, options);
445 assert_eq!(detected.len(), 1);
446 let line = detected[0];
447 assert_eq!(line.r, 1f32);
448 assert_eq!(line.angle_in_degrees, 90);
449 }
450
451 fn image_with_polar_line(
452 width: u32,
453 height: u32,
454 r: f32,
455 angle_in_degrees: u32,
456 color: Luma<u8>,
457 ) -> GrayImage {
458 let mut image = GrayImage::new(width, height);
459 draw_polar_line(
460 &mut image,
461 PolarLine {
462 r,
463 angle_in_degrees,
464 },
465 color,
466 );
467 image
468 }
469
470 #[test]
471 fn draw_polar_line_horizontal() {
472 let actual = image_with_polar_line(5, 5, 2.0, 90, Luma([1]));
473 let expected = gray_image!(
474 0, 0, 0, 0, 0;
475 0, 0, 0, 0, 0;
476 1, 1, 1, 1, 1;
477 0, 0, 0, 0, 0;
478 0, 0, 0, 0, 0);
479 assert_pixels_eq!(actual, expected);
480 }
481
482 #[test]
483 fn draw_polar_line_vertical() {
484 let actual = image_with_polar_line(5, 5, 2.0, 0, Luma([1]));
485 let expected = gray_image!(
486 0, 0, 1, 0, 0;
487 0, 0, 1, 0, 0;
488 0, 0, 1, 0, 0;
489 0, 0, 1, 0, 0;
490 0, 0, 1, 0, 0);
491 assert_pixels_eq!(actual, expected);
492 }
493
494 #[test]
495 fn draw_polar_line_bottom_left_to_top_right() {
496 let actual = image_with_polar_line(5, 5, 3.0, 45, Luma([1]));
497 let expected = gray_image!(
498 0, 0, 0, 0, 1;
499 0, 0, 0, 1, 0;
500 0, 0, 1, 0, 0;
501 0, 1, 0, 0, 0;
502 1, 0, 0, 0, 0);
503 assert_pixels_eq!(actual, expected);
504 }
505
506 #[test]
507 fn draw_polar_line_top_left_to_bottom_right() {
508 let actual = image_with_polar_line(5, 5, 0.0, 135, Luma([1]));
509 let expected = gray_image!(
510 1, 0, 0, 0, 0;
511 0, 1, 0, 0, 0;
512 0, 0, 1, 0, 0;
513 0, 0, 0, 1, 0;
514 0, 0, 0, 0, 1);
515 assert_pixels_eq!(actual, expected);
516 }
517
518 macro_rules! test_detect_line {
519 ($name:ident, $r:expr, $angle:expr) => {
520 #[test]
521 fn $name() {
522 let options = LineDetectionOptions {
523 vote_threshold: 10,
524 suppression_radius: 8,
525 };
526 let image = image_with_polar_line(100, 100, $r, $angle, Luma([255]));
527 let detected = detect_lines(&image, options);
528 assert_eq!(detected.len(), 1);
529
530 let line = detected[0];
531 assert_approx_eq!(line.r, $r, 1.1);
532 assert_approx_eq!(line.angle_in_degrees as f32, $angle as f32, 5.0);
533 }
534 };
535 }
536
537 test_detect_line!(detect_line_50_45, 50.0, 45);
538 test_detect_line!(detect_line_eps_135, 0.001, 135);
539 test_detect_line!(detect_line_neg10_120, -10.0, 120);
541
542 macro_rules! bench_detect_lines {
543 ($name:ident, $r:expr, $angle:expr) => {
544 #[bench]
545 fn $name(b: &mut Bencher) {
546 let options = LineDetectionOptions {
547 vote_threshold: 10,
548 suppression_radius: 8,
549 };
550 let mut image = GrayImage::new(100, 100);
551 draw_polar_line(
552 &mut image,
553 PolarLine {
554 r: $r,
555 angle_in_degrees: $angle,
556 },
557 Luma([255u8]),
558 );
559
560 b.iter(|| {
561 let lines = detect_lines(&image, options);
562 black_box(lines);
563 });
564 }
565 };
566 }
567
568 bench_detect_lines!(bench_detect_line_50_45, 50.0, 45);
569 bench_detect_lines!(bench_detect_line_eps_135, 0.001, 135);
570 bench_detect_lines!(bench_detect_line_neg10_120, -10.0, 120);
571
572 fn chessboard(width: u32, height: u32) -> GrayImage {
573 ImageBuffer::from_fn(width, height, |x, y| {
574 if (x + y) % 2 == 0 {
575 Luma([255u8])
576 } else {
577 Luma([0u8])
578 }
579 })
580 }
581
582 #[bench]
583 fn bench_detect_lines(b: &mut Bencher) {
584 let image = chessboard(100, 100);
585
586 let options = LineDetectionOptions {
587 vote_threshold: 10,
588 suppression_radius: 3,
589 };
590
591 b.iter(|| {
592 let lines = detect_lines(&image, options);
593 black_box(lines);
594 });
595 }
596}