Skip to main content

visual_cortex_vision/
template.rs

1use visual_cortex_capture::FrameView;
2
3use crate::detector::{Detector, DetectorError, DetectorOutput};
4use crate::luma::view_to_luma;
5
6/// Variances below this are treated as "flat" (degenerate for NCC).
7const FLAT_EPSILON: f64 = 1e-6;
8/// Mean-luma difference (0..255 scale) considered equal for flat comparisons.
9const FLAT_MEAN_TOLERANCE: f64 = 2.0;
10
11/// Grayscale normalized cross-correlation template matcher.
12///
13/// `evaluate` slides the template over the region and returns the best match
14/// as `DetectorOutput::Score` in `0.0..=1.0` (negative correlation clamps to
15/// 0). Naive O(region * template) search — fine for the small regions and
16/// icons this library targets; optimize (integral images / FFT) only if
17/// profiling ever demands it.
18///
19/// Degenerate (flat) inputs fall back to mean comparison: a flat template
20/// scores 1.0 on a window of the same mean brightness, else 0.0.
21///
22/// NCC is contrast-invariant: a dim copy of the pattern scores as high as a
23/// bright one, so the score says nothing about absolute brightness/contrast.
24/// A template larger than the watched region scores 0.0 ("not found") rather
25/// than erroring — check your region size if a template watcher never fires.
26pub struct TemplateMatcher {
27    width: u32,
28    height: u32,
29    mean: f64,
30    /// Template luma with the mean subtracted, row-major.
31    centered: Vec<f64>,
32    /// sqrt(sum(centered^2)) — 0.0 for a flat template.
33    norm: f64,
34}
35
36impl TemplateMatcher {
37    /// Build from a row-major grayscale buffer. Errors if `luma.len()` does
38    /// not equal `width * height` or either dimension is zero.
39    pub fn from_luma(width: u32, height: u32, luma: Vec<u8>) -> Result<Self, DetectorError> {
40        let expected = width as usize * height as usize;
41        if width == 0 || height == 0 || luma.len() != expected {
42            return Err(DetectorError::Other(format!(
43                "template must be non-empty and {width}x{height} = {expected} bytes, got {}",
44                luma.len()
45            )));
46        }
47        let mean = luma.iter().map(|&p| p as f64).sum::<f64>() / expected as f64;
48        let centered: Vec<f64> = luma.iter().map(|&p| p as f64 - mean).collect();
49        let norm = centered.iter().map(|d| d * d).sum::<f64>().sqrt();
50        Ok(Self {
51            width,
52            height,
53            mean,
54            centered,
55            norm,
56        })
57    }
58
59    /// Build from encoded image bytes (PNG). Color images are converted to
60    /// grayscale; matching is luma-only.
61    /// Alpha is ignored (not composited); export icons with an opaque background.
62    pub fn from_png_bytes(bytes: &[u8]) -> Result<Self, DetectorError> {
63        let img = image::load_from_memory(bytes)
64            .map_err(|e| DetectorError::Other(format!("template decode failed: {e}")))?
65            .to_luma8();
66        let (w, h) = img.dimensions();
67        Self::from_luma(w, h, img.into_raw())
68    }
69
70    fn best_score(&self, region: &[u8], rw: usize, rh: usize) -> f32 {
71        let (tw, th) = (self.width as usize, self.height as usize);
72        if tw > rw || th > rh {
73            return 0.0;
74        }
75        let area = (tw * th) as f64;
76        let mut best = 0.0f64;
77        for oy in 0..=(rh - th) {
78            for ox in 0..=(rw - tw) {
79                let mut sum = 0u64;
80                for y in 0..th {
81                    let start = (oy + y) * rw + ox;
82                    for &p in &region[start..start + tw] {
83                        sum += p as u64;
84                    }
85                }
86                let wmean = sum as f64 / area;
87                let mut dot = 0.0f64;
88                let mut wnorm2 = 0.0f64;
89                for y in 0..th {
90                    let start = (oy + y) * rw + ox;
91                    for x in 0..tw {
92                        let wv = region[start + x] as f64 - wmean;
93                        dot += self.centered[y * tw + x] * wv;
94                        wnorm2 += wv * wv;
95                    }
96                }
97                let wnorm = wnorm2.sqrt();
98                let score = if self.norm < FLAT_EPSILON && wnorm < FLAT_EPSILON {
99                    // Both flat: match iff the brightness matches.
100                    if (self.mean - wmean).abs() <= FLAT_MEAN_TOLERANCE {
101                        1.0
102                    } else {
103                        0.0
104                    }
105                } else if self.norm < FLAT_EPSILON || wnorm < FLAT_EPSILON {
106                    // One flat, one textured: no meaningful correlation.
107                    0.0
108                } else {
109                    (dot / (self.norm * wnorm)).max(0.0)
110                };
111                if score > best {
112                    best = score;
113                }
114            }
115        }
116        best.min(1.0) as f32
117    }
118}
119
120impl Detector for TemplateMatcher {
121    fn evaluate(&mut self, view: &FrameView<'_>) -> Result<DetectorOutput, DetectorError> {
122        let (region, w, h) = view_to_luma(view);
123        Ok(DetectorOutput::Score(
124            self.best_score(&region, w as usize, h as usize),
125        ))
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132    use visual_cortex_capture::{Frame, PxRect};
133
134    fn full_view(frame: &Frame) -> FrameView<'_> {
135        frame
136            .view(PxRect {
137                x: 0,
138                y: 0,
139                w: frame.width(),
140                h: frame.height(),
141            })
142            .unwrap()
143    }
144
145    /// 4x4 black/white checkerboard as raw luma.
146    fn checker_luma() -> Vec<u8> {
147        (0..16u32)
148            .map(|i| {
149                let (x, y) = (i % 4, i / 4);
150                if (x + y) % 2 == 0 {
151                    255
152                } else {
153                    0
154                }
155            })
156            .collect()
157    }
158
159    /// 12x10 dark frame with the checkerboard stamped at (5, 3).
160    fn frame_with_icon() -> Frame {
161        Frame::from_fn(12, 10, |x, y| {
162            if (5..9).contains(&x) && (3..7).contains(&y) {
163                if ((x - 5) + (y - 3)) % 2 == 0 {
164                    [255, 255, 255, 255]
165                } else {
166                    [0, 0, 0, 255]
167                }
168            } else {
169                [30, 30, 30, 255]
170            }
171        })
172    }
173
174    #[test]
175    fn exact_match_scores_one() {
176        let mut det = TemplateMatcher::from_luma(4, 4, checker_luma()).unwrap();
177        let f = frame_with_icon();
178        let DetectorOutput::Score(s) = det.evaluate(&full_view(&f)).unwrap() else {
179            panic!("expected Score");
180        };
181        assert!(s > 0.99, "expected near-perfect match, got {s}");
182        assert!(s <= 1.0);
183    }
184
185    #[test]
186    fn absent_template_scores_low() {
187        let mut det = TemplateMatcher::from_luma(4, 4, checker_luma()).unwrap();
188        let f = Frame::solid(12, 10, [30, 30, 30, 255]);
189        let DetectorOutput::Score(s) = det.evaluate(&full_view(&f)).unwrap() else {
190            panic!("expected Score");
191        };
192        assert!(s < 0.1, "uniform region must not match, got {s}");
193    }
194
195    #[test]
196    fn template_larger_than_region_scores_zero() {
197        let mut det = TemplateMatcher::from_luma(8, 8, vec![128; 64]).unwrap();
198        let f = Frame::solid(4, 4, [30, 30, 30, 255]);
199        assert_eq!(
200            det.evaluate(&full_view(&f)).unwrap(),
201            DetectorOutput::Score(0.0)
202        );
203    }
204
205    #[test]
206    fn flat_template_matches_equal_flat_region_only() {
207        // Flat gray template: NCC is degenerate, falls back to mean comparison.
208        let mut det = TemplateMatcher::from_luma(2, 2, vec![100; 4]).unwrap();
209        let same = Frame::from_fn(4, 4, |_, _| [100, 100, 100, 255]);
210        let DetectorOutput::Score(s) = det.evaluate(&full_view(&same)).unwrap() else {
211            panic!("expected Score");
212        };
213        assert_eq!(s, 1.0);
214
215        let brighter = Frame::from_fn(4, 4, |_, _| [200, 200, 200, 255]);
216        let DetectorOutput::Score(s) = det.evaluate(&full_view(&brighter)).unwrap() else {
217            panic!("expected Score");
218        };
219        assert_eq!(s, 0.0);
220    }
221
222    #[test]
223    fn from_luma_validates_input() {
224        assert!(TemplateMatcher::from_luma(4, 4, vec![0; 15]).is_err());
225        assert!(TemplateMatcher::from_luma(0, 4, vec![]).is_err());
226    }
227
228    /// The checkerboard encoded as PNG bytes, via the image crate.
229    fn checker_png() -> Vec<u8> {
230        let img = image::ImageBuffer::from_fn(4, 4, |x, y| {
231            if (x + y) % 2 == 0 {
232                image::Luma([255u8])
233            } else {
234                image::Luma([0u8])
235            }
236        });
237        let mut out = Vec::new();
238        img.write_to(&mut std::io::Cursor::new(&mut out), image::ImageFormat::Png)
239            .expect("png encode");
240        out
241    }
242
243    #[test]
244    fn png_roundtrip_matches_like_raw_luma() {
245        let mut det = TemplateMatcher::from_png_bytes(&checker_png()).unwrap();
246        let f = frame_with_icon();
247        let DetectorOutput::Score(s) = det.evaluate(&full_view(&f)).unwrap() else {
248            panic!("expected Score");
249        };
250        assert!(s > 0.99, "png-loaded template must match, got {s}");
251    }
252
253    #[test]
254    fn invalid_png_bytes_error() {
255        assert!(TemplateMatcher::from_png_bytes(b"definitely not a png").is_err());
256    }
257}