docling_pdf/ocr.rs
1//! OCR for scanned pages, via the PP-OCRv3 recognition model (CRNN/SVTR) run
2//! with `ort`. The layout model already locates text regions on the page image
3//! (it works without a text layer), so OCR only needs *recognition*: each text
4//! region is cropped, split into lines by horizontal projection, and each line
5//! is recognised and decoded with CTC — producing [`TextCell`]s the normal
6//! layout assembly then consumes. This avoids a separate text-detection model.
7
8use image::RgbImage;
9use ort::session::Session;
10use ort::value::Tensor;
11
12use crate::layout::Region;
13// The ONNX-free half (line prep, batching, CTC decode) lives in `ocr_prep`
14// so the wasm build shares it verbatim (#79 phase 2).
15use crate::ocr_prep::{
16 batch_input, decode_row, dict_chars, prep_region_lines, width_batches, PrepLine, REC_HEIGHT,
17};
18use crate::pdfium_backend::TextCell;
19
20pub struct OcrModel {
21 rec: Session,
22 /// CTC classes: index 0 = blank, 1..=6623 = dictionary, 6624 = space.
23 chars: Vec<String>,
24}
25
26/// OCR recognition language: which PP-OCRv3 model + dictionary pair runs.
27///
28/// The default is **English** (`models/ocr_rec_en.onnx` + `models/en_dict.txt`):
29/// the multilingual `ch_` model reads Latin scripts with badly degraded word
30/// spacing (glued words on ordinary English scans), which is the common
31/// real-world case. `Ch` selects the `ch_` pair (`models/ocr_rec.onnx` +
32/// `models/ppocr_keys_v1.txt`) — that is what upstream docling conformance is
33/// measured with, and `scripts/conformance/pdf_*.sh` pin it explicitly (by
34/// path, which wins over this selector).
35#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
36pub enum OcrLang {
37 /// en_PP-OCRv3 — English-only, proper Latin word spacing.
38 #[default]
39 En,
40 /// ch_PP-OCRv3 — multilingual; the docling-conformance model.
41 Ch,
42}
43
44impl OcrLang {
45 /// Parse a user-supplied language id. `None` for anything but `en`/`ch`
46 /// (trimmed, case-insensitive) — callers surface their own error/warning.
47 pub fn parse(s: &str) -> Option<Self> {
48 match s.trim().to_ascii_lowercase().as_str() {
49 "en" => Some(Self::En),
50 "ch" => Some(Self::Ch),
51 _ => None,
52 }
53 }
54
55 /// The process-level choice from `DOCLING_RS_OCR_LANG` (empty/unset → the
56 /// English default; unknown values warn and use English).
57 pub fn from_env() -> Self {
58 let raw = std::env::var("DOCLING_RS_OCR_LANG").unwrap_or_default();
59 if raw.trim().is_empty() {
60 return Self::default();
61 }
62 Self::parse(&raw).unwrap_or_else(|| {
63 eprintln!("docling-pdf: DOCLING_RS_OCR_LANG={raw:?} is not en|ch; using en");
64 Self::default()
65 })
66 }
67}
68
69/// Resolve the recognition model + dictionary pair for `lang`. An English
70/// default that isn't on disk (older model checkouts) degrades to the `ch_`
71/// pair with a warning rather than failing — the usual missing-optional-asset
72/// convention. Explicit `DOCLING_OCR_REC_ONNX` / `DOCLING_OCR_DICT` paths win
73/// over all of this; they are a pair, so set both together.
74pub(crate) fn resolve_rec_pair(lang: OcrLang) -> (String, String) {
75 const CH: (&str, &str) = ("models/ocr_rec.onnx", "models/ppocr_keys_v1.txt");
76 const EN: (&str, &str) = ("models/ocr_rec_en.onnx", "models/en_dict.txt");
77 let want_ch = lang == OcrLang::Ch;
78 let pick = if want_ch { CH } else { EN };
79 let (mut rec, mut dict) = (crate::resolve_asset(pick.0), crate::resolve_asset(pick.1));
80 if !want_ch && (!std::path::Path::new(&rec).exists() || !std::path::Path::new(&dict).exists()) {
81 let (ch_rec, ch_dict) = (crate::resolve_asset(CH.0), crate::resolve_asset(CH.1));
82 if std::path::Path::new(&ch_rec).exists() && std::path::Path::new(&ch_dict).exists() {
83 eprintln!(
84 "docling-pdf: English OCR model not found ({rec}); falling back to the \
85 multilingual ch_ model — expect weak Latin word spacing. Fetch it with \
86 scripts/install/download_dependencies.sh"
87 );
88 (rec, dict) = (ch_rec, ch_dict);
89 }
90 }
91 (
92 std::env::var("DOCLING_OCR_REC_ONNX").unwrap_or(rec),
93 std::env::var("DOCLING_OCR_DICT").unwrap_or(dict),
94 )
95}
96
97impl OcrModel {
98 /// Load the recognition model and its character dictionary for `lang` —
99 /// see [`resolve_rec_pair`] for the selection rules (explicit
100 /// `DOCLING_OCR_REC_ONNX`/`DOCLING_OCR_DICT` paths win).
101 pub fn load(lang: OcrLang) -> Result<Self, String> {
102 let (rec_path, dict_path) = resolve_rec_pair(lang);
103 // Single-threaded: ORT's multi-threaded float-reduction order varies
104 // across runs, which flips the CTC argmax on low-confidence characters
105 // (e.g. noisy faxes) and makes the snapshot output non-deterministic. The
106 // recognition inputs are tiny per-line crops, so the throughput cost is
107 // negligible.
108 let builder = Session::builder()
109 .map_err(|e| format!("ocr: builder: {e}"))?
110 .with_intra_threads(1)
111 .map_err(|e| format!("ocr: intra_threads: {e}"))?;
112 let rec = crate::ep::apply(builder)
113 .map_err(|e| format!("ocr: {e}"))?
114 .commit_from_file(&rec_path)
115 .map_err(|e| format!("ocr: load {rec_path}: {e}"))?;
116 let dict = std::fs::read_to_string(&dict_path)
117 .map_err(|e| format!("ocr: read dict {dict_path}: {e}"))?;
118 Ok(Self {
119 rec,
120 chars: dict_chars(&dict),
121 })
122 }
123
124 /// Recognise a batch of prepared *same-width* lines in one session run.
125 ///
126 /// Only equal widths ever share a run: same-width batching is
127 /// bit-identical to one-at-a-time recognition (each sample keeps its own
128 /// data and per-sample kernel reduction order — verified empirically on
129 /// the scanned corpus), whereas width-padding leaks into the real
130 /// timesteps through the model's global-attention blocks and measurably
131 /// changes low-confidence characters.
132 fn recognize_batch(
133 &mut self,
134 w: usize,
135 chunk: &[usize],
136 lines: &[PrepLine],
137 ) -> Result<Vec<String>, String> {
138 let n = chunk.len();
139 let data = batch_input(w, chunk, lines);
140 let input = Tensor::from_array(([n, 3, REC_HEIGHT as usize, w], data))
141 .map_err(|e| format!("ocr: input tensor: {e}"))?;
142 let outputs = self
143 .rec
144 .run(ort::inputs!["x" => input])
145 .map_err(|e| format!("ocr: rec inference: {e}"))?;
146 let (shape, probs) = outputs[0]
147 .try_extract_tensor::<f32>()
148 .map_err(|e| format!("ocr: extract rec: {e}"))?;
149 let t_len = shape[1] as usize;
150 let nc = shape[2] as usize;
151 Ok((0..n)
152 .map(|i| {
153 decode_row(
154 &self.chars,
155 &probs[i * t_len * nc..(i + 1) * t_len * nc],
156 nc,
157 )
158 })
159 .collect())
160 }
161
162 /// OCR a page: produce text cells (page points) for every line found inside
163 /// the text regions. `scale` is image-px per page-point.
164 pub fn ocr_page(
165 &mut self,
166 img: &RgbImage,
167 regions: &[Region],
168 scale: f32,
169 ) -> Result<Vec<TextCell>, String> {
170 // Gather every line crop on the page first (shared with the browser
171 // path), so equal-width lines can share a recognition run regardless
172 // of which region they came from.
173 let (bboxes, lines) = prep_region_lines(img, regions, scale);
174
175 // Deterministic width-batching (shared with the wasm path).
176 let mut texts = vec![String::new(); lines.len()];
177 for (w, chunk) in width_batches(&lines) {
178 for (&i, text) in chunk.iter().zip(self.recognize_batch(w, &chunk, &lines)?) {
179 texts[i] = text;
180 }
181 }
182
183 // Emit cells in page order, exactly as the sequential walk did.
184 let mut cells = Vec::new();
185 for ((l, t, r, b), text) in bboxes.into_iter().zip(texts) {
186 let text = text.trim().to_string();
187 if text.is_empty() {
188 continue;
189 }
190 cells.push(TextCell { text, l, t, r, b });
191 }
192 Ok(cells)
193 }
194}