1use visual_cortex_capture::FrameView;
2
3use crate::detector::{Detector, DetectorError, DetectorOutput};
4use crate::luma::view_to_luma;
5
6const FLAT_EPSILON: f64 = 1e-6;
8const FLAT_MEAN_TOLERANCE: f64 = 2.0;
10
11pub struct TemplateMatcher {
27 width: u32,
28 height: u32,
29 mean: f64,
30 centered: Vec<f64>,
32 norm: f64,
34}
35
36impl TemplateMatcher {
37 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 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 ®ion[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 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 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(®ion, 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 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 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 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 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}