projective_grid/shared/
grow_extend.rs1use crate::shared::grow::{
13 any_cardinal_edge_ok, choose_unambiguous, collect_candidates, collect_labelled_neighbours,
14 enqueue_cardinal_neighbours, is_extrapolating, predict_from_neighbours, CandidateChoice,
15 GrowParams, GrowResult, SquareAttachPolicy,
16};
17use kiddo::KdTree;
18use nalgebra::Point2;
19use std::collections::{HashSet, VecDeque};
20
21#[non_exhaustive]
28#[derive(Clone, Debug, Default)]
29pub struct BfsExtensionStats {
30 pub attached: usize,
32 pub rejected_no_candidate: usize,
34 pub rejected_ambiguous: usize,
36 pub rejected_edge: usize,
38 pub attached_indices: Vec<usize>,
41 pub attached_cells: Vec<(i32, i32)>,
44}
45
46#[cfg_attr(
64 feature = "tracing",
65 tracing::instrument(
66 level = "info",
67 skip_all,
68 fields(num_corners = positions.len(), num_labelled = grow.labelled.len(), cell_size = cell_size),
69 )
70)]
71pub fn extend_from_labelled<V: SquareAttachPolicy>(
72 positions: &[Point2<f32>],
73 grow: &mut GrowResult,
74 cell_size: f32,
75 params: &GrowParams,
76 policy: &V,
77) -> BfsExtensionStats {
78 let mut tree: KdTree<f32, 2> = KdTree::new();
80 let mut tree_slot_to_corner: Vec<usize> = Vec::new();
81 for (idx, pos) in positions.iter().enumerate() {
82 if policy.is_eligible(idx) && !grow.by_corner.contains_key(&idx) {
83 tree.add(&[pos.x, pos.y], tree_slot_to_corner.len() as u64);
84 tree_slot_to_corner.push(idx);
85 }
86 }
87
88 let mut boundary: VecDeque<(i32, i32)> = VecDeque::new();
90 let mut seen_boundary: HashSet<(i32, i32)> = HashSet::new();
91 for &pos in grow.labelled.keys() {
92 enqueue_cardinal_neighbours(pos, &grow.labelled, &mut boundary, &mut seen_boundary);
93 }
94
95 let mut stats = BfsExtensionStats::default();
96
97 while let Some(pos) = boundary.pop_front() {
98 if grow.labelled.contains_key(&pos) {
99 continue;
100 }
101
102 let neighbours = collect_labelled_neighbours(pos, 1, &grow.labelled, positions);
103 if neighbours.is_empty() {
104 stats.rejected_no_candidate += 1;
105 grow.holes.insert(pos);
106 continue;
107 }
108
109 let prediction = predict_from_neighbours(
110 pos,
111 &neighbours,
112 grow.axis_i,
113 grow.axis_j,
114 cell_size,
115 &grow.labelled,
116 positions,
117 );
118
119 let search_r = params.attach_search_rel * cell_size;
120 let extrapolating = is_extrapolating(pos, &neighbours);
121 let local_search_r = if extrapolating {
122 search_r * params.boundary_search_factor
123 } else {
124 search_r
125 };
126
127 let required_label = policy.required_label_at(pos.0, pos.1);
128 let candidates = collect_candidates(
129 &tree,
130 &tree_slot_to_corner,
131 prediction,
132 local_search_r,
133 policy,
134 required_label,
135 &grow.by_corner,
136 );
137
138 let choice = choose_unambiguous(
139 &candidates,
140 params.attach_ambiguity_factor,
141 prediction,
142 positions,
143 policy,
144 pos,
145 &neighbours,
146 );
147
148 match choice {
149 CandidateChoice::None => {
150 stats.rejected_no_candidate += 1;
151 grow.holes.insert(pos);
152 }
153 CandidateChoice::Ambiguous => {
154 stats.rejected_ambiguous += 1;
155 grow.ambiguous.insert(pos);
156 }
157 CandidateChoice::Unique(c_idx) => {
158 if !any_cardinal_edge_ok(c_idx, pos, &grow.labelled, policy) {
159 stats.rejected_edge += 1;
160 grow.holes.insert(pos);
161 } else {
162 grow.labelled.insert(pos, c_idx);
163 grow.by_corner.insert(c_idx, pos);
164 enqueue_cardinal_neighbours(
165 pos,
166 &grow.labelled,
167 &mut boundary,
168 &mut seen_boundary,
169 );
170 stats.attached += 1;
171 stats.attached_indices.push(c_idx);
172 stats.attached_cells.push(pos);
173 }
174 }
175 }
176 }
177
178 stats
179}