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

visual_cortex_vision/
ocr.rs
