Ball detection fixes

HULKs · Jul 21, 2024 · 219bf8e · 219bf8e
1 parent 77badc4
commit 219bf8e
Showing 1 changed file with 13 additions and 144 deletions.
diff --git a/crates/vision/src/perspective_grid_candidates_provider.rs b/crates/vision/src/perspective_grid_candidates_provider.rs
@@ -1,4 +1,4 @@
-use std::collections::{BTreeMap, HashSet};
+use std::collections::{HashMap, HashSet};
 
 use color_eyre::Result;
 use serde::{Deserialize, Serialize};
@@ -63,6 +63,7 @@ impl PerspectiveGridCandidatesProvider {
         );
 
         let candidates = generate_candidates(
+            context.camera_matrix,
             vertical_scanlines,
             skip_segments,
             &perspective_grid_ball_sizes,
@@ -122,12 +123,13 @@ fn find_matching_row(rows: &[Row], segment: &Segment) -> Option<(usize, Row)> {
 }
 
 fn generate_candidates(
+    camera_matrix: &CameraMatrix,
     vertical_scanlines: &[ScanLine],
     skip_segments: &HashSet<Point2<Pixel, u16>>,
     rows: &[Row],
     minimum_number_of_segments_per_circle: usize,
 ) -> PerspectiveGridCandidates {
-    let mut segments_per_circles = BTreeMap::new();
+    let mut segments_per_circles = HashMap::new();
 
     for scan_line in vertical_scanlines {
         for segment in &scan_line.segments {
@@ -149,9 +151,13 @@ fn generate_candidates(
         }
     }
 
-    let mut candidates = segments_per_circles
-        .into_iter()
-        .filter_map(|((row_index, index_in_row), segments_per_circle)| {
+    let mut candidates = Vec::new();
+    let center = point![320.0, 240.0];
+    if let Ok(radius) = camera_matrix.get_pixel_radius(0.05, center) {
+        candidates.push(Circle { center, radius });
+    };
+    candidates.extend(segments_per_circles.into_iter().filter_map(
+        |((row_index, index_in_row), segments_per_circle)| {
             if segments_per_circle >= minimum_number_of_segments_per_circle {
                 let row = rows[row_index];
                 Some(Circle {
@@ -164,8 +170,8 @@ fn generate_candidates(
             } else {
                 None
             }
-        })
-        .collect::<Vec<_>>();
+        },
+    ));
 
     candidates.sort_by(|a, b| b.center.y().total_cmp(&a.center.y()));
 
@@ -177,10 +183,6 @@ mod tests {
     use approx::assert_relative_eq;
     use linear_algebra::{vector, IntoTransform, Isometry3};
     use nalgebra::{Translation, UnitQuaternion};
-    use types::{
-        color::{Intensity, YCbCr444},
-        image_segments::EdgeType,
-    };
 
     use super::*;
 
@@ -239,137 +241,4 @@ mod tests {
             current
         });
     }
-
-    #[test]
-    fn candidates_correct_single_segment() {
-        let rows = vec![
-            Row {
-                circle_radius: 10.0,
-                center_y: 10.0,
-            },
-            Row {
-                circle_radius: 10.0,
-                center_y: 30.0,
-            },
-            Row {
-                circle_radius: 10.0,
-                center_y: 50.0,
-            },
-        ];
-        let vertical_scan_lines = vec![ScanLine {
-            position: 42,
-            segments: vec![Segment {
-                start: 20,
-                end: 50,
-                start_edge_type: EdgeType::ImageBorder,
-                end_edge_type: EdgeType::ImageBorder,
-                color: YCbCr444 { y: 0, cb: 0, cr: 0 },
-                field_color: Intensity::Low,
-            }],
-        }];
-        let skip_segments = HashSet::new();
-        let candidates = generate_candidates(&vertical_scan_lines, &skip_segments, &rows, 0);
-        assert_relative_eq!(
-            candidates,
-            PerspectiveGridCandidates {
-                candidates: vec![Circle {
-                    center: point![50.0, 30.0],
-                    radius: 10.0
-                }]
-            }
-        );
-    }
-
-    // TODO: Fix and reenable
-    // #[test]
-    // fn candidates_correct_multi_segment() {
-    //     let rows = vec![
-    //         Row {
-    //             circle_radius: 10.0,
-    //             center_y: 10.0,
-    //         },
-    //         Row {
-    //             circle_radius: 10.0,
-    //             center_y: 30.0,
-    //         },
-    //         Row {
-    //             circle_radius: 10.0,
-    //             center_y: 50.0,
-    //         },
-    //     ];
-    //     let segments = vec![
-    //         Segment {
-    //             start: 5,
-    //             end: 12,
-    //             start_edge_type: EdgeType::ImageBorder,
-    //             end_edge_type: EdgeType::ImageBorder,
-    //             color: YCbCr444 { y: 0, cb: 0, cr: 0 },
-    //             field_color: Intensity::Low,
-    //         },
-    //         Segment {
-    //             start: 18,
-    //             end: 28,
-    //             start_edge_type: EdgeType::ImageBorder,
-    //             end_edge_type: EdgeType::ImageBorder,
-    //             color: YCbCr444 { y: 0, cb: 0, cr: 0 },
-    //             field_color: Intensity::Low,
-    //         },
-    //         Segment {
-    //             start: 45,
-    //             end: 50,
-    //             start_edge_type: EdgeType::ImageBorder,
-    //             end_edge_type: EdgeType::ImageBorder,
-    //             color: YCbCr444 { y: 0, cb: 0, cr: 0 },
-    //             field_color: Intensity::Low,
-    //         },
-    //     ];
-    //     let vertical_scan_lines = vec![
-    //         ScanLine {
-    //             position: 0,
-    //             segments: segments.clone(),
-    //         },
-    //         ScanLine {
-    //             position: 42,
-    //             segments: segments.clone(),
-    //         },
-    //         ScanLine {
-    //             position: 110,
-    //             segments,
-    //         },
-    //     ];
-    //     let skip_segments = HashSet::from_iter(
-    //         [
-    //             point![0, 18],
-    //             point![42, 5],
-    //             point![42, 45],
-    //             point![110, 5],
-    //             point![110, 18],
-    //         ]
-    //         .map(|point| point),
-    //     );
-    //     let candidates = generate_candidates(&vertical_scan_lines, &skip_segments, &rows, 0);
-    //     assert_relative_eq!(
-    //         candidates,
-    //         PerspectiveGridCandidates {
-    //             candidates: vec![
-    //                 Circle {
-    //                     center: point![10.0, 50.0],
-    //                     radius: 10.0
-    //                 },
-    //                 Circle {
-    //                     center: point![110.0, 50.0],
-    //                     radius: 10.0
-    //                 },
-    //                 Circle {
-    //                     center: point![50.0, 30.0],
-    //                     radius: 10.0
-    //                 },
-    //                 Circle {
-    //                     center: point![10.0, 10.0],
-    //                     radius: 10.0
-    //                 },
-    //             ]
-    //         }
-    //     );
-    // }
 }