visual_cortex_vision/

ocr.rs

use visual_cortex_capture::{FrameView, PxRect};

use crate::detector::DetectorError;

/// One recognized text region. `bbox` is relative to the evaluated view.
#[derive(Debug, Clone, PartialEq)]
pub struct TextSpan {
    pub text: String,
    /// Mean character confidence in 0.0..=1.0. Never NaN.
    pub confidence: f32,
    pub bbox: PxRect,
}

/// Pluggable OCR backend. The default ONNX/PaddleOCR implementation lives in
/// the `visual-cortex-ocr-onnx` crate; any engine can slot in here. Returning an empty
/// vec means "no text present" and is not an error.
pub trait OcrEngine: Send + 'static {
    fn recognize(&mut self, view: &FrameView<'_>) -> Result<Vec<TextSpan>, DetectorError>;
}

use crate::detector::{Detector, DetectorOutput};

type NumberParser = Box<dyn Fn(&str) -> Option<f64> + Send + 'static>;

enum OcrMode {
    Text,
    Number(NumberParser),
}

/// Adapts any [`OcrEngine`] into a [`Detector`].
///
/// Text mode sorts spans top-to-bottom then left-to-right and joins them with
/// `\n` into `DetectorOutput::Text`; an empty span list becomes
/// `DetectorOutput::None`, so, under `Mode::Change`, "text appeared" is the
/// `None` -> `Text` transition.
/// Number mode additionally parses the joined text; parse failures and
/// non-finite values become `DetectorOutput::None` (the never-NaN invariant).
pub struct OcrDetector {
    engine: Box<dyn OcrEngine>,
    mode: OcrMode,
}

impl OcrDetector {
    pub fn text(engine: impl OcrEngine) -> Self {
        Self {
            engine: Box::new(engine),
            mode: OcrMode::Text,
        }
    }

    pub fn number(
        engine: impl OcrEngine,
        parser: impl Fn(&str) -> Option<f64> + Send + 'static,
    ) -> Self {
        Self {
            engine: Box::new(engine),
            mode: OcrMode::Number(Box::new(parser)),
        }
    }
}

impl Detector for OcrDetector {
    fn evaluate(&mut self, view: &FrameView<'_>) -> Result<DetectorOutput, DetectorError> {
        let mut spans = self.engine.recognize(view)?;
        if spans.is_empty() {
            return Ok(DetectorOutput::None);
        }
        spans.sort_by_key(|s| (s.bbox.y, s.bbox.x));
        let joined = spans
            .iter()
            .map(|s| s.text.as_str())
            .collect::<Vec<_>>()
            .join("\n");
        match &self.mode {
            OcrMode::Text => Ok(DetectorOutput::Text(joined)),
            OcrMode::Number(parser) => Ok(match parser(&joined) {
                Some(n) if n.is_finite() => DetectorOutput::Number(n),
                _ => DetectorOutput::None,
            }),
        }
    }

    fn is_heavy(&self) -> bool {
        true
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use visual_cortex_capture::{Frame, PxRect};

    struct FixedOcr(Vec<TextSpan>);

    impl OcrEngine for FixedOcr {
        fn recognize(&mut self, _view: &FrameView<'_>) -> Result<Vec<TextSpan>, DetectorError> {
            Ok(self.0.clone())
        }
    }

    #[test]
    fn trait_is_object_safe_and_returns_spans() {
        let span = TextSpan {
            text: "HP 300".to_string(),
            confidence: 0.98,
            bbox: PxRect {
                x: 1,
                y: 2,
                w: 40,
                h: 10,
            },
        };
        let mut engine: Box<dyn OcrEngine> = Box::new(FixedOcr(vec![span]));
        let frame = Frame::solid(64, 16, [0, 0, 0, 255]);
        let view = frame
            .view(PxRect {
                x: 0,
                y: 0,
                w: 64,
                h: 16,
            })
            .unwrap();
        let spans = engine.recognize(&view).unwrap();
        assert_eq!(spans.len(), 1);
        assert_eq!(spans[0].text, "HP 300");
        assert!(spans[0].confidence > 0.9);
        assert_eq!(spans[0].bbox.w, 40);
    }

    fn view_frame() -> Frame {
        Frame::solid(64, 32, [0, 0, 0, 255])
    }

    fn full_view(frame: &Frame) -> FrameView<'_> {
        frame
            .view(PxRect {
                x: 0,
                y: 0,
                w: frame.width(),
                h: frame.height(),
            })
            .unwrap()
    }

    fn span(text: &str, x: u32, y: u32) -> TextSpan {
        TextSpan {
            text: text.to_string(),
            confidence: 0.9,
            bbox: PxRect { x, y, w: 10, h: 8 },
        }
    }

    #[test]
    fn text_mode_sorts_spans_and_joins_lines() {
        use crate::detector::{Detector, DetectorOutput};
        // Deliberately out of order: second line first, then two spans on line one.
        let mut det = OcrDetector::text(FixedOcr(vec![
            span("Sword", 2, 20),
            span("Doombringer", 2, 2),
            span("Unique", 30, 2),
        ]));
        let frame = view_frame();
        assert_eq!(
            det.evaluate(&full_view(&frame)).unwrap(),
            DetectorOutput::Text("Doombringer\nUnique\nSword".to_string())
        );
        assert!(det.is_heavy());
    }

    #[test]
    fn text_mode_empty_spans_is_none() {
        use crate::detector::{Detector, DetectorOutput};
        let mut det = OcrDetector::text(FixedOcr(vec![]));
        let frame = view_frame();
        assert_eq!(
            det.evaluate(&full_view(&frame)).unwrap(),
            DetectorOutput::None
        );
    }

    #[test]
    fn number_mode_parses_and_guards() {
        use crate::detector::{Detector, DetectorOutput};
        let mut det = OcrDetector::number(FixedOcr(vec![span("HP 1,234", 0, 0)]), |t| {
            crate::patterns::number()(t)
        });
        let frame = view_frame();
        assert_eq!(
            det.evaluate(&full_view(&frame)).unwrap(),
            DetectorOutput::Number(1234.0)
        );

        // Parse failure -> None
        let mut det = OcrDetector::number(FixedOcr(vec![span("no digits", 0, 0)]), |t| {
            crate::patterns::number()(t)
        });
        assert_eq!(
            det.evaluate(&full_view(&frame)).unwrap(),
            DetectorOutput::None
        );

        // A parser returning NaN must be swallowed to None (never-NaN invariant).
        let mut det = OcrDetector::number(FixedOcr(vec![span("x", 0, 0)]), |_| Some(f64::NAN));
        assert_eq!(
            det.evaluate(&full_view(&frame)).unwrap(),
            DetectorOutput::None
        );
    }

    #[test]
    fn engine_errors_propagate_unchanged() {
        use crate::detector::Detector;

        struct FailingOcr;
        impl OcrEngine for FailingOcr {
            fn recognize(&mut self, _view: &FrameView<'_>) -> Result<Vec<TextSpan>, DetectorError> {
                Err(DetectorError::Ocr("engine exploded".to_string()))
            }
        }

        let mut det = OcrDetector::text(FailingOcr);
        let frame = view_frame();
        let err = det.evaluate(&full_view(&frame)).unwrap_err();
        assert!(matches!(err, DetectorError::Ocr(msg) if msg == "engine exploded"));
    }
}