1use crate::types::PlateConstants;
25
26const MIN_PAIR_PX: f64 = 8.0;
29
30pub const CENSUS_SIGNIFICANCE: f64 = 2.0;
36
37struct PosGrid {
40 min_x: f64,
41 min_y: f64,
42 inv_cell: f64,
43 nx: usize,
44 ny: usize,
45 start: Vec<u32>,
47 items: Vec<u32>,
48}
49
50impl PosGrid {
51 fn new(pos: &[(f64, f64)], cell: f64) -> Self {
52 let (mut min_x, mut min_y) = (f64::INFINITY, f64::INFINITY);
53 let (mut max_x, mut max_y) = (f64::NEG_INFINITY, f64::NEG_INFINITY);
54 for &(x, y) in pos {
55 min_x = min_x.min(x);
56 max_x = max_x.max(x);
57 min_y = min_y.min(y);
58 max_y = max_y.max(y);
59 }
60 if pos.is_empty() {
61 (min_x, min_y, max_x, max_y) = (0.0, 0.0, 0.0, 0.0);
62 }
63 let cell = cell.max(0.5);
64 let inv_cell = 1.0 / cell;
65 let nx = ((max_x - min_x) * inv_cell) as usize + 1;
66 let ny = ((max_y - min_y) * inv_cell) as usize + 1;
67 let cell_of = |x: f64, y: f64| {
68 let gx = (((x - min_x) * inv_cell) as usize).min(nx - 1);
69 let gy = (((y - min_y) * inv_cell) as usize).min(ny - 1);
70 gy * nx + gx
71 };
72 let mut start = vec![0u32; nx * ny + 1];
73 for &(x, y) in pos {
74 start[cell_of(x, y) + 1] += 1;
75 }
76 for c in 0..nx * ny {
77 start[c + 1] += start[c];
78 }
79 let mut fill = start.clone();
80 let mut items = vec![0u32; pos.len()];
81 for (i, &(x, y)) in pos.iter().enumerate() {
82 let c = cell_of(x, y);
83 items[fill[c] as usize] = i as u32;
84 fill[c] += 1;
85 }
86 Self {
87 min_x,
88 min_y,
89 inv_cell,
90 nx,
91 ny,
92 start,
93 items,
94 }
95 }
96
97 fn nearest(&self, pos: &[(f64, f64)], x: f64, y: f64, tol: f64) -> Option<usize> {
99 let gx0 = ((x - tol - self.min_x) * self.inv_cell).floor();
100 let gy0 = ((y - tol - self.min_y) * self.inv_cell).floor();
101 let gx1 = ((x + tol - self.min_x) * self.inv_cell).floor();
102 let gy1 = ((y + tol - self.min_y) * self.inv_cell).floor();
103 if gx1 < 0.0 || gy1 < 0.0 || gx0 >= self.nx as f64 || gy0 >= self.ny as f64 {
104 return None;
105 }
106 let (gx0, gy0) = (gx0.max(0.0) as usize, gy0.max(0.0) as usize);
107 let gx1 = (gx1 as usize).min(self.nx - 1);
108 let gy1 = (gy1 as usize).min(self.ny - 1);
109 let mut best = None;
110 let mut best_d2 = tol * tol;
111 for gy in gy0..=gy1 {
112 let row = gy * self.nx;
113 for c in row + gx0..=row + gx1 {
114 for &i in &self.items[self.start[c] as usize..self.start[c + 1] as usize] {
115 let (px, py) = pos[i as usize];
116 let d2 = (px - x) * (px - x) + (py - y) * (py - y);
117 if d2 <= best_d2 {
118 best_d2 = d2;
119 best = Some(i as usize);
120 }
121 }
122 }
123 }
124 best
125 }
126}
127
128const NEAR_MAX_CELLS_PER_SIDE: f64 = 1024.0;
132
133struct NearMap {
137 min_x: f32,
138 min_y: f32,
139 inv_cell: f32,
140 nx: usize,
141 ny: usize,
142 bits: Vec<u64>,
143}
144
145impl NearMap {
146 fn new(pos: &[(f64, f64)], grid: &PosGrid, probe_tol: f64) -> Self {
147 let (w, h) = (
148 grid.nx as f64 / grid.inv_cell,
149 grid.ny as f64 / grid.inv_cell,
150 );
151 let cell = (1.01 * probe_tol)
154 .max(w.max(h) / NEAR_MAX_CELLS_PER_SIDE)
155 .max(0.5);
156 let inv_cell = 1.0 / cell;
157 let nx = (w * inv_cell) as usize + 1;
158 let ny = (h * inv_cell) as usize + 1;
159 let mut bits = vec![0u64; (nx * ny).div_ceil(64)];
160 for &(x, y) in pos {
161 let gx = ((x - grid.min_x) * inv_cell) as usize;
162 let gy = ((y - grid.min_y) * inv_cell) as usize;
163 for cy in gy.saturating_sub(1)..=(gy + 1).min(ny - 1) {
164 for cx in gx.saturating_sub(1)..=(gx + 1).min(nx - 1) {
165 let c = cy * nx + cx;
166 bits[c / 64] |= 1 << (c % 64);
167 }
168 }
169 }
170 Self {
171 min_x: grid.min_x as f32,
172 min_y: grid.min_y as f32,
173 inv_cell: inv_cell as f32,
174 nx,
175 ny,
176 bits,
177 }
178 }
179
180 #[inline]
182 fn maybe(&self, x: f32, y: f32) -> bool {
183 let fx = (x - self.min_x) * self.inv_cell;
184 let fy = (y - self.min_y) * self.inv_cell;
185 if !(fx >= 0.0 && fy >= 0.0) {
186 return fx > -1.0 && fy > -1.0;
188 }
189 let (cx, cy) = (fx as usize, fy as usize);
190 if cx >= self.nx || cy >= self.ny {
191 return cx <= self.nx && cy <= self.ny;
192 }
193 let c = cy * self.nx + cx;
194 self.bits[c / 64] & (1 << (c % 64)) != 0
195 }
196}
197
198pub struct ImageIndex {
201 pos: Vec<(f64, f64)>,
202 grid: PosGrid,
203 near: NearMap,
205 pair_len: Vec<f32>,
207 pair_vec: Vec<[f32; 4]>,
209 pair_idx: Vec<(u32, u32)>,
211 probe_tol: f64,
212}
213
214impl ImageIndex {
215 #[must_use]
218 pub fn new(pos: Vec<(f64, f64)>, max_pair_px: f64, probe_tol: f64) -> Self {
219 let grid = PosGrid::new(&pos, probe_tol);
220 let near = NearMap::new(&pos, &grid, probe_tol);
221 let mut order: Vec<u32> = (0..pos.len() as u32).collect();
223 order.sort_unstable_by(|&a, &b| pos[a as usize].0.total_cmp(&pos[b as usize].0));
224 let max2 = max_pair_px * max_pair_px;
225 let min2 = MIN_PAIR_PX * MIN_PAIR_PX;
226 let mut pairs = Vec::new();
227 for (k, &a) in order.iter().enumerate() {
228 let (ax, ay) = pos[a as usize];
229 for &b in &order[k + 1..] {
230 let (bx, by) = pos[b as usize];
231 if bx - ax > max_pair_px {
232 break;
233 }
234 let d2 = (bx - ax) * (bx - ax) + (by - ay) * (by - ay);
235 if d2 >= min2 && d2 <= max2 {
236 pairs.push((d2.sqrt() as f32, a, b));
237 }
238 }
239 }
240 pairs.sort_unstable_by(|p, q| p.0.total_cmp(&q.0));
241 let pair_len = pairs.iter().map(|p| p.0).collect();
242 let pair_vec = pairs
243 .iter()
244 .map(|&(_, a, b)| {
245 let ((ax, ay), (bx, by)) = (pos[a as usize], pos[b as usize]);
246 [ax as f32, ay as f32, (bx - ax) as f32, (by - ay) as f32]
247 })
248 .collect();
249 let pair_idx = pairs.iter().map(|&(_, a, b)| (a, b)).collect();
250 Self {
251 pos,
252 grid,
253 near,
254 pair_len,
255 pair_vec,
256 pair_idx,
257 probe_tol,
258 }
259 }
260
261 #[must_use]
263 pub fn len(&self) -> usize {
264 self.pos.len()
265 }
266
267 #[must_use]
269 pub fn is_empty(&self) -> bool {
270 self.pos.is_empty()
271 }
272
273 #[must_use]
275 pub fn n_pairs(&self) -> usize {
276 self.pair_len.len()
277 }
278
279 fn hit(&self, x: f64, y: f64, tol: f64) -> Option<usize> {
280 self.grid.nearest(&self.pos, x, y, tol)
281 }
282}
283
284#[derive(Debug, Clone)]
286pub struct SeedParams {
287 pub scale: f64,
289 pub scale_tol: f64,
291 pub width: f64,
293 pub height: f64,
295 pub seed_stars: usize,
297 pub max_quads: usize,
299 pub census_stars: usize,
301 pub min_census: usize,
306 pub verify_cost: u64,
308}
309
310#[derive(Debug, Clone)]
313pub struct Candidate {
314 pub plate: PlateConstants,
316 pub img: Vec<(f64, f64)>,
318 pub cat: Vec<(f64, f64)>,
320}
321
322#[derive(Debug, Clone, Copy)]
325struct Similarity {
326 sr: f64,
327 si: f64,
328 tr: f64,
329 ti: f64,
330 mirrored: bool,
331}
332
333impl Similarity {
334 #[cfg(test)]
337 fn from_pair(
338 p1: (f64, f64),
339 p2: (f64, f64),
340 a: (f64, f64),
341 b: (f64, f64),
342 mirrored: bool,
343 ) -> Option<Self> {
344 let flip = |p: (f64, f64)| if mirrored { (p.0, -p.1) } else { p };
345 let (w1, w2) = (flip(p1), flip(p2));
346 let (dwr, dwi) = (w2.0 - w1.0, w2.1 - w1.1);
347 let den = dwr * dwr + dwi * dwi;
348 if den <= 0.0 {
349 return None;
350 }
351 let (dzr, dzi) = (b.0 - a.0, b.1 - a.1);
352 let sr = (dzr * dwr + dzi * dwi) / den;
354 let si = (dzi * dwr - dzr * dwi) / den;
355 let tr = a.0 - (sr * w1.0 - si * w1.1);
356 let ti = a.1 - (sr * w1.1 + si * w1.0);
357 Some(Self {
358 sr,
359 si,
360 tr,
361 ti,
362 mirrored,
363 })
364 }
365
366 #[inline]
367 fn apply(&self, p: (f64, f64)) -> (f64, f64) {
368 let wy = if self.mirrored { -p.1 } else { p.1 };
369 (
370 self.sr * p.0 - self.si * wy + self.tr,
371 self.sr * wy + self.si * p.0 + self.ti,
372 )
373 }
374
375 fn plate(&self) -> Option<PlateConstants> {
377 let det = self.sr * self.sr + self.si * self.si;
379 if det <= 0.0 {
380 return None;
381 }
382 let (ir, ii) = (self.sr / det, -self.si / det); let (a, b, c) = (ir, -ii, -(ir * self.tr - ii * self.ti));
385 let (d, e, f) = (ii, ir, -(ir * self.ti + ii * self.tr));
386 Some(if self.mirrored {
387 PlateConstants {
388 a,
389 b,
390 c,
391 d: -d,
392 e: -e,
393 f: -f,
394 }
395 } else {
396 PlateConstants { a, b, c, d, e, f }
397 })
398 }
399}
400
401fn seed_quads(cat: &[(f64, f64)], seed_stars: usize, max_quads: usize) -> Vec<[usize; 4]> {
404 let n = cat.len().min(seed_stars);
405 let mut quads = Vec::new();
406 for a in 0..n {
407 let mut near: Vec<(f64, usize)> = (0..n)
408 .filter(|&b| b != a)
409 .map(|b| {
410 let d = (cat[b].0 - cat[a].0).hypot(cat[b].1 - cat[a].1);
411 (d, b)
412 })
413 .collect();
414 near.sort_unstable_by(|p, q| p.0.total_cmp(&q.0));
415 near.truncate(4);
416 if near.len() < 3 {
417 continue;
418 }
419 for skip in 0..near.len() {
420 let mut q = [a; 4];
421 let mut k = 1;
422 for (m, &(_, b)) in near.iter().enumerate() {
423 if m != skip && k < 4 {
424 q[k] = b;
425 k += 1;
426 }
427 }
428 if k < 4 {
429 continue;
430 }
431 let mut widest = (0, 1, -1.0);
433 for i in 0..4 {
434 for j in i + 1..4 {
435 let d = (cat[q[i]].0 - cat[q[j]].0).hypot(cat[q[i]].1 - cat[q[j]].1);
436 if d > widest.2 {
437 widest = (i, j, d);
438 }
439 }
440 }
441 let rest: Vec<usize> = (0..4).filter(|&k| k != widest.0 && k != widest.1).collect();
442 quads.push([q[widest.0], q[widest.1], q[rest[0]], q[rest[1]]]);
443 if quads.len() >= max_quads {
444 break;
445 }
446 }
447 if quads.len() >= max_quads {
448 break;
449 }
450 }
451 quads
452}
453
454#[must_use]
457pub fn max_backbone_px(cat: &[(f64, f64)], p: &SeedParams) -> f64 {
458 let quads = seed_quads(cat, p.seed_stars, p.max_quads);
459 let longest = quads
460 .iter()
461 .map(|q| (cat[q[0]].0 - cat[q[1]].0).hypot(cat[q[0]].1 - cat[q[1]].1))
462 .fold(0.0, f64::max);
463 longest / (p.scale * (1.0 - p.scale_tol)) + 1.0
464}
465
466pub fn search(
474 index: &ImageIndex,
475 cat: &[(f64, f64)],
476 p: &SeedParams,
477 budget: &mut u64,
478 mut accept: impl FnMut(&Candidate) -> bool,
479) -> Option<Candidate> {
480 if index.len() < 4 || cat.len() < 4 {
481 return None;
482 }
483 let tol = index.probe_tol;
484 let census: Vec<(f64, f64)> = cat.iter().take(p.census_stars).copied().collect();
485 let margin = 1.5 * tol;
486 let density = index.len() as f64 / (p.width * p.height).max(1.0);
487 let p_chance = 1.0 - (-density * core::f64::consts::PI * margin * margin).exp();
488 let in_frame = |q: (f64, f64)| {
489 q.0 >= -margin && q.1 >= -margin && q.0 < p.width + margin && q.1 < p.height + margin
490 };
491 let mut seen: std::collections::HashSet<(i64, i64, i64, bool)> = Default::default();
493
494 for quad in seed_quads(cat, p.seed_stars, p.max_quads) {
495 let (p1, p2, p3, p4) = (cat[quad[0]], cat[quad[1]], cat[quad[2]], cat[quad[3]]);
496 let backbone = (p2.0 - p1.0).hypot(p2.1 - p1.1);
497 let lo = (backbone / (p.scale * (1.0 + p.scale_tol)) - tol) as f32;
498 let hi = (backbone / (p.scale * (1.0 - p.scale_tol)) + tol) as f32;
499 let from = index.pair_len.partition_point(|&l| l < lo);
500 let to = index.pair_len.partition_point(|&l| l <= hi);
501 for mirrored in [false, true] {
502 let flip = |q: (f64, f64)| if mirrored { (q.0, -q.1) } else { q };
506 let (w1, w2, w3, w4) = (flip(p1), flip(p2), flip(p3), flip(p4));
507 let (dwr, dwi) = (w2.0 - w1.0, w2.1 - w1.1);
508 let den = dwr * dwr + dwi * dwi;
509 if den <= 0.0 {
510 continue;
511 }
512 let (ir, ii) = (dwr / den, -dwi / den); let (d3, d4) = ((w3.0 - w1.0, w3.1 - w1.1), (w4.0 - w1.0, w4.1 - w1.1));
514 let (ir32, ii32) = (ir as f32, ii as f32);
516 let (d3x, d3y, d4x, d4y) = (d3.0 as f32, d3.1 as f32, d4.0 as f32, d4.1 as f32);
517 for (k, &[ax, ay, dx, dy]) in index.pair_vec[from..to].iter().enumerate() {
518 let (sr, si) = (dx * ir32 - dy * ii32, dx * ii32 + dy * ir32);
521 let (e3x, e3y) = (sr * d3x - si * d3y, sr * d3y + si * d3x);
522 let (e4x, e4y) = (sr * d4x - si * d4y, sr * d4y + si * d4x);
523 let (bx, by) = (ax + dx, ay + dy);
524 for (forward, zx, zy, sign) in [(true, ax, ay, 1.0f32), (false, bx, by, -1.0)] {
525 if *budget == 0 {
526 return None;
527 }
528 *budget -= 1;
529 if !(index.near.maybe(zx + sign * e3x, zy + sign * e3y)
530 && index.near.maybe(zx + sign * e4x, zy + sign * e4y))
531 {
532 continue;
533 }
534 let (i, j) = index.pair_idx[from + k];
536 let (a, b) = (index.pos[i as usize], index.pos[j as usize]);
537 let (za, zb) = if forward { (a, b) } else { (b, a) };
538 let (dzr, dzi) = (zb.0 - za.0, zb.1 - za.1);
539 let (sr, si) = (dzr * ir - dzi * ii, dzr * ii + dzi * ir);
540 let q3 = (za.0 + sr * d3.0 - si * d3.1, za.1 + sr * d3.1 + si * d3.0);
541 let q4 = (za.0 + sr * d4.0 - si * d4.1, za.1 + sr * d4.1 + si * d4.0);
542 if index.hit(q3.0, q3.1, tol).is_none() || index.hit(q4.0, q4.1, tol).is_none()
543 {
544 continue;
545 }
546 let t = Similarity {
547 sr,
548 si,
549 tr: za.0 - (sr * w1.0 - si * w1.1),
550 ti: za.1 - (sr * w1.1 + si * w1.0),
551 mirrored,
552 };
553 let key = (
554 (t.tr / 2.0).round() as i64,
555 (t.ti / 2.0).round() as i64,
556 (t.si.atan2(t.sr) * 200.0).round() as i64,
557 mirrored,
558 );
559 if !seen.insert(key) {
560 continue;
561 }
562 *budget = budget.saturating_sub(census.len() as u64);
564 let mut n_in = 0usize;
565 let mut img = Vec::new();
566 let mut catp = Vec::new();
567 let mut used = std::collections::HashSet::new();
568 for &c in &census {
569 let q = t.apply(c);
570 if !in_frame(q) {
571 continue;
572 }
573 n_in += 1;
574 if let Some(m) = index.hit(q.0, q.1, margin)
575 && used.insert(m)
576 {
577 img.push(index.pos[m]);
578 catp.push(c);
579 }
580 }
581 let need = (p.min_census.min((n_in * 2 / 3).max(4)) as f64)
583 .max(CENSUS_SIGNIFICANCE * p_chance * n_in.saturating_sub(4) as f64 + 4.0);
584 if (img.len() as f64) < need {
585 continue;
586 }
587 let plate = crate::math::lsq::fit_affine(&img, &catp)
588 .ok()
589 .or_else(|| t.plate());
590 let Some(plate) = plate else { continue };
591 let cand = Candidate {
592 plate,
593 img,
594 cat: catp,
595 };
596 *budget = budget.saturating_sub(p.verify_cost);
597 if accept(&cand) {
598 return Some(cand);
599 }
600 }
601 }
602 }
603 }
604 None
605}
606
607#[cfg(test)]
608mod tests {
609 use super::*;
610
611 fn lcg(seed: &mut u64) -> f64 {
612 *seed = seed
613 .wrapping_mul(6_364_136_223_846_793_005)
614 .wrapping_add(1_442_695_040_888_963_407);
615 (*seed >> 11) as f64 / (1u64 << 53) as f64
616 }
617
618 #[test]
619 fn similarity_maps_the_pair_and_inverts_to_a_plate() {
620 for mirrored in [false, true] {
621 let t = Similarity::from_pair(
622 (10.0, 20.0),
623 (110.0, -5.0),
624 (300.0, 400.0),
625 (350.0, 480.0),
626 mirrored,
627 )
628 .unwrap();
629 let a = t.apply((10.0, 20.0));
630 let b = t.apply((110.0, -5.0));
631 assert!((a.0 - 300.0).abs() < 1e-9 && (a.1 - 400.0).abs() < 1e-9);
632 assert!((b.0 - 350.0).abs() < 1e-9 && (b.1 - 480.0).abs() < 1e-9);
633 let pl = t.plate().unwrap();
634 let q = (37.0, -12.0);
635 let z = t.apply(q);
636 let back = (
637 pl.a * z.0 + pl.b * z.1 + pl.c,
638 pl.d * z.0 + pl.e * z.1 + pl.f,
639 );
640 assert!((back.0 - q.0).abs() < 1e-9 && (back.1 - q.1).abs() < 1e-9);
641 let det = pl.a * pl.e - pl.b * pl.d;
643 assert_eq!(det < 0.0, mirrored);
644 }
645 }
646
647 #[test]
650 fn finds_the_transform_whatever_the_image_ranking() {
651 let mut seed = 7u64;
652 let (w, h) = (1000.0, 800.0);
653 let scale = 2.0; let rot = 0.6f64;
655 let (c, s) = (rot.cos(), rot.sin());
656 let cat: Vec<(f64, f64)> = (0..200)
658 .map(|_| {
659 (
660 (lcg(&mut seed) - 0.5) * w * scale,
661 (lcg(&mut seed) - 0.5) * h * scale,
662 )
663 })
664 .collect();
665 let to_pix = |p: (f64, f64)| {
666 let (x, y) = (p.0 / scale, p.1 / scale);
667 (c * x - s * y + w / 2.0, s * x + c * y + h / 2.0)
668 };
669 let mut pos: Vec<(f64, f64)> = cat.iter().step_by(3).map(|&p| to_pix(p)).collect();
672 for _ in 0..1000 {
673 pos.push((lcg(&mut seed) * w, lcg(&mut seed) * h));
674 }
675 for i in (1..pos.len()).rev() {
676 let j = (lcg(&mut seed) * (i + 1) as f64) as usize;
677 pos.swap(i, j);
678 }
679 let p = SeedParams {
680 scale,
681 scale_tol: 0.05,
682 width: w,
683 height: h,
684 seed_stars: 100,
685 max_quads: 400,
686 census_stars: 100,
687 min_census: 10,
688 verify_cost: 0,
689 };
690 let index = ImageIndex::new(pos, max_backbone_px(&cat, &p), 2.5);
691 let mut budget = 50_000_000u64;
692 let found = search(&index, &cat, &p, &mut budget, |c| c.img.len() >= 25).expect("found");
695 let pl = &found.plate;
696 for &(x, y) in &[(0.0, 0.0), (w, 0.0), (0.0, h), (w, h)] {
698 let got = (pl.a * x + pl.b * y + pl.c, pl.d * x + pl.e * y + pl.f);
699 let (dx, dy) = (x - w / 2.0, y - h / 2.0);
700 let want = ((c * dx + s * dy) * scale, (-s * dx + c * dy) * scale);
701 assert!(
702 (got.0 - want.0).hypot(got.1 - want.1) < 2.0 * scale,
703 "corner ({x},{y}) {got:?} vs {want:?}"
704 );
705 }
706 }
707
708 #[test]
709 fn the_budget_bounds_a_search_that_cannot_succeed() {
710 let mut seed = 11u64;
711 let cat: Vec<(f64, f64)> = (0..200)
712 .map(|_| (lcg(&mut seed) * 2000.0, lcg(&mut seed) * 2000.0))
713 .collect();
714 let pos: Vec<(f64, f64)> = (0..800)
715 .map(|_| (lcg(&mut seed) * 1000.0, lcg(&mut seed) * 1000.0))
716 .collect();
717 let p = SeedParams {
718 scale: 2.0,
719 scale_tol: 0.1,
720 width: 1000.0,
721 height: 1000.0,
722 seed_stars: 100,
723 max_quads: 400,
724 census_stars: 100,
725 min_census: 10,
726 verify_cost: 0,
727 };
728 let index = ImageIndex::new(pos, max_backbone_px(&cat, &p), 2.5);
729 let mut budget = 1_000u64;
730 assert!(search(&index, &cat, &p, &mut budget, |_| false).is_none());
731 assert_eq!(budget, 0, "a refused search runs until the budget is spent");
732 }
733}