1use serde::Serialize;
18use thiserror::Error;
19
20use crate::geometry::{Point, Shape, Size};
21use crate::session::SessionFile;
22
23pub const FIND_SCHEMA_VERSION: u32 = 1;
24
25pub const SCORE_FLOOR: f64 = 0.9;
29
30pub const AMBIGUITY_GAP: f64 = 0.03;
33
34#[derive(Debug, Clone)]
36pub struct GrayImage {
37 pub w: usize,
38 pub h: usize,
39 pub px: Vec<f32>,
40}
41
42impl GrayImage {
43 pub fn from_rgba(w: usize, h: usize, rgba: &[u8]) -> Self {
45 assert_eq!(rgba.len(), w * h * 4, "rgba buffer matches dimensions");
46 let px = rgba
47 .chunks_exact(4)
48 .map(|p| luma(p[0], p[1], p[2]))
49 .collect();
50 Self { w, h, px }
51 }
52}
53
54fn luma(r: u8, g: u8, b: u8) -> f32 {
55 (0.299 * f32::from(r) + 0.587 * f32::from(g) + 0.114 * f32::from(b)) / 255.0
56}
57
58#[derive(Debug, Clone)]
62pub struct Template {
63 pub gray: GrayImage,
64 pub mask: Vec<bool>,
65}
66
67impl Template {
68 pub fn from_rgba(w: usize, h: usize, rgba: &[u8]) -> Self {
70 let gray = GrayImage::from_rgba(w, h, rgba);
71 let mask = rgba.chunks_exact(4).map(|p| p[3] >= 128).collect();
72 Self { gray, mask }
73 }
74}
75
76#[derive(Debug, Clone, Copy, PartialEq)]
78pub struct Located {
79 pub x: i32,
81 pub y: i32,
82 pub score: f64,
84 pub runner_up: f64,
87 pub ambiguous: bool,
90}
91
92#[derive(Debug, Error, PartialEq, Eq)]
93pub enum LocateError {
94 #[error("the crop is larger than the capture it is searched in")]
95 TemplateLargerThanScreen,
96 #[error("the crop has no visible pixels to match")]
97 EmptyTemplate,
98 #[error("the crop is a flat color, which matches anywhere rather than somewhere")]
99 FlatTemplate,
100}
101
102struct TemplateStats {
104 coords: Vec<(usize, usize)>,
105 values: Vec<f64>,
106 mean: f64,
107 var: f64,
108}
109
110fn template_stats(tpl: &Template) -> Result<TemplateStats, LocateError> {
111 let mut coords = Vec::new();
112 let mut values = Vec::new();
113 for ty in 0..tpl.gray.h {
114 for tx in 0..tpl.gray.w {
115 if !tpl.mask[ty * tpl.gray.w + tx] {
116 continue;
117 }
118 coords.push((tx, ty));
119 values.push(f64::from(tpl.gray.px[ty * tpl.gray.w + tx]));
120 }
121 }
122 if coords.is_empty() {
123 return Err(LocateError::EmptyTemplate);
124 }
125 let n = values.len() as f64;
126 let mean = values.iter().sum::<f64>() / n;
127 let var = values.iter().map(|v| (v - mean) * (v - mean)).sum::<f64>();
128 if var <= f64::EPSILON {
129 return Err(LocateError::FlatTemplate);
130 }
131 Ok(TemplateStats {
132 coords,
133 values,
134 mean,
135 var,
136 })
137}
138
139fn ncc_at(screen: &GrayImage, stats: &TemplateStats, ox: usize, oy: usize) -> f64 {
141 let n = stats.values.len() as f64;
142 let mut sum_screen = 0.0;
143 let mut sum_screen_sq = 0.0;
144 let mut sum_cross = 0.0;
145 for (i, &(tx, ty)) in stats.coords.iter().enumerate() {
146 let s = f64::from(screen.px[(oy + ty) * screen.w + ox + tx]);
147 sum_screen += s;
148 sum_screen_sq += s * s;
149 sum_cross += stats.values[i] * s;
150 }
151 let mean_screen = sum_screen / n;
152 let var_screen = sum_screen_sq - n * mean_screen * mean_screen;
153 if var_screen <= f64::EPSILON {
154 return 0.0;
155 }
156 let cov = sum_cross - n * stats.mean * mean_screen;
157 cov / (stats.var * var_screen).sqrt()
158}
159
160fn best_in_window(
162 screen: &GrayImage,
163 stats: &TemplateStats,
164 x0: usize,
165 y0: usize,
166 x1: usize,
167 y1: usize,
168) -> (usize, usize, f64) {
169 let mut best = (x0, y0, f64::from(f32::MIN));
170 for oy in y0..=y1 {
171 for ox in x0..=x1 {
172 let score = ncc_at(screen, stats, ox, oy);
173 if score > best.2 {
174 best = (ox, oy, score);
175 }
176 }
177 }
178 best
179}
180
181pub fn locate(
184 screen: &GrayImage,
185 tpl: &Template,
186 hint: Option<Point>,
187) -> Result<Located, LocateError> {
188 let (tw, th) = (tpl.gray.w, tpl.gray.h);
189 if tw > screen.w || th > screen.h || tw == 0 || th == 0 {
190 return Err(LocateError::TemplateLargerThanScreen);
191 }
192 let stats = template_stats(tpl)?;
193 let max_x = screen.w - tw;
194 let max_y = screen.h - th;
195
196 let factor = match tw.min(th) {
201 0..16 => 1,
202 16..32 => 2,
203 _ => 4,
204 };
205 let mut candidates = coarse_candidates(screen, tpl, factor)?;
206 if let Some(p) = hint {
207 candidates.push((
208 p.x.clamp(0, max_x as i32) as usize,
209 p.y.clamp(0, max_y as i32) as usize,
210 ));
211 }
212
213 let margin = factor * 2;
215 let mut refined: Vec<(usize, usize, f64)> = candidates
216 .iter()
217 .map(|&(cx, cy)| {
218 let x0 = cx.saturating_sub(margin);
219 let y0 = cy.saturating_sub(margin);
220 best_in_window(
221 screen,
222 &stats,
223 x0,
224 y0,
225 (cx + margin).min(max_x),
226 (cy + margin).min(max_y),
227 )
228 })
229 .collect();
230 refined.sort_by(|a, b| b.2.total_cmp(&a.2));
231 let best = refined[0];
232
233 let separated = |x: usize, y: usize| x.abs_diff(best.0) > tw / 2 || y.abs_diff(best.1) > th / 2;
236 let runner_up = refined[1..]
237 .iter()
238 .filter(|&&(x, y, _)| separated(x, y))
239 .map(|&(_, _, s)| s)
240 .fold(-1.0f64, f64::max);
241
242 let ambiguous = runner_up >= SCORE_FLOOR && best.2 - runner_up <= AMBIGUITY_GAP;
243 Ok(Located {
244 x: best.0 as i32,
245 y: best.1 as i32,
246 score: best.2,
247 runner_up,
248 ambiguous,
249 })
250}
251
252fn coarse_candidates(
257 screen: &GrayImage,
258 tpl: &Template,
259 factor: usize,
260) -> Result<Vec<(usize, usize)>, LocateError> {
261 let (small_screen, small_stats) = if factor == 1 {
262 (screen.clone(), template_stats(tpl)?)
263 } else {
264 let small_tpl = downsample_template(tpl, factor);
265 (downsample(screen, factor), template_stats(&small_tpl)?)
266 };
267 let tw = small_stats.coords.iter().map(|c| c.0).max().unwrap_or(0) + 1;
268 let th = small_stats.coords.iter().map(|c| c.1).max().unwrap_or(0) + 1;
269 if tw > small_screen.w || th > small_screen.h {
270 return Err(LocateError::TemplateLargerThanScreen);
271 }
272 let max_x = small_screen.w - tw;
273 let max_y = small_screen.h - th;
274
275 let mut scores = vec![0.0f64; (max_x + 1) * (max_y + 1)];
276 for oy in 0..=max_y {
277 for ox in 0..=max_x {
278 scores[oy * (max_x + 1) + ox] = ncc_at(&small_screen, &small_stats, ox, oy);
279 }
280 }
281 let mut out: Vec<(usize, usize)> = Vec::new();
286 for _ in 0..3 {
287 let mut best = (0usize, 0usize, f64::from(f32::MIN));
288 for oy in 0..=max_y {
289 for ox in 0..=max_x {
290 let suppressed = out.iter().any(|&(px, py)| {
291 ox.abs_diff(px / factor) <= tw / 2 && oy.abs_diff(py / factor) <= th / 2
292 });
293 if suppressed {
294 continue;
295 }
296 let s = scores[oy * (max_x + 1) + ox];
297 if s > best.2 {
298 best = (ox, oy, s);
299 }
300 }
301 }
302 if best.2 <= f64::from(f32::MIN) {
303 break;
304 }
305 out.push((best.0 * factor, best.1 * factor));
306 }
307 Ok(out)
308}
309
310fn downsample(img: &GrayImage, factor: usize) -> GrayImage {
312 let w = img.w / factor;
313 let h = img.h / factor;
314 let mut px = Vec::with_capacity(w * h);
315 for y in 0..h {
316 for x in 0..w {
317 let mut sum = 0.0f32;
318 for sy in 0..factor {
319 for sx in 0..factor {
320 sum += img.px[(y * factor + sy) * img.w + x * factor + sx];
321 }
322 }
323 px.push(sum / (factor * factor) as f32);
324 }
325 }
326 GrayImage { w, h, px }
327}
328
329fn downsample_template(tpl: &Template, factor: usize) -> Template {
332 let w = tpl.gray.w / factor;
333 let h = tpl.gray.h / factor;
334 let mut px = Vec::with_capacity(w * h);
335 let mut mask = Vec::with_capacity(w * h);
336 for y in 0..h {
337 for x in 0..w {
338 let mut sum = 0.0f32;
339 let mut n = 0usize;
340 for sy in 0..factor {
341 for sx in 0..factor {
342 let i = (y * factor + sy) * tpl.gray.w + x * factor + sx;
343 if !tpl.mask[i] {
344 continue;
345 }
346 sum += tpl.gray.px[i];
347 n += 1;
348 }
349 }
350 let keep = n * 2 >= factor * factor;
351 mask.push(keep);
352 px.push(if keep { sum / n as f32 } else { 0.0 });
353 }
354 }
355 Template {
356 gray: GrayImage { w, h, px },
357 mask,
358 }
359}
360
361pub fn crop_origin(shape: &Shape, rot_deg: i32, frame: Size) -> Point {
364 let bbox = shape.rotated_bbox(rot_deg);
365 Point::new(bbox.x.max(0).min(frame.w), bbox.y.max(0).min(frame.h))
366}
367
368pub struct Relocation {
370 pub crop_origin: Point,
372 pub outcome: Result<Located, LocateError>,
373}
374
375#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
377pub struct Delta {
378 pub dx: i32,
379 pub dy: i32,
380}
381
382#[derive(Debug, Clone, PartialEq, Serialize)]
384pub struct FindResult {
385 pub index: usize,
386 pub label: String,
387 pub monitor: usize,
388 pub found: bool,
389 pub ambiguous: bool,
390 pub score: f64,
391 #[serde(skip_serializing_if = "Option::is_none")]
392 pub reason: Option<String>,
393 pub old_px: Shape,
394 #[serde(skip_serializing_if = "Option::is_none")]
395 pub new_px: Option<Shape>,
396 #[serde(skip_serializing_if = "Option::is_none")]
397 pub new_global_px: Option<Shape>,
398 #[serde(skip_serializing_if = "Option::is_none")]
399 pub delta: Option<Delta>,
400}
401
402#[derive(Debug, Clone, PartialEq, Serialize)]
404pub struct FindReport {
405 pub schema: u32,
406 pub captured_utc: String,
407 pub all_relocated: bool,
409 pub results: Vec<FindResult>,
410}
411
412pub fn report(
417 session: &SessionFile,
418 relocations: &[(usize, Relocation)],
419 captured_utc: String,
420) -> FindReport {
421 let results: Vec<FindResult> = relocations
422 .iter()
423 .map(|(index, reloc)| {
424 let index = *index;
425 let record = session
426 .selections
427 .get(index)
428 .expect("relocation index within the session");
429 let base = FindResult {
430 index,
431 label: record.label.clone(),
432 monitor: record.monitor,
433 found: false,
434 ambiguous: false,
435 score: 0.0,
436 reason: None,
437 old_px: record.px.clone(),
438 new_px: None,
439 new_global_px: None,
440 delta: None,
441 };
442 match &reloc.outcome {
443 Err(e) => FindResult {
444 reason: Some(e.to_string()),
445 ..base
446 },
447 Ok(loc) => {
448 let found = loc.score >= SCORE_FLOOR;
449 let moved = found && !loc.ambiguous;
450 let dx = loc.x - reloc.crop_origin.x;
451 let dy = loc.y - reloc.crop_origin.y;
452 FindResult {
453 found,
454 ambiguous: found && loc.ambiguous,
455 score: loc.score,
456 new_px: moved.then(|| record.px.translated(dx, dy)),
457 new_global_px: moved.then(|| record.global_px.translated(dx, dy)),
458 delta: moved.then_some(Delta { dx, dy }),
459 ..base
460 }
461 }
462 }
463 })
464 .collect();
465 let all_relocated = !results.is_empty() && results.iter().all(|r| r.found && !r.ambiguous);
466 FindReport {
467 schema: FIND_SCHEMA_VERSION,
468 captured_utc,
469 all_relocated,
470 results,
471 }
472}
473
474#[cfg(test)]
475mod tests {
476 use super::*;
477 use crate::geometry::Rect;
478
479 fn textured(w: usize, h: usize, seed: u32) -> GrayImage {
484 let mut state = seed | 1;
485 let mut px: Vec<f32> = (0..w * h)
486 .map(|_| {
487 state ^= state << 13;
489 state ^= state >> 17;
490 state ^= state << 5;
491 (state % 1000) as f32 / 1000.0
492 })
493 .collect();
494 for _ in 0..2 {
495 px = blur3(&px, w, h);
496 }
497 GrayImage { w, h, px }
498 }
499
500 fn blur3(px: &[f32], w: usize, h: usize) -> Vec<f32> {
501 let mut out = Vec::with_capacity(w * h);
502 for y in 0..h {
503 for x in 0..w {
504 out.push(blurred_at(px, w, h, x, y));
505 }
506 }
507 out
508 }
509
510 fn blurred_at(px: &[f32], width: usize, height: usize, col: usize, row: usize) -> f32 {
511 let mut sum = 0.0f32;
512 let mut count = 0.0f32;
513 for dy in -1i32..=1 {
514 for dx in -1i32..=1 {
515 let nx = col as i32 + dx;
516 let ny = row as i32 + dy;
517 if nx < 0 || ny < 0 || nx >= width as i32 || ny >= height as i32 {
518 continue;
519 }
520 sum += px[ny as usize * width + nx as usize];
521 count += 1.0;
522 }
523 }
524 sum / count
525 }
526
527 fn cut(img: &GrayImage, x: usize, y: usize, w: usize, h: usize) -> Template {
528 let mut px = Vec::with_capacity(w * h);
529 for ty in 0..h {
530 for tx in 0..w {
531 px.push(img.px[(y + ty) * img.w + x + tx]);
532 }
533 }
534 Template {
535 gray: GrayImage { w, h, px },
536 mask: vec![true; w * h],
537 }
538 }
539
540 #[test]
541 fn a_cut_template_is_found_where_it_was_cut() {
542 let screen = textured(200, 150, 7);
543 let tpl = cut(&screen, 63, 41, 24, 18);
544 let loc = locate(&screen, &tpl, Some(Point::new(5, 5))).unwrap();
546 assert_eq!((loc.x, loc.y), (63, 41));
547 assert!(loc.score > 0.999, "exact cut scores ~1, got {}", loc.score);
548 assert!(!loc.ambiguous);
549 }
550
551 #[test]
552 fn a_large_template_takes_the_pyramid_path_and_still_lands_exactly() {
553 let screen = textured(400, 300, 11);
554 let tpl = cut(&screen, 137, 92, 80, 64);
555 let loc = locate(&screen, &tpl, None).unwrap();
556 assert_eq!((loc.x, loc.y), (137, 92));
557 assert!(loc.score > 0.999);
558 }
559
560 #[test]
561 fn a_masked_template_ignores_its_transparent_corners() {
562 let screen = textured(160, 120, 23);
563 let mut tpl = cut(&screen, 50, 40, 32, 32);
564 for (i, m) in tpl.mask.iter_mut().enumerate() {
567 let (x, y) = (i % 32, i / 32);
568 let (cx, cy) = (16i32, 16i32);
569 let (dx, dy) = (x as i32 - cx, y as i32 - cy);
570 if dx * dx + dy * dy > 16 * 16 {
571 *m = false;
572 }
573 }
574 for (i, p) in tpl.gray.px.iter_mut().enumerate() {
575 if !tpl.mask[i] {
576 *p = 1.0 - *p;
577 }
578 }
579 let loc = locate(&screen, &tpl, None).unwrap();
580 assert_eq!((loc.x, loc.y), (50, 40));
581 assert!(loc.score > 0.999);
582 }
583
584 #[test]
585 fn changed_pixels_score_below_the_floor() {
586 let screen = textured(160, 120, 31);
587 let tpl = cut(&screen, 50, 40, 24, 24);
588 let elsewhere = textured(160, 120, 97);
589 let loc = locate(&elsewhere, &tpl, None).unwrap();
590 assert!(
591 loc.score < SCORE_FLOOR,
592 "unrelated content must not match, got {}",
593 loc.score
594 );
595 }
596
597 #[test]
598 fn a_duplicated_region_is_ambiguous_not_silently_resolved() {
599 let mut screen = textured(300, 100, 43);
600 for ty in 0..24 {
602 for tx in 0..24 {
603 let v = screen.px[(30 + ty) * 300 + 20 + tx];
604 screen.px[(30 + ty) * 300 + 200 + tx] = v;
605 }
606 }
607 let tpl = cut(&screen, 20, 30, 24, 24);
608 let loc = locate(&screen, &tpl, None).unwrap();
609 assert!(loc.ambiguous, "two identical regions: {loc:?}");
610 assert!(loc.runner_up > 0.999);
611 }
612
613 #[test]
614 fn degenerate_templates_are_refused_with_reasons() {
615 let screen = textured(50, 50, 3);
616 let flat = Template {
617 gray: GrayImage {
618 w: 8,
619 h: 8,
620 px: vec![0.5; 64],
621 },
622 mask: vec![true; 64],
623 };
624 assert_eq!(
625 locate(&screen, &flat, None).unwrap_err(),
626 LocateError::FlatTemplate
627 );
628 let empty = Template {
629 gray: GrayImage {
630 w: 8,
631 h: 8,
632 px: vec![0.5; 64],
633 },
634 mask: vec![false; 64],
635 };
636 assert_eq!(
637 locate(&screen, &empty, None).unwrap_err(),
638 LocateError::EmptyTemplate
639 );
640 let huge = cut(&textured(80, 80, 5), 0, 0, 80, 80);
641 assert_eq!(
642 locate(&screen, &huge, None).unwrap_err(),
643 LocateError::TemplateLargerThanScreen
644 );
645 }
646
647 #[test]
648 fn rgba_conversion_masks_transparency_and_weights_luma() {
649 let rgba = [0u8, 255, 0, 255, 255, 255, 255, 0];
651 let tpl = Template::from_rgba(2, 1, &rgba);
652 assert!(tpl.mask[0] && !tpl.mask[1]);
653 assert!((tpl.gray.px[0] - 0.587).abs() < 1e-4);
654 assert!((tpl.gray.px[1] - 1.0).abs() < 1e-4);
655 }
656
657 #[test]
658 fn crop_origin_is_the_clipped_rotated_bbox() {
659 let frame = Size::new(1920, 1080);
660 assert_eq!(
661 crop_origin(&Shape::Rect(Rect::new(100, 50, 40, 30)), 0, frame),
662 Point::new(100, 50)
663 );
664 assert_eq!(
667 crop_origin(&Shape::Rect(Rect::new(100, 50, 40, 30)), 90, frame),
668 Point::new(105, 45)
669 );
670 assert_eq!(
672 crop_origin(&Shape::Rect(Rect::new(-20, -10, 40, 30)), 0, frame),
673 Point::new(0, 0)
674 );
675 }
676
677 fn session_of_one() -> SessionFile {
678 use crate::selection::Selection;
679 use crate::session::MonitorRecord;
680 let mut sel = Selection::new(Shape::Rect(Rect::new(10, 20, 30, 40)), 0);
681 sel.label = "submit".into();
682 SessionFile::build(
683 "test",
684 "2026-07-27T00:00:00Z".into(),
685 vec![MonitorRecord {
686 index: 0,
687 name: "Main".into(),
688 primary: true,
689 origin_px: Point::new(100, 0),
690 size_px: Size::new(1920, 1080),
691 scale: 1.0,
692 }],
693 &[sel],
694 &["c0.png".into()],
695 None,
696 )
697 }
698
699 #[test]
700 fn report_translates_found_selections_and_flags_the_rest() {
701 let session = session_of_one();
702 let found = report(
703 &session,
704 &[(
705 0,
706 Relocation {
707 crop_origin: Point::new(10, 20),
708 outcome: Ok(Located {
709 x: 14,
710 y: 8,
711 score: 0.97,
712 runner_up: 0.1,
713 ambiguous: false,
714 }),
715 },
716 )],
717 "2026-07-27T01:00:00Z".into(),
718 );
719 assert!(found.all_relocated);
720 let r = &found.results[0];
721 assert_eq!(r.delta, Some(Delta { dx: 4, dy: -12 }));
722 assert_eq!(r.new_px, Some(Shape::Rect(Rect::new(14, 8, 30, 40))));
723 assert_eq!(
724 r.new_global_px,
725 Some(Shape::Rect(Rect::new(114, 8, 30, 40))),
726 "global keeps the monitor origin offset"
727 );
728
729 let miss = report(
730 &session,
731 &[(
732 0,
733 Relocation {
734 crop_origin: Point::new(10, 20),
735 outcome: Ok(Located {
736 x: 0,
737 y: 0,
738 score: 0.4,
739 runner_up: 0.1,
740 ambiguous: false,
741 }),
742 },
743 )],
744 "t".into(),
745 );
746 assert!(!miss.all_relocated);
747 assert!(!miss.results[0].found);
748 assert!(miss.results[0].new_px.is_none());
749
750 let ambiguous = report(
751 &session,
752 &[(
753 0,
754 Relocation {
755 crop_origin: Point::new(10, 20),
756 outcome: Ok(Located {
757 x: 14,
758 y: 8,
759 score: 0.99,
760 runner_up: 0.98,
761 ambiguous: true,
762 }),
763 },
764 )],
765 "t".into(),
766 );
767 assert!(!ambiguous.all_relocated);
768 assert!(ambiguous.results[0].found && ambiguous.results[0].ambiguous);
769 assert!(
770 ambiguous.results[0].new_px.is_none(),
771 "an ambiguous match must not hand out coordinates"
772 );
773
774 let errored = report(
775 &session,
776 &[(
777 0,
778 Relocation {
779 crop_origin: Point::new(10, 20),
780 outcome: Err(LocateError::FlatTemplate),
781 },
782 )],
783 "t".into(),
784 );
785 assert!(!errored.all_relocated);
786 assert!(
787 errored.results[0]
788 .reason
789 .as_deref()
790 .unwrap()
791 .contains("flat color")
792 );
793 }
794
795 #[test]
796 fn find_report_json_shape_is_stable() {
797 let session = session_of_one();
798 let rep = report(
799 &session,
800 &[(
801 0,
802 Relocation {
803 crop_origin: Point::new(10, 20),
804 outcome: Ok(Located {
805 x: 10,
806 y: 20,
807 score: 1.0,
808 runner_up: 0.0,
809 ambiguous: false,
810 }),
811 },
812 )],
813 "2026-07-27T01:00:00Z".into(),
814 );
815 let json = serde_json::to_value(&rep).unwrap();
816 assert_eq!(json["schema"], 1);
817 assert_eq!(json["all_relocated"], true);
818 assert_eq!(json["results"][0]["label"], "submit");
819 assert_eq!(json["results"][0]["delta"]["dx"], 0);
820 assert_eq!(json["results"][0]["new_px"]["x"], 10);
821 assert!(json["results"][0].get("reason").is_none());
822 }
823}