visual_cortex_vision/
ocr.rs

1use visual_cortex_capture::{FrameView, PxRect};
2
3use crate::detector::DetectorError;
4
5/// One recognized text region. `bbox` is relative to the evaluated view.
6#[derive(Debug, Clone, PartialEq)]
7pub struct TextSpan {
8    pub text: String,
9    /// Mean character confidence in 0.0..=1.0. Never NaN.
10    pub confidence: f32,
11    pub bbox: PxRect,
12}
13
14/// Pluggable OCR backend. The default ONNX/PaddleOCR implementation lives in
15/// the `visual-cortex-ocr-onnx` crate; any engine can slot in here. Returning an empty
16/// vec means "no text present" and is not an error.
17pub trait OcrEngine: Send + 'static {
18    fn recognize(&mut self, view: &FrameView<'_>) -> Result<Vec<TextSpan>, DetectorError>;
19}
20
21/// Boxed engines are engines, so builder plumbing can defer construction.
22impl OcrEngine for Box<dyn OcrEngine> {
23    fn recognize(&mut self, view: &FrameView<'_>) -> Result<Vec<TextSpan>, DetectorError> {
24        (**self).recognize(view)
25    }
26}
27
28use crate::detector::{Detector, DetectorOutput};
29
30type NumberParser = Box<dyn Fn(&str) -> Option<f64> + Send + 'static>;
31
32enum OcrMode {
33    Text,
34    Number(NumberParser),
35    Spans,
36}
37
38/// Adapts any [`OcrEngine`] into a [`Detector`].
39///
40/// All modes sort spans top-to-bottom then left-to-right, and an empty span
41/// list becomes `DetectorOutput::None` — so, under `Mode::Change`, "text
42/// appeared" is the `None` -> `Text`/`Spans` transition.
43/// Text mode joins span texts with `\n` into `DetectorOutput::Text`.
44/// Number mode additionally parses the joined text; parse failures and
45/// non-finite values become `DetectorOutput::None` (the never-NaN invariant).
46/// Spans mode preserves per-span geometry and confidence in
47/// `DetectorOutput::Spans` (whose equality is text-only — see the variant
48/// docs — so OCR jitter does not fire change events).
49type SpanFilter = Box<dyn Fn(&TextSpan) -> bool + Send + 'static>;
50
51pub struct OcrDetector {
52    engine: Box<dyn OcrEngine>,
53    mode: OcrMode,
54    filters: Vec<SpanFilter>,
55}
56
57impl OcrDetector {
58    pub fn text(engine: impl OcrEngine) -> Self {
59        Self {
60            engine: Box::new(engine),
61            mode: OcrMode::Text,
62            filters: Vec::new(),
63        }
64    }
65
66    /// Add a span filter (`true` = keep), applied BEFORE sorting, joining,
67    /// and change detection — a filtered span can never fire a `Changed`
68    /// event. Filters compose; a span must pass all of them. Stock filters
69    /// live in [`crate::filters`].
70    pub fn with_span_filter(mut self, filter: impl Fn(&TextSpan) -> bool + Send + 'static) -> Self {
71        self.filters.push(Box::new(filter));
72        self
73    }
74
75    pub fn number(
76        engine: impl OcrEngine,
77        parser: impl Fn(&str) -> Option<f64> + Send + 'static,
78    ) -> Self {
79        Self {
80            engine: Box::new(engine),
81            mode: OcrMode::Number(Box::new(parser)),
82            filters: Vec::new(),
83        }
84    }
85
86    pub fn spans(engine: impl OcrEngine) -> Self {
87        Self {
88            engine: Box::new(engine),
89            mode: OcrMode::Spans,
90            filters: Vec::new(),
91        }
92    }
93}
94
95impl Detector for OcrDetector {
96    fn evaluate(&mut self, view: &FrameView<'_>) -> Result<DetectorOutput, DetectorError> {
97        let mut spans = self.engine.recognize(view)?;
98        spans.retain(|s| self.filters.iter().all(|keep| keep(s)));
99        if spans.is_empty() {
100            return Ok(DetectorOutput::None);
101        }
102        spans.sort_by_key(|s| (s.bbox.y, s.bbox.x));
103        if matches!(self.mode, OcrMode::Spans) {
104            return Ok(DetectorOutput::Spans(spans));
105        }
106        let joined = spans
107            .iter()
108            .map(|s| s.text.as_str())
109            .collect::<Vec<_>>()
110            .join("\n");
111        match &self.mode {
112            OcrMode::Text => Ok(DetectorOutput::Text(joined)),
113            OcrMode::Number(parser) => Ok(match parser(&joined) {
114                Some(n) if n.is_finite() => DetectorOutput::Number(n),
115                _ => DetectorOutput::None,
116            }),
117            OcrMode::Spans => unreachable!("handled above"),
118        }
119    }
120
121    fn is_heavy(&self) -> bool {
122        true
123    }
124}
125
126#[cfg(test)]
127mod tests {
128    use super::*;
129    use visual_cortex_capture::{Frame, PxRect};
130
131    struct FixedOcr(Vec<TextSpan>);
132
133    impl OcrEngine for FixedOcr {
134        fn recognize(&mut self, _view: &FrameView<'_>) -> Result<Vec<TextSpan>, DetectorError> {
135            Ok(self.0.clone())
136        }
137    }
138
139    #[test]
140    fn trait_is_object_safe_and_returns_spans() {
141        let span = TextSpan {
142            text: "HP 300".to_string(),
143            confidence: 0.98,
144            bbox: PxRect {
145                x: 1,
146                y: 2,
147                w: 40,
148                h: 10,
149            },
150        };
151        let mut engine: Box<dyn OcrEngine> = Box::new(FixedOcr(vec![span]));
152        let frame = Frame::solid(64, 16, [0, 0, 0, 255]);
153        let view = frame
154            .view(PxRect {
155                x: 0,
156                y: 0,
157                w: 64,
158                h: 16,
159            })
160            .unwrap();
161        let spans = engine.recognize(&view).unwrap();
162        assert_eq!(spans.len(), 1);
163        assert_eq!(spans[0].text, "HP 300");
164        assert!(spans[0].confidence > 0.9);
165        assert_eq!(spans[0].bbox.w, 40);
166    }
167
168    fn view_frame() -> Frame {
169        Frame::solid(64, 32, [0, 0, 0, 255])
170    }
171
172    fn full_view(frame: &Frame) -> FrameView<'_> {
173        frame
174            .view(PxRect {
175                x: 0,
176                y: 0,
177                w: frame.width(),
178                h: frame.height(),
179            })
180            .unwrap()
181    }
182
183    fn span(text: &str, x: u32, y: u32) -> TextSpan {
184        TextSpan {
185            text: text.to_string(),
186            confidence: 0.9,
187            bbox: PxRect { x, y, w: 10, h: 8 },
188        }
189    }
190
191    #[test]
192    fn text_mode_sorts_spans_and_joins_lines() {
193        use crate::detector::{Detector, DetectorOutput};
194        // Deliberately out of order: second line first, then two spans on line one.
195        let mut det = OcrDetector::text(FixedOcr(vec![
196            span("Sword", 2, 20),
197            span("Doombringer", 2, 2),
198            span("Unique", 30, 2),
199        ]));
200        let frame = view_frame();
201        assert_eq!(
202            det.evaluate(&full_view(&frame)).unwrap(),
203            DetectorOutput::Text("Doombringer\nUnique\nSword".to_string())
204        );
205        assert!(det.is_heavy());
206    }
207
208    #[test]
209    fn text_mode_empty_spans_is_none() {
210        use crate::detector::{Detector, DetectorOutput};
211        let mut det = OcrDetector::text(FixedOcr(vec![]));
212        let frame = view_frame();
213        assert_eq!(
214            det.evaluate(&full_view(&frame)).unwrap(),
215            DetectorOutput::None
216        );
217    }
218
219    #[test]
220    fn spans_mode_preserves_metadata_in_bbox_order() {
221        use crate::detector::{Detector, DetectorOutput};
222        let mut det = OcrDetector::spans(FixedOcr(vec![
223            span("Sword", 2, 20),
224            span("Doombringer", 2, 2),
225        ]));
226        let frame = view_frame();
227        let DetectorOutput::Spans(spans) = det.evaluate(&full_view(&frame)).unwrap() else {
228            panic!("expected Spans output");
229        };
230        // Sorted top-to-bottom like text mode, with geometry/confidence intact.
231        assert_eq!(spans.len(), 2);
232        assert_eq!(spans[0].text, "Doombringer");
233        assert_eq!(spans[0].bbox.y, 2);
234        assert_eq!(spans[1].text, "Sword");
235        assert_eq!(spans[1].bbox.y, 20);
236        assert!(spans.iter().all(|s| s.confidence > 0.0));
237        assert!(det.is_heavy());
238    }
239
240    #[test]
241    fn spans_mode_empty_is_none() {
242        use crate::detector::{Detector, DetectorOutput};
243        let mut det = OcrDetector::spans(FixedOcr(vec![]));
244        let frame = view_frame();
245        assert_eq!(
246            det.evaluate(&full_view(&frame)).unwrap(),
247            DetectorOutput::None
248        );
249    }
250
251    #[test]
252    fn number_mode_parses_and_guards() {
253        use crate::detector::{Detector, DetectorOutput};
254        let mut det = OcrDetector::number(FixedOcr(vec![span("HP 1,234", 0, 0)]), |t| {
255            crate::patterns::number()(t)
256        });
257        let frame = view_frame();
258        assert_eq!(
259            det.evaluate(&full_view(&frame)).unwrap(),
260            DetectorOutput::Number(1234.0)
261        );
262
263        // Parse failure -> None
264        let mut det = OcrDetector::number(FixedOcr(vec![span("no digits", 0, 0)]), |t| {
265            crate::patterns::number()(t)
266        });
267        assert_eq!(
268            det.evaluate(&full_view(&frame)).unwrap(),
269            DetectorOutput::None
270        );
271
272        // A parser returning NaN must be swallowed to None (never-NaN invariant).
273        let mut det = OcrDetector::number(FixedOcr(vec![span("x", 0, 0)]), |_| Some(f64::NAN));
274        assert_eq!(
275            det.evaluate(&full_view(&frame)).unwrap(),
276            DetectorOutput::None
277        );
278    }
279
280    #[test]
281    fn engine_errors_propagate_unchanged() {
282        use crate::detector::Detector;
283
284        struct FailingOcr;
285        impl OcrEngine for FailingOcr {
286            fn recognize(&mut self, _view: &FrameView<'_>) -> Result<Vec<TextSpan>, DetectorError> {
287                Err(DetectorError::Ocr("engine exploded".to_string()))
288            }
289        }
290
291        let mut det = OcrDetector::text(FailingOcr);
292        let frame = view_frame();
293        let err = det.evaluate(&full_view(&frame)).unwrap_err();
294        assert!(matches!(err, DetectorError::Ocr(msg) if msg == "engine exploded"));
295    }
296}
visual_cortex_vision/ocr.rs

visual_cortex_vision/
ocr.rs
