visual_cortex_ocr_onnx/
engine.rs

1use std::path::Path;
2use std::sync::{Arc, Mutex};
3
4use ndarray::Array4;
5use ort::session::Session;
6use ort::value::TensorRef;
7use visual_cortex_capture::{FrameView, PxRect};
8use visual_cortex_vision::{DetectorError, OcrEngine, TextSpan};
9
10use crate::det::{det_target_size, extract_boxes, preprocess_det, view_to_rgb};
11use crate::error::OcrError;
12use crate::models::ensure_all_cached;
13use crate::rec::{build_charset, ctc_decode, preprocess_rec};
14
15/// Binarization threshold for the det probability map. Tuned against the
16/// hp_1234 e2e fixture: lowered from PaddleOCR's stock 0.3 while diagnosing a
17/// clipped leading digit (no box change on its own, kept as the more
18/// permissive tested value; false positives are filtered by MIN_CONFIDENCE).
19const DET_THRESHOLD: f32 = 0.2;
20/// Box padding in det-map pixels (the DB unclip substitute). Tuned 3 -> 5
21/// against the hp_1234 e2e fixture: at 3 the rec crop clipped the leading
22/// "1" and read "HP 234".
23const BOX_PAD: u32 = 5;
24/// Spans below this confidence are discarded.
25const MIN_CONFIDENCE: f32 = 0.3;
26
27/// PaddleOCR det+rec pipeline behind visual-cortex-vision's `OcrEngine`.
28///
29/// Cheap to clone: clones share the loaded ONNX sessions and charset
30/// (issue #15 — N OCR watchers, one model in memory). Inference across
31/// clones is serialized on the shared sessions (ort 2.0.0-rc.12's safe
32/// `Session::run` takes `&mut self`, verified against vendored source);
33/// OCR already runs on the blocking pool, so contending watchers queue
34/// briefly instead of multiplying model memory.
35#[derive(Clone)]
36pub struct PaddleOcr {
37    det: Arc<Mutex<Session>>,
38    rec: Arc<Mutex<Session>>,
39    charset: Arc<Vec<String>>,
40}
41
42impl PaddleOcr {
43    /// Download (if absent), verify, and load the pinned models from the
44    /// platform cache dir (`~/.cache/visual_cortex` or the macOS/Windows equivalent).
45    pub async fn new() -> Result<Self, OcrError> {
46        let (det, rec, dict) = ensure_all_cached().await?;
47        Self::from_paths(&det, &rec, &dict)
48    }
49
50    /// Load models from explicit local paths. Never touches the network.
51    pub fn from_paths(det: &Path, rec: &Path, dict: &Path) -> Result<Self, OcrError> {
52        let det = Session::builder()
53            .and_then(|mut b| b.commit_from_file(det))
54            .map_err(|e| OcrError::ModelLoad(format!("det: {e}")))?;
55        let rec = Session::builder()
56            .and_then(|mut b| b.commit_from_file(rec))
57            .map_err(|e| OcrError::ModelLoad(format!("rec: {e}")))?;
58        let dict =
59            std::fs::read_to_string(dict).map_err(|e| OcrError::ModelLoad(format!("dict: {e}")))?;
60        Ok(Self {
61            det: Arc::new(Mutex::new(det)),
62            rec: Arc::new(Mutex::new(rec)),
63            charset: Arc::new(build_charset(&dict)),
64        })
65    }
66
67    fn run_det(&mut self, input: Array4<f32>, tw: u32, th: u32) -> Result<Vec<f32>, OcrError> {
68        let mut det = self.det.lock().unwrap();
69        let outputs = det
70            .run(ort::inputs!["x" => TensorRef::from_array_view(input.view())
71                .map_err(|e| OcrError::Inference(format!("det input: {e}")))?])
72            .map_err(|e| OcrError::Inference(format!("det run: {e}")))?;
73        let map = outputs["sigmoid_0.tmp_0"]
74            .try_extract_array::<f32>()
75            .map_err(|e| OcrError::Inference(format!("det output: {e}")))?;
76        let flat: Vec<f32> = map.iter().copied().collect();
77        let expected = tw as usize * th as usize;
78        if flat.len() != expected {
79            return Err(OcrError::Inference(format!(
80                "det output has {} elements, expected {tw}x{th} = {expected}",
81                flat.len()
82            )));
83        }
84        Ok(flat)
85    }
86
87    fn run_rec(&mut self, input: Array4<f32>) -> Result<(String, f32), OcrError> {
88        let mut rec = self.rec.lock().unwrap();
89        let outputs = rec
90            .run(ort::inputs!["x" => TensorRef::from_array_view(input.view())
91                .map_err(|e| OcrError::Inference(format!("rec input: {e}")))?])
92            .map_err(|e| OcrError::Inference(format!("rec run: {e}")))?;
93        let probs = outputs["softmax_2.tmp_0"]
94            .try_extract_array::<f32>()
95            .map_err(|e| OcrError::Inference(format!("rec output: {e}")))?;
96        let shape = probs.shape().to_vec(); // [1, T, C]
97        if shape.len() != 3 || shape[2] != self.charset.len() {
98            return Err(OcrError::Inference(format!(
99                "rec output shape {shape:?} does not match charset len {}",
100                self.charset.len()
101            )));
102        }
103        let flat: Vec<f32> = probs.iter().copied().collect();
104        Ok(ctc_decode(&flat, shape[1], shape[2], &self.charset))
105    }
106}
107
108impl OcrEngine for PaddleOcr {
109    fn recognize(&mut self, view: &FrameView<'_>) -> Result<Vec<TextSpan>, DetectorError> {
110        let rgb = view_to_rgb(view);
111        let (vw, vh) = (rgb.width(), rgb.height());
112        let (tw, th) = det_target_size(vw, vh);
113        let det_input = preprocess_det(&rgb, tw, th);
114        let prob = self
115            .run_det(det_input, tw, th)
116            .map_err(|e| DetectorError::Ocr(e.to_string()))?;
117        let boxes = extract_boxes(&prob, tw as usize, th as usize, DET_THRESHOLD, BOX_PAD);
118
119        // Map det-map boxes back to view coordinates.
120        let (sx, sy) = (vw as f32 / tw as f32, vh as f32 / th as f32);
121        let mut spans = Vec::new();
122        for b in boxes {
123            let x = ((b.x as f32 * sx) as u32).min(vw - 1);
124            let y = ((b.y as f32 * sy) as u32).min(vh - 1);
125            let w = ((b.w as f32 * sx).ceil() as u32).clamp(1, vw - x);
126            let h = ((b.h as f32 * sy).ceil() as u32).clamp(1, vh - y);
127            let crop = image::imageops::crop_imm(&rgb, x, y, w, h).to_image();
128            let (text, confidence) = self
129                .run_rec(preprocess_rec(&crop))
130                .map_err(|e| DetectorError::Ocr(e.to_string()))?;
131            if !text.is_empty() && confidence >= MIN_CONFIDENCE {
132                spans.push(TextSpan {
133                    text,
134                    confidence,
135                    bbox: PxRect { x, y, w, h },
136                });
137            }
138        }
139        Ok(spans)
140    }
141}
142
143#[cfg(test)]
144mod tests {
145    #[test]
146    fn paddle_ocr_is_cheaply_cloneable() {
147        fn assert_clone<T: Clone>() {}
148        assert_clone::<super::PaddleOcr>();
149    }
150}
visual_cortex_ocr_onnx/engine.rs

visual_cortex_ocr_onnx/
engine.rs
