visual_cortex_ocr_onnx/
engine.rs

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

visual_cortex_ocr_onnx/
engine.rs
