Skip to main content

oar_ocr/indra/
mod.rs

1//! Indra OCR: Datalab's CPU-first page reader on ONNX Runtime.
2//!
3//! [Indra OCR](https://huggingface.co/datalab-to/indra-ocr) is a set of small
4//! ONNX models that turn a page image into lines with boxes, layout blocks
5//! in reading order, and Markdown: a page orientation classifier, a line
6//! detector, a script classifier that routes each line to the recognizer
7//! of its script group (Latin, Cyrillic and Greek; Chinese, Japanese and
8//! Korean; Arabic and Hebrew; Indic, Southeast Asian and other scripts),
9//! n-gram language models for lines a recognizer is unsure of, a
10//! reading-order model with layout heads, and an optional math reader that
11//! writes lines with math as inline LaTeX.
12//!
13//! [`IndraOcr`] loads the models from a local copy of the Hugging Face
14//! repository: `det.onnx`, `rec.onnx`, `order.onnx` and `meta.json`, with
15//! `orient.onnx`, `script.onnx`, the group recognizers `rec_<group>.onnx`
16//! and the language models `lm.npz` and `lm_<group>.npz` used when present.
17//! It reproduces the reference pipeline's OCR mode (`indra-ocr` 1.1.2;
18//! the earlier v3.1 model set loads too) with two fixes: token positions
19//! of a tall-window recognizer map back through the crop it read, and a
20//! table built from the table heads never spans over a neighbouring cell. Reading a PDF's own text layer is
21//! not supported: render PDF pages to images first.
22//!
23//! [`IndraOcrBuilder::workers`] enables CPU page concurrency with separate
24//! model sessions and shared language models. Select execution providers
25//! through [`IndraOcrBuilder::ort_session`]; accelerators require one worker.
26//!
27//! The model weights use a modified OpenRAIL-M license; see the model
28//! repository before using them commercially.
29//!
30//! ```no_run
31//! use oar_ocr::indra::IndraOcrBuilder;
32//! use oar_ocr::utils::load_image;
33//!
34//! # fn main() -> Result<(), Box<dyn std::error::Error>> {
35//! let ocr = IndraOcrBuilder::new("models/datalab-to/indra-ocr").build()?;
36//! let page = ocr.predict(vec![load_image("page.png")?])?.remove(0);
37//! println!("{}", page.markdown);
38//! for line in &page.lines {
39//!     println!("{:?} {}", line.bbox, line.text);
40//! }
41//! # Ok(())
42//! # }
43//! ```
44
45mod cv;
46mod detect;
47mod layout;
48mod lm;
49mod math;
50mod order;
51mod recognize;
52mod rtl;
53
54use image::{GrayImage, RgbImage};
55use oar_ocr_core::core::OCRError;
56use oar_ocr_core::core::config::{ModelInferenceConfig, OrtSessionConfig};
57use oar_ocr_core::core::inference::OrtInfer;
58use ort::session::{SessionInputValue, SessionOutputs};
59use ort::value::Tensor;
60use rayon::prelude::*;
61use recognize::{Crop, LineText, Token};
62use serde::Serialize;
63use std::borrow::Cow;
64use std::path::{Path, PathBuf};
65use std::sync::Arc;
66use std::sync::atomic::{AtomicUsize, Ordering};
67use unicode_normalization::UnicodeNormalization;
68
69/// A median line height (detector input pixels) under this triggers a
70/// second detection pass at a larger scale.
71const ADAPTIVE_MIN_H: f32 = 9.0;
72const ADAPTIVE_TARGET_H: f32 = 12.0;
73const ADAPTIVE_MAX_SIDE: usize = 2048;
74/// Cut a detected line at a gap this many line heights wide (between CTC
75/// character centers).
76const SPLIT_GAP: f64 = 3.0;
77/// Cut a table line into cells at a gap this many line heights wide.
78const CELL_GAP: f64 = 1.8;
79/// `orient.onnx` input side: the page thumbnail, padded square.
80const ORIENT_SIDE: usize = 768;
81/// A quarter or half turn is applied only above this probability.
82const ORIENT_MIN: f32 = 0.95;
83/// The `orient.onnx` class (of five) of a page set in vertical columns.
84const VERTICAL: u32 = 4;
85/// The skew search and the orientation model read the page shrunk to at
86/// most this many pixels a side.
87const PREVIEW_SIDE: f64 = 1600.0;
88/// A recognized line under this mean CTC peak probability gets no text.
89const MIN_CONF: f64 = 0.5;
90/// Degrees of tilt searched by the deskew (quarter turns cover the rest).
91const DESKEW_MAX: f64 = 45.0;
92/// Smaller tilts are left alone.
93const DESKEW_MIN: f64 = 0.3;
94/// Lines per recognizer batch.
95const REC_BATCH: usize = 64;
96/// A line goes to one group's recognizer when the script classifier is
97/// surer than this; otherwise the main recognizer reads it as well and the
98/// more confident reading wins.
99const ROUTE_SURE: f32 = 0.95;
100/// Script groups in `script.onnx`'s class order unless `meta.json` names
101/// them (`script_groups`). `base` is `rec.onnx`.
102const GROUPS: [&str; 4] = ["base", "cjk", "rtl", "other"];
103/// Language-model defaults (`meta.json` `lm_params` overrides them per
104/// group): lines under the gate confidence are read by beam search.
105const LM_GATE: f64 = 0.9;
106const LM_ALPHA: f64 = 0.5;
107const LM_BETA: f64 = 1.5;
108const LM_BEAM: usize = 8;
109/// A class under this probability at a step is not tried by the beam search.
110const LM_CAND_FLOOR: f64 = 3e-3;
111
112/// Builder for [`IndraOcr`].
113#[derive(Debug, Clone)]
114pub struct IndraOcrBuilder {
115    model_dir: PathBuf,
116    math: bool,
117    keep_furniture: bool,
118    table_heads: bool,
119    workers: usize,
120    ort_session: Option<OrtSessionConfig>,
121}
122
123impl IndraOcrBuilder {
124    /// Starts a builder for the models in `model_dir`, a local copy of the
125    /// `datalab-to/indra-ocr` Hugging Face repository.
126    pub fn new(model_dir: impl Into<PathBuf>) -> Self {
127        Self {
128            model_dir: model_dir.into(),
129            math: false,
130            keep_furniture: false,
131            table_heads: false,
132            workers: 1,
133            ort_session: None,
134        }
135    }
136
137    /// Reads lines that hold math as text with inline LaTeX (off by default;
138    /// slower on math pages). Needs `mathcls.onnx`, `math_enc.onnx`,
139    /// `math_dec_init.onnx`, `math_dec_step.onnx` and `math_meta.json`.
140    pub fn math(mut self, enable: bool) -> Self {
141        self.math = enable;
142        self
143    }
144
145    /// Keeps page headers and footers in the Markdown (left out by default;
146    /// they are always listed in [`IndraPage::blocks`]).
147    pub fn keep_furniture(mut self, enable: bool) -> Self {
148        self.keep_furniture = enable;
149        self
150    }
151
152    /// Builds tables from the order model's row, column and cell heads
153    /// instead of from line geometry (off by default, as in the reference).
154    /// Table lines are first split into cells and the order model runs
155    /// again; a table whose cell head joins side-by-side lines falls back
156    /// to geometry. Needs an `order.onnx` with the `tab` output; without
157    /// it, tables come from geometry.
158    pub fn table_heads(mut self, enable: bool) -> Self {
159        self.table_heads = enable;
160        self
161    }
162
163    /// Sets the ONNX Runtime session configuration (execution providers,
164    /// threads) applied to every model.
165    pub fn ort_session(mut self, config: OrtSessionConfig) -> Self {
166        self.ort_session = Some(config);
167        self
168    }
169
170    /// Sets the number of CPU page workers (default: one).
171    ///
172    /// Each worker owns a copy of every ONNX session; vocabularies and
173    /// language models are shared. More workers increase memory use for
174    /// sessions and page buffers. With multiple workers, unset ONNX intra-
175    /// and inter-op thread counts default to one. Explicit settings are
176    /// retained: keep workers times threads within the available CPU cores.
177    /// A process-wide ORT thread pool retains its own thread configuration.
178    /// Accelerator execution providers currently require one worker.
179    /// Zero workers is a configuration error.
180    pub fn workers(mut self, workers: usize) -> Self {
181        self.workers = workers;
182        self
183    }
184
185    /// Loads the models.
186    pub fn build(self) -> Result<IndraOcr, OCRError> {
187        if self.workers == 0 {
188            return Err(OCRError::config_error(
189                "Indra OCR needs at least one worker",
190            ));
191        }
192        let dir = self.model_dir.as_path();
193        let mut common = ModelInferenceConfig::new();
194        common.ort_session = self.ort_session.map(OrtSessionConfig::resolve_auto);
195        let pool = if self.workers > 1 {
196            let config = common.ort_session.get_or_insert_with(OrtSessionConfig::new);
197            if config.has_accelerator_provider() {
198                return Err(OCRError::config_error(
199                    "Indra OCR accelerator execution requires workers(1); multiple workers are CPU-only",
200                ));
201            }
202            config.intra_threads.get_or_insert(1);
203            config.inter_threads.get_or_insert(1);
204            Some(
205                rayon::ThreadPoolBuilder::new()
206                    .num_threads(self.workers)
207                    .thread_name(|i| format!("indra-{i}"))
208                    .build()
209                    .map_err(|e| OCRError::config_error(format!("Indra OCR worker pool: {e}")))?,
210            )
211        } else {
212            None
213        };
214        let meta = read_json(&dir.join("meta.json"))?;
215        let vocab = |key: &str| -> Result<Vec<String>, OCRError> {
216            let inner: serde_json::Value = serde_json::from_str(meta[key].as_str().unwrap_or(""))
217                .map_err(|e| {
218                OCRError::config_error(format!("meta.json: {key} is not a vocabulary: {e}"))
219            })?;
220            string_list(&inner["chars"], &format!("meta.json: {key}.chars"))
221        };
222        let size = |key: &str, default: Option<usize>| {
223            match meta.get(key) {
224                None => default,
225                Some(v) => v.as_u64().filter(|&v| v > 0).map(|v| v as usize),
226            }
227            .ok_or_else(|| {
228                OCRError::config_error(format!("meta.json: {key} must be a positive integer"))
229            })
230        };
231        // A recognizer, with its language model when the folder has one and
232        // the recognizer exports the beam search's inputs.
233        let recognizer = |name: &str, file: &str, lm_file: &str, vocab: Vec<String>| {
234            let model = model(dir, file, &common)?;
235            let lm_path = dir.join(lm_file);
236            let outputs: Vec<String> = model.output_shapes().into_iter().map(|o| o.0).collect();
237            let lm = if lm_path.exists()
238                && ["blank", "topv", "topi"]
239                    .iter()
240                    .all(|o| outputs.iter().any(|n| n == o))
241            {
242                let params = &meta["lm_params"][name];
243                let param = |key: &str, default: f64| params[key].as_f64().unwrap_or(default);
244                Some(Arc::new(Lm {
245                    model: lm::NgramLm::load(&lm_path)?,
246                    gate: param("gate", LM_GATE),
247                    weights: lm::Weights {
248                        alpha: param("alpha", LM_ALPHA),
249                        beta: param("beta", LM_BETA),
250                        beam: LM_BEAM,
251                        cand_floor: LM_CAND_FLOOR,
252                    },
253                }))
254            } else {
255                None
256            };
257            let jamo = vocab.iter().any(|c| c == "\u{1100}");
258            Ok::<_, OCRError>(Recognizer {
259                model,
260                vocab: vocab.into(),
261                jamo,
262                lm,
263            })
264        };
265        // A script group is read by its own recognizer when the folder has
266        // it; otherwise by rec.onnx.
267        let mut groups = Vec::new();
268        if let Some(map) = meta.as_object() {
269            for name in map.keys().filter_map(|k| k.strip_prefix("vocab_")) {
270                let file = format!("rec_{name}.onnx");
271                if dir.join(&file).exists() {
272                    let rec = recognizer(
273                        name,
274                        &file,
275                        &format!("lm_{name}.npz"),
276                        vocab(&format!("vocab_{name}"))?,
277                    )?;
278                    groups.push((name.to_string(), rec));
279                }
280            }
281        }
282        let script = if !groups.is_empty() && dir.join("script.onnx").exists() {
283            Some(model(dir, "script.onnx", &common)?)
284        } else {
285            None
286        };
287        let script_groups = match &meta["script_groups"] {
288            serde_json::Value::Null => GROUPS.map(str::to_string).to_vec(),
289            v => string_list(v, "meta.json: script_groups")?,
290        };
291        let optional = |name: &str| -> Result<Option<OrtInfer>, OCRError> {
292            if dir.join(name).exists() {
293                model(dir, name, &common).map(Some)
294            } else {
295                Ok(None)
296            }
297        };
298        let first = Engine {
299            det: model(dir, "det.onnx", &common)?,
300            base: recognizer("base", "rec.onnx", "lm.npz", vocab("vocab")?)?,
301            order: model(dir, "order.onnx", &common)?,
302            orient: optional("orient.onnx")?,
303            groups,
304            script,
305            script_groups,
306            math: if self.math {
307                Some(math::MathReader::load(dir, &common)?)
308            } else {
309                None
310            },
311            long_side: size("det_long_side", None)?,
312            rec_height: size("rec_height", None)?,
313            window: size("rec_window", Some(1))?,
314            keep_furniture: self.keep_furniture,
315            table_heads: self.table_heads,
316        };
317        let mut engines = Vec::with_capacity(self.workers);
318        engines.push(first);
319        for _ in 1..self.workers {
320            engines.push(engines[0].new_worker(&common)?);
321        }
322        Ok(IndraOcr { engines, pool })
323    }
324}
325
326/// A CTC line recognizer and its vocabulary.
327#[derive(Debug)]
328struct Recognizer {
329    model: OrtInfer,
330    vocab: Arc<[String]>,
331    /// The vocabulary reads Hangul as jamo, composed back into syllables.
332    jamo: bool,
333    lm: Option<Arc<Lm>>,
334}
335
336impl Recognizer {
337    fn new_worker(&self, common: &ModelInferenceConfig) -> Result<Self, OCRError> {
338        Ok(Self {
339            model: reload_model(&self.model, common)?,
340            vocab: Arc::clone(&self.vocab),
341            jamo: self.jamo,
342            lm: self.lm.clone(),
343        })
344    }
345}
346
347/// A recognizer's language model and when to use it.
348#[derive(Debug)]
349struct Lm {
350    model: lm::NgramLm,
351    /// Lines under this confidence are read by beam search.
352    gate: f64,
353    weights: lm::Weights,
354}
355
356/// Compact greedy output and, only below the LM gate, retained beam inputs.
357struct LineRead {
358    ids: Vec<i64>,
359    steps: Vec<usize>,
360    confidence: f64,
361    lm: Option<LmInput>,
362}
363
364struct LmInput {
365    blank: Vec<f32>,
366    topv: Vec<f32>,
367    topi: Vec<i64>,
368    k: usize,
369}
370
371/// The Indra OCR pipeline. Build it with [`IndraOcrBuilder`].
372#[derive(Debug)]
373pub struct IndraOcr {
374    engines: Vec<Engine>,
375    pool: Option<rayon::ThreadPool>,
376}
377
378/// One worker's ONNX sessions and shared recognition data.
379#[derive(Debug)]
380struct Engine {
381    det: OrtInfer,
382    /// `rec.onnx`: Latin, Cyrillic and Greek, and every line no group
383    /// recognizer reads.
384    base: Recognizer,
385    order: OrtInfer,
386    orient: Option<OrtInfer>,
387    /// The script groups' recognizers.
388    groups: Vec<(String, Recognizer)>,
389    script: Option<OrtInfer>,
390    script_groups: Vec<String>,
391    math: Option<math::MathReader>,
392    long_side: usize,
393    rec_height: usize,
394    /// `rec.onnx` reads this many line bands: the line and the page above
395    /// and below it.
396    window: usize,
397    keep_furniture: bool,
398    table_heads: bool,
399}
400
401/// One recognized line.
402#[derive(Debug, Clone, Serialize)]
403#[non_exhaustive]
404pub struct IndraLine {
405    /// `[x0, y0, x1, y1]` in pixels of the upright page.
406    pub bbox: [f32; 4],
407    /// The line's text; empty when the recognizer was not confident.
408    pub text: String,
409    /// Mean CTC peak probability (1.0 for a math reading).
410    pub confidence: f32,
411}
412
413/// One layout block.
414#[derive(Debug, Clone, Serialize)]
415#[non_exhaustive]
416pub struct IndraBlock {
417    /// The block class: `Text`, `Section-Header`, `Page-Header`,
418    /// `Page-Footer`, `List-Group`, `Table`, `Caption`, `Footnote`,
419    /// `Table-Of-Contents` or `Form`.
420    pub label: String,
421    /// `[x0, y0, x1, y1]` in pixels of the upright page: the union of the
422    /// block's lines, or on a straightened page that union's rectangle
423    /// mapped back across the skew (as the reference does), which can be
424    /// slightly larger than the union of the mapped line boxes.
425    pub bbox: [f32; 4],
426    /// The block's lines, newline-joined.
427    pub text: String,
428    /// The table HTML of a table block.
429    #[serde(skip)]
430    html: Option<String>,
431}
432
433/// One page read by [`IndraOcr`].
434#[derive(Debug, Clone, Serialize)]
435#[non_exhaustive]
436pub struct IndraPage {
437    /// Position of the page in the [`IndraOcr::predict`] batch.
438    pub index: usize,
439    /// Width of the upright page the boxes are in.
440    pub width: u32,
441    /// Height of the upright page the boxes are in.
442    pub height: u32,
443    /// Degrees counter-clockwise the input was turned (0, 90, 180 or 270);
444    /// the page was turned upright before reading.
445    pub rotation: u32,
446    /// Degrees of tilt straightened before reading; boxes are mapped back
447    /// to the page before the straightening.
448    pub skew: f32,
449    /// The page is set in vertical columns (Chinese, Japanese): it was read
450    /// turned a quarter turn, each column a line, and its boxes are on the
451    /// page as it is ([`rotation`](Self::rotation) stays 0).
452    pub vertical: bool,
453    /// Lines in reading order.
454    pub lines: Vec<IndraLine>,
455    /// Layout blocks in reading order (empty for a page with fewer than two
456    /// lines).
457    pub blocks: Vec<IndraBlock>,
458    /// Headings, paragraphs, lists, and tables as HTML. Page headers and
459    /// footers are left out unless [`IndraOcrBuilder::keep_furniture`].
460    pub markdown: String,
461}
462
463impl IndraPage {
464    /// Exports this page in the shared page JSON format
465    /// (`docs/page-format.md`) with `pipeline` `"indra"`.
466    ///
467    /// Block labels map onto the canonical vocabulary (`Section-Header` →
468    /// `paragraph_title`, `Page-Header` → `header`, `Page-Footer` →
469    /// `footer`, `List-Group` → `list`, `Caption` →
470    /// `figure_table_chart_title`, `Table-Of-Contents` → `content`), with
471    /// the original label in `label`; `Form` passes through. Table blocks
472    /// carry their HTML. Boxes are in the upright page, whose size is
473    /// `page.width` and `page.height`, and `page.angle` is
474    /// [`rotation`](Self::rotation) when the page was turned.
475    pub fn to_json(&self) -> serde_json::Value {
476        let blocks: Vec<_> = self
477            .blocks
478            .iter()
479            .enumerate()
480            .map(|(i, block)| {
481                let canonical = match block.label.as_str() {
482                    "Text" => "text",
483                    "Section-Header" => "paragraph_title",
484                    "Page-Header" => "header",
485                    "Page-Footer" => "footer",
486                    "List-Group" => "list",
487                    "Table" => "table",
488                    "Caption" => "figure_table_chart_title",
489                    "Footnote" => "footnote",
490                    "Table-Of-Contents" => "content",
491                    other => other,
492                };
493                let mut value = serde_json::json!({
494                    "type": canonical,
495                    // Boxes are rounded to 0.1 px; widen without float noise.
496                    "bbox": block.bbox.map(|v| (v as f64 * 10.0).round() / 10.0),
497                    "order": i + 1,
498                    "confidence": null,
499                    "angle": null,
500                    "content": block.html.as_deref().unwrap_or(&block.text),
501                });
502                if !block.label.eq_ignore_ascii_case(canonical) {
503                    value["label"] = serde_json::Value::String(block.label.clone());
504                }
505                value
506            })
507            .collect();
508        serde_json::json!({
509            "version": 1,
510            "pipeline": "indra",
511            "page": {
512                "index": self.index,
513                "width": self.width,
514                "height": self.height,
515                "angle": (self.rotation != 0).then_some(self.rotation),
516            },
517            "blocks": blocks,
518            "markdown": self.markdown,
519            "raw_output": null,
520            "diagnostics": [],
521        })
522    }
523}
524
525impl IndraOcr {
526    /// Reads pages; one [`IndraPage`] per image, in input order. An empty
527    /// image is an [`OCRError::InvalidInput`].
528    /// With multiple workers, pages are assigned as engines become free;
529    /// results (or the first page error) are returned in input order.
530    pub fn predict(&self, images: Vec<RgbImage>) -> Result<Vec<IndraPage>, OCRError> {
531        if let Some(i) = images
532            .iter()
533            .position(|im| im.width() == 0 || im.height() == 0)
534        {
535            return Err(OCRError::invalid_input(format!(
536                "image {i} has no pixels ({}x{})",
537                images[i].width(),
538                images[i].height()
539            )));
540        }
541        let read = |engine: &Engine, index: usize| {
542            let mut page = engine.read_page(cv::rgb_to_gray(&images[index]))?;
543            page.index = index;
544            Ok(page)
545        };
546        let Some(pool) = &self.pool else {
547            return (0..images.len())
548                .map(|index| read(&self.engines[0], index))
549                .collect();
550        };
551        if images.is_empty() {
552            return Ok(Vec::new());
553        }
554        let next = AtomicUsize::new(self.engines.len());
555        // One loop per engine. Nested Rayon image processing stays in this
556        // pool, and never creates an additional global preprocessing pool.
557        let mut results: Vec<_> = pool
558            .broadcast(|worker| {
559                let engine = &self.engines[worker.index()];
560                let mut results = Vec::new();
561                // Give each engine its first page before taking more work.
562                // A batch of at least N pages therefore warms every engine.
563                let mut index = worker.index();
564                while index < images.len() {
565                    results.push((index, read(engine, index)));
566                    index = next.fetch_add(1, Ordering::Relaxed);
567                }
568                results
569            })
570            .into_iter()
571            .flatten()
572            .collect();
573        results.sort_unstable_by_key(|(index, _)| *index);
574        results.into_iter().map(|(_, result)| result).collect()
575    }
576}
577
578impl Engine {
579    fn new_worker(&self, common: &ModelInferenceConfig) -> Result<Self, OCRError> {
580        let optional =
581            |model: &Option<OrtInfer>| model.as_ref().map(|m| reload_model(m, common)).transpose();
582        Ok(Self {
583            det: reload_model(&self.det, common)?,
584            base: self.base.new_worker(common)?,
585            order: reload_model(&self.order, common)?,
586            orient: optional(&self.orient)?,
587            groups: self
588                .groups
589                .iter()
590                .map(|(name, rec)| Ok((name.clone(), rec.new_worker(common)?)))
591                .collect::<Result<_, OCRError>>()?,
592            script: optional(&self.script)?,
593            script_groups: self.script_groups.clone(),
594            math: self
595                .math
596                .as_ref()
597                .map(|m| m.new_worker(common))
598                .transpose()?,
599            long_side: self.long_side,
600            rec_height: self.rec_height,
601            window: self.window,
602            keep_furniture: self.keep_furniture,
603            table_heads: self.table_heads,
604        })
605    }
606
607    fn read_page(&self, gray: GrayImage) -> Result<IndraPage, OCRError> {
608        let small = preview(&gray);
609        let a = estimate_skew(&small);
610        // The orientation model reads the page straightened by its row
611        // angle (a tilted page could read as turned).
612        let k = if a.abs() >= DESKEW_MIN {
613            self.orientation(&cv::rotate_u8(&small, a))?
614        } else {
615            self.orientation(&small)?
616        };
617        let vertical = k == VERTICAL;
618        // Quarter turns clockwise to read the page in: a vertical page
619        // turns a quarter counter-clockwise, so its columns become lines,
620        // the first (right) column on top.
621        let turns = if vertical { 3 } else { k % 4 };
622        let gray = cv::rot90_cw(&gray, turns);
623        // A half turn keeps the rows, so the angle stands; after a quarter
624        // turn the rows were columns, so it is measured again.
625        let mut skew = if turns % 2 == 1 {
626            estimate_skew(&cv::rot90_cw(&small, turns))
627        } else {
628            a
629        };
630        let gray = if skew.abs() >= DESKEW_MIN {
631            cv::rotate_u8(&gray, skew)
632        } else {
633            skew = 0.0;
634            gray
635        };
636        let boxes = self.detect(&gray)?;
637        let lines = self.recognize(&gray, &boxes, vertical)?;
638        let mut page = self.assemble(gray.dimensions(), boxes, lines, vertical)?;
639        page.rotation = if vertical { 0 } else { 90 * turns };
640        page.skew = ((skew * 100.0).round_ties_even() / 100.0) as f32;
641        if skew != 0.0 {
642            let wh = gray.dimensions();
643            for line in &mut page.lines {
644                line.bbox = unrotate_box(line.bbox, skew, wh);
645            }
646            for block in &mut page.blocks {
647                block.bbox = unrotate_box(block.bbox, skew, wh);
648            }
649        }
650        if vertical {
651            // Back on the page as it is: (x, y) -> (W - y, x).
652            let w = gray.height() as f64;
653            let back = |b: [f32; 4]| {
654                [
655                    (w - b[3] as f64) as f32,
656                    b[0],
657                    (w - b[1] as f64) as f32,
658                    b[2],
659                ]
660            };
661            for line in &mut page.lines {
662                line.bbox = back(line.bbox);
663            }
664            for block in &mut page.blocks {
665                block.bbox = back(block.bbox);
666            }
667            (page.width, page.height) = (gray.height(), gray.width());
668            page.vertical = true;
669        }
670        Ok(page)
671    }
672
673    /// Quarter turns (counter-clockwise) the page content is turned by, or
674    /// [`VERTICAL`] for a page set in vertical columns.
675    fn orientation(&self, gray: &GrayImage) -> Result<u32, OCRError> {
676        let Some(orient) = &self.orient else {
677            return Ok(0);
678        };
679        let (w, h) = (gray.width() as f64, gray.height() as f64);
680        let s = ORIENT_SIDE as f64 / w.max(h);
681        let (sw, sh) = (
682            ((w * s).round_ties_even() as u32).max(1),
683            ((h * s).round_ties_even() as u32).max(1),
684        );
685        let small = cv::resize(gray, sw, sh, cv::Interp::Area);
686        let mut x = vec![0f32; ORIENT_SIDE * ORIENT_SIDE];
687        for (y, row) in small.as_raw().chunks_exact(sw as usize).enumerate() {
688            for (o, &v) in x[y * ORIENT_SIDE..].iter_mut().zip(row) {
689                *o = 1.0 - v as f32 / 255.0;
690            }
691        }
692        let p = run(
693            orient,
694            ort::inputs!["x" => Tensor::from_array(([1, 1, ORIENT_SIDE, ORIENT_SIDE], x)).map_err(input_error)?],
695            |o| Ok(o["p"].try_extract_tensor::<f32>()?.1.to_vec()),
696        )?;
697        let k = argmax(&p);
698        Ok(if p[k] > ORIENT_MIN { k as u32 } else { 0 })
699    }
700
701    fn detect_at(
702        &self,
703        gray: &GrayImage,
704        long_side: usize,
705    ) -> Result<(Vec<[f32; 5]>, f64), OCRError> {
706        let (x, h, w, scale) = detect::page_to_input(gray, long_side);
707        let boxes = run(
708            &self.det,
709            ort::inputs!["x" => Tensor::from_array(([1, 1, h, w], x)).map_err(input_error)?],
710            |o| {
711                let (shape, out) = o["out"].try_extract_tensor::<f32>()?;
712                Ok(detect::decode(out, shape[2] as usize, shape[3] as usize))
713            },
714        )?;
715        Ok((boxes, scale))
716    }
717
718    /// Line boxes `[x0, y0, x1, y1, score]` in page pixels.
719    fn detect(&self, gray: &GrayImage) -> Result<Vec<[f32; 5]>, OCRError> {
720        let (mut boxes, mut scale) = self.detect_at(gray, self.long_side)?;
721        if boxes.len() >= 8 {
722            let mut heights: Vec<f32> = boxes.iter().map(|b| b[3] - b[1]).collect();
723            let med_h = order::median_f32(&mut heights);
724            if med_h < ADAPTIVE_MIN_H {
725                let target = (self.long_side as f32 * ADAPTIVE_TARGET_H / med_h.max(1.0)) as f64;
726                let ls = (target.min(ADAPTIVE_MAX_SIDE as f64) as usize) / 32 * 32;
727                if ls > self.long_side {
728                    (boxes, scale) = self.detect_at(gray, ls)?;
729                }
730            }
731        }
732        let (w, h) = (gray.width() as f32, gray.height() as f32);
733        let s = scale as f32;
734        for b in &mut boxes {
735            b[0] = (b[0] / s).clamp(0.0, w);
736            b[1] = (b[1] / s).clamp(0.0, h);
737            b[2] = (b[2] / s).clamp(0.0, w);
738            b[3] = (b[3] / s).clamp(0.0, h);
739        }
740        Ok(boxes)
741    }
742
743    /// Text, confidence and tokens per line; each token's x is its CTC peak
744    /// mapped back through the crop. Each line goes to the recognizer of
745    /// its script group (every line of a vertical page to the CJK one); a
746    /// line the script classifier is unsure of is also read by `rec.onnx`,
747    /// and the more confident reading wins. A line under `MIN_CONF` comes
748    /// back empty.
749    fn recognize(
750        &self,
751        gray: &GrayImage,
752        boxes: &[[f32; 5]],
753        vertical: bool,
754    ) -> Result<Vec<LineText>, OCRError> {
755        let height = self.rec_height;
756        let b4 = |b: &[f32; 5]| [b[0], b[1], b[2], b[3]];
757        // Light-on-dark is decided once per line box, for every crop.
758        let dark: Vec<bool> = boxes
759            .par_iter()
760            .map(|b| recognize::is_dark_line(gray, b4(b)))
761            .collect();
762        let crop_all = |window: usize| -> Vec<Crop> {
763            boxes
764                .par_iter()
765                .zip(&dark)
766                .map(|(b, &d)| recognize::crop_line(gray, b4(b), height, window, d))
767                .collect()
768        };
769        let crops = crop_all(1);
770        // rec.onnx may read a tall window: the line and the page around it.
771        let tall = (self.window > 1).then(|| crop_all(self.window));
772        let n = crops.len();
773        // (recognizer, the lines it reads, whether its tokens are characters)
774        let mut reads: Vec<(&Recognizer, Vec<usize>, bool)> =
775            vec![(&self.base, (0..n).collect(), false)];
776        let cjk = self.groups.iter().find(|(name, _)| name == "cjk");
777        if vertical && let Some((_, rec)) = cjk {
778            // A vertical page's lines are CJK columns, characters on their
779            // side; the CJK recognizer reads those.
780            reads = vec![(rec, (0..n).collect(), true)];
781        } else if let Some(script) = &self.script
782            && n > 0
783        {
784            let (route, sure) = self.line_groups(script, &crops)?;
785            reads = vec![(
786                &self.base,
787                (0..n).filter(|&j| route[j].is_none() || !sure[j]).collect(),
788                false,
789            )];
790            for (g, (_, rec)) in self.groups.iter().enumerate() {
791                reads.push((rec, (0..n).filter(|&j| route[j] == Some(g)).collect(), true));
792            }
793        }
794        let mut out = vec![LineText::default(); n];
795        let mut best_conf = vec![-1.0; n];
796        for (rec, sel, char_tokens) in reads {
797            let base = std::ptr::eq(rec, &self.base);
798            let (src, rows) = match &tall {
799                Some(tall) if base => (tall, height * self.window),
800                _ => (&crops, height),
801            };
802            let widths: Vec<usize> = sel.iter().map(|&j| crops[j].width).collect();
803            for part in recognize::width_batches(&widths, REC_BATCH, 1.3, 64 * 640) {
804                let idx: Vec<usize> = part.iter().map(|&i| sel[i]).collect();
805                let batch: Vec<&Crop> = idx.iter().map(|&j| &src[j]).collect();
806                let (data, wm, crop_widths) = recognize::pad_batch(&batch, rows);
807                let b = batch.len();
808                let read = run(
809                    &rec.model,
810                    ort::inputs![
811                        "x" => Tensor::from_array(([b, 1, rows, wm], data)).map_err(input_error)?,
812                        "widths" => Tensor::from_array(([b], crop_widths)).map_err(input_error)?,
813                    ],
814                    |o| {
815                        let (shape, ids) = o["ids"].try_extract_tensor::<i64>()?;
816                        let (_, conf) = o["conf"].try_extract_tensor::<f32>()?;
817                        let (_, lengths) = o["lengths"].try_extract_tensor::<i64>()?;
818                        // Borrow all outputs while reducing greedy CTC. Only rows
819                        // inside the unchanged confidence gate retain LM arrays.
820                        let lm_outputs = if rec.lm.is_some() {
821                            let (_, blank) = o["blank"].try_extract_tensor::<f32>()?;
822                            let (top_shape, topv) = o["topv"].try_extract_tensor::<f32>()?;
823                            let (_, topi) = o["topi"].try_extract_tensor::<i64>()?;
824                            Some((blank, topv, topi, top_shape[2] as usize))
825                        } else {
826                            None
827                        };
828                        let t_len = shape[1] as usize;
829                        let reads = (0..b)
830                            .map(|k| {
831                                let len = (lengths[k].max(0) as usize).min(t_len);
832                                let row = k * t_len;
833                                let seq = &ids[row..row + len];
834                                let pm = &conf[row..row + len];
835                                let keep: Vec<usize> = (0..len)
836                                    .filter(|&t| seq[t] != 0 && (t == 0 || seq[t] != seq[t - 1]))
837                                    .collect();
838                                let confidence = if keep.is_empty() {
839                                    0.0
840                                } else {
841                                    mean_f32(&keep.iter().map(|&t| pm[t]).collect::<Vec<_>>())
842                                };
843                                let lm = match (&rec.lm, lm_outputs) {
844                                    (Some(lm), Some((blank, topv, topi, k)))
845                                        if (MIN_CONF..lm.gate).contains(&confidence) =>
846                                    {
847                                        Some(LmInput {
848                                            blank: blank[row..row + len].to_vec(),
849                                            topv: topv[row * k..(row + len) * k].to_vec(),
850                                            topi: topi[row * k..(row + len) * k].to_vec(),
851                                            k,
852                                        })
853                                    }
854                                    _ => None,
855                                };
856                                LineRead {
857                                    ids: keep.iter().map(|&t| seq[t]).collect(),
858                                    steps: keep,
859                                    confidence,
860                                    lm,
861                                }
862                            })
863                            .collect::<Vec<_>>();
864                        Ok((t_len, reads))
865                    },
866                )?;
867                let (t_len, reads) = read;
868                let stride = wm as f64 / t_len.max(1) as f64;
869                for (read, &j) in reads.into_iter().zip(&idx) {
870                    let LineRead {
871                        ids: mut line_ids,
872                        steps: mut keep,
873                        confidence: line_conf,
874                        lm: lm_input,
875                    } = read;
876                    let [x0, y0, x1, _] =
877                        [boxes[j][0], boxes[j][1], boxes[j][2], boxes[j][3]].map(f64::from);
878                    let h = (boxes[j][3] as f64 - y0).max(1.0);
879                    let ix0 = (x0 - recognize::PAD_X * h).floor();
880                    let cw = (x1 + recognize::PAD_X * h).ceil() - ix0;
881                    // Timesteps map back through the crop the recognizer
882                    // read (the reference scales by the one-band crop even
883                    // for a tall window, which stretches token positions).
884                    let f = cw / src[j].width as f64;
885                    // Expensive beam search stays outside the session lock.
886                    if let (Some(lm), Some(input)) = (&rec.lm, lm_input) {
887                        let steps = lm::Steps {
888                            blank: &input.blank,
889                            topv: &input.topv,
890                            topi: &input.topi,
891                            k: input.k,
892                        };
893                        (line_ids, keep) = lm::beam_search(&steps, &lm.model, lm.weights);
894                    }
895                    if line_conf <= best_conf[j] {
896                        continue;
897                    }
898                    best_conf[j] = line_conf;
899                    // An id past the vocabulary means meta.json does not
900                    // belong to this recognizer.
901                    let tokens: Vec<Token> = line_ids
902                        .iter()
903                        .zip(&keep)
904                        .map(|(&id, &t)| {
905                            let text = usize::try_from(id - 1)
906                                .ok()
907                                .and_then(|i| rec.vocab.get(i))
908                                .ok_or_else(|| {
909                                    OCRError::config_error(format!(
910                                        "{}: token id {id} is outside its {}-entry vocabulary",
911                                        rec.model.model_name(),
912                                        rec.vocab.len()
913                                    ))
914                                })?;
915                            Ok(Token {
916                                text: text.clone(),
917                                x: ix0 + (t as f64 + 0.5) * stride * f,
918                            })
919                        })
920                        .collect::<Result<_, OCRError>>()?;
921                    out[j] = if line_conf >= MIN_CONF {
922                        let text = if char_tokens {
923                            let text: String = tokens.iter().map(|t| t.text.as_str()).collect();
924                            if rec.jamo { text.nfc().collect() } else { text }
925                        } else {
926                            recognize::join_tokens(
927                                &tokens.iter().map(|t| t.text.as_str()).collect::<Vec<_>>(),
928                            )
929                        };
930                        LineText {
931                            text,
932                            conf: line_conf,
933                            tokens,
934                            char_tokens,
935                        }
936                    } else {
937                        LineText {
938                            conf: line_conf,
939                            char_tokens,
940                            ..LineText::default()
941                        }
942                    };
943                }
944            }
945        }
946        if let Some(math) = &self.math
947            && !crops.is_empty()
948        {
949            let flagged: Vec<usize> = math
950                .flags(&crops, height)?
951                .iter()
952                .enumerate()
953                .filter_map(|(j, &f)| f.then_some(j))
954                .collect();
955            if !flagged.is_empty() {
956                let mboxes: Vec<[f32; 4]> = flagged.iter().map(|&j| b4(&boxes[j])).collect();
957                for (&j, (text, conf)) in flagged.iter().zip(math.read(gray, &mboxes)?) {
958                    // Only a confident reading with math replaces the line,
959                    // as one token that is never split.
960                    if text.contains('$') && conf >= math::MATH_MIN_CONF {
961                        let x = (boxes[j][0] + boxes[j][2]) as f64 / 2.0;
962                        out[j] = LineText {
963                            text: text.clone(),
964                            conf: 1.0,
965                            tokens: vec![Token { text, x }],
966                            char_tokens: true,
967                        };
968                    }
969                }
970            }
971        }
972        Ok(out)
973    }
974
975    /// Per line crop, the script group whose recognizer reads it (`None`:
976    /// `rec.onnx`), and whether the script classifier is sure of it. An
977    /// older `script.onnx` has one output: the chance of Chinese characters.
978    fn line_groups(
979        &self,
980        script: &OrtInfer,
981        crops: &[Crop],
982    ) -> Result<(Vec<Option<usize>>, Vec<bool>), OCRError> {
983        let height = self.rec_height;
984        let mut route = vec![None; crops.len()];
985        let mut sure = vec![true; crops.len()];
986        for idx in recognize::line_batches(crops) {
987            let batch: Vec<&Crop> = idx.iter().map(|&j| &crops[j]).collect();
988            let (data, wm, widths) = recognize::pad_batch(&batch, height);
989            let n = batch.len();
990            let (p, classes) = run(
991                script,
992                ort::inputs![
993                    "x" => Tensor::from_array(([n, 1, height, wm], data)).map_err(input_error)?,
994                    "widths" => Tensor::from_array(([n], widths)).map_err(input_error)?,
995                ],
996                |o| {
997                    let (shape, p) = o["p"].try_extract_tensor::<f32>()?;
998                    Ok((p.to_vec(), shape.get(1).map(|&c| c as usize)))
999                },
1000            )?;
1001            for (k, &j) in idx.iter().enumerate() {
1002                let (name, pk) = match classes {
1003                    None => {
1004                        let pc = p[k];
1005                        (if pc > 0.5 { "cjk" } else { "base" }, pc.max(1.0 - pc))
1006                    }
1007                    Some(c) => {
1008                        let row = &p[k * c..(k + 1) * c];
1009                        let g = argmax(row);
1010                        let name = self.script_groups.get(g).ok_or_else(|| {
1011                            OCRError::config_error(format!(
1012                                "script.onnx has {c} classes; meta.json script_groups names {}",
1013                                self.script_groups.len()
1014                            ))
1015                        })?;
1016                        (name.as_str(), row[g])
1017                    }
1018                };
1019                route[j] = self.groups.iter().position(|(g, _)| g == name);
1020                sure[j] = pk > ROUTE_SURE;
1021            }
1022        }
1023        Ok((route, sure))
1024    }
1025
1026    /// Lines split at wide gaps, tidied, reading order, layout blocks and
1027    /// Markdown. A right-to-left page reads its lines right to left, and
1028    /// its reading order and Markdown work on the page mirrored left to
1029    /// right, so the first column is the right one.
1030    fn assemble(
1031        &self,
1032        (pw, ph): (u32, u32),
1033        boxes: Vec<[f32; 5]>,
1034        lines: Vec<LineText>,
1035        vertical: bool,
1036    ) -> Result<IndraPage, OCRError> {
1037        let boxes: Vec<[f32; 4]> = boxes.iter().map(|b| [b[0], b[1], b[2], b[3]]).collect();
1038        let (letters, n_rtl) = lines
1039            .iter()
1040            .flat_map(|l| l.text.chars())
1041            .filter(|&c| rtl::is_letter(c))
1042            .fold((0, 0), |(n, r), c| (n + 1, r + rtl::is_rtl(c) as usize));
1043        let page_rtl = n_rtl > letters - n_rtl;
1044        let width = pw as f32;
1045        let geo = |b: [f32; 4]| {
1046            if page_rtl {
1047                [width - b[2], b[1], width - b[0], b[3]]
1048            } else {
1049                b
1050            }
1051        };
1052        let cuts = recognize::gutter_cuts(&boxes, &lines);
1053        let (mut nb, mut nl) = (Vec::new(), Vec::new());
1054        for ((b, line), at) in boxes.iter().zip(&lines).zip(&cuts) {
1055            for (bx, text, tokens) in recognize::split_line(*b, line, SPLIT_GAP, at, page_rtl) {
1056                nb.push(bx);
1057                nl.push(LineText {
1058                    text,
1059                    conf: line.conf,
1060                    tokens,
1061                    char_tokens: line.char_tokens,
1062                });
1063            }
1064        }
1065        let space = self.base.vocab.iter().any(|c| c == " ");
1066        let (mut boxes, mut lines) = recognize::tidy_lines(nb, nl, page_rtl, space);
1067        let mut texts: Vec<String> = lines.iter().map(|l| l.text.clone()).collect();
1068        let round1 = |v: f32| ((v as f64 * 10.0).round_ties_even() / 10.0) as f32;
1069        let mut page = IndraPage {
1070            index: 0,
1071            width: pw,
1072            height: ph,
1073            rotation: 0,
1074            skew: 0.0,
1075            vertical: false,
1076            lines: Vec::new(),
1077            blocks: Vec::new(),
1078            markdown: String::new(),
1079        };
1080        if boxes.len() < 2 {
1081            page.lines = boxes
1082                .iter()
1083                .zip(&lines)
1084                .map(|(b, l)| IndraLine {
1085                    bbox: b.map(round1),
1086                    text: l.text.clone(),
1087                    confidence: ((l.conf * 1e4).round_ties_even() / 1e4) as f32,
1088                })
1089                .collect();
1090            page.markdown = texts.join("\n");
1091            return Ok(page);
1092        }
1093        let mut geo_boxes: Vec<[f32; 4]> = boxes.iter().map(|&b| geo(b)).collect();
1094        let mut layout = self.layout(&geo_boxes, &texts, (pw, ph))?;
1095        let mut n = boxes.len();
1096        if vertical {
1097            // The order model saw no vertical pages: turned, their columns
1098            // of short lines look like one table read row by row. Blocks and
1099            // order come from XY-cut instead, without grid classes.
1100            let blocks = order::xycut_blocks(&boxes);
1101            layout.same = vec![-1.0; n * n];
1102            for b in &blocks {
1103                for &i in b {
1104                    for &j in b {
1105                        layout.same[i * n + j] = 1.0;
1106                    }
1107                }
1108            }
1109            let nc = order::BLOCK_CLASSES.len();
1110            for row in layout.cls.chunks_exact_mut(nc) {
1111                for c in order::GRID_CLASSES {
1112                    row[c] = -1e4;
1113                }
1114            }
1115            layout.order = blocks.concat();
1116            layout.tab = None;
1117        } else if layout.tab.is_some() {
1118            // The table heads learned one line per cell, but the detector
1119            // can join cells: split table lines into cells, then run the
1120            // order model again.
1121            let nc = order::BLOCK_CLASSES.len();
1122            let is_table: Vec<bool> = layout
1123                .cls
1124                .chunks_exact(nc)
1125                .map(|row| argmax(row) == layout::TABLE)
1126                .collect();
1127            if is_table.contains(&true) {
1128                let (mut nb, mut nl) = (Vec::new(), Vec::new());
1129                for ((b, line), table) in boxes.iter().zip(&lines).zip(is_table) {
1130                    if !table {
1131                        nb.push(*b);
1132                        nl.push(line.clone());
1133                        continue;
1134                    }
1135                    for (bx, text, tokens) in
1136                        recognize::split_line(*b, line, CELL_GAP, &[], page_rtl)
1137                    {
1138                        nb.push(bx);
1139                        nl.push(LineText {
1140                            text,
1141                            conf: line.conf,
1142                            tokens,
1143                            char_tokens: line.char_tokens,
1144                        });
1145                    }
1146                }
1147                if nb.len() != n {
1148                    (boxes, lines) = (nb, nl);
1149                    texts = lines.iter().map(|l| l.text.clone()).collect();
1150                    geo_boxes = boxes.iter().map(|&b| geo(b)).collect();
1151                    layout = self.layout(&geo_boxes, &texts, (pw, ph))?;
1152                    n = boxes.len();
1153                }
1154            }
1155        }
1156        let Layout {
1157            order: reading,
1158            cls,
1159            same,
1160            tab,
1161        } = layout;
1162        page.lines = reading
1163            .iter()
1164            .map(|&i| IndraLine {
1165                bbox: boxes[i].map(round1),
1166                text: lines[i].text.clone(),
1167                confidence: ((lines[i].conf * 1e4).round_ties_even() / 1e4) as f32,
1168            })
1169            .collect();
1170        // Inside tables, cells split at smaller gaps.
1171        let cells: Vec<Vec<([f32; 4], String)>> = boxes
1172            .iter()
1173            .zip(&lines)
1174            .map(|(b, l)| {
1175                recognize::split_line(*b, l, CELL_GAP, &[], page_rtl)
1176                    .into_iter()
1177                    .map(|(bx, text, _)| (geo(bx), text))
1178                    .collect()
1179            })
1180            .collect();
1181        let blocks = layout::page_blocks(
1182            &geo_boxes,
1183            &texts,
1184            &reading,
1185            &layout::Heads {
1186                cls: &cls,
1187                same: order::Square { n, data: &same },
1188                tab: tab.as_deref().map(|data| layout::TableHeads { n, data }),
1189            },
1190            self.keep_furniture,
1191            &cells,
1192        );
1193        page.markdown = layout::page_markdown(&blocks);
1194        let round = |v: f32| (v as f64 * 10.0).round_ties_even() / 10.0;
1195        page.blocks = blocks
1196            .into_iter()
1197            .map(|block| {
1198                let geo_boxes = &geo_boxes;
1199                let bx = |k: usize| block.lines.iter().map(move |&i| geo_boxes[i][k]);
1200                let mut bbox = [
1201                    round(bx(0).fold(f32::INFINITY, f32::min)),
1202                    round(bx(1).fold(f32::INFINITY, f32::min)),
1203                    round(bx(2).fold(f32::NEG_INFINITY, f32::max)),
1204                    round(bx(3).fold(f32::NEG_INFINITY, f32::max)),
1205                ];
1206                if page_rtl {
1207                    let w = pw as f64;
1208                    (bbox[0], bbox[2]) = (w - bbox[2], w - bbox[0]);
1209                }
1210                IndraBlock {
1211                    label: layout::label(block.class).to_string(),
1212                    bbox: bbox.map(|v| round1(v as f32)),
1213                    text: block
1214                        .lines
1215                        .iter()
1216                        .map(|&i| texts[i].as_str())
1217                        .collect::<Vec<_>>()
1218                        .join("\n"),
1219                    html: block.html,
1220                }
1221            })
1222            .collect();
1223        Ok(page)
1224    }
1225}
1226
1227/// The order model's reading of a page's lines.
1228struct Layout {
1229    order: Vec<usize>,
1230    /// Block class logits, `[n, classes]`.
1231    cls: Vec<f32>,
1232    /// Same-block logits, `[n, n]`.
1233    same: Vec<f32>,
1234    /// Same row, column and cell logits, `[3, n, n]`, with the table heads.
1235    tab: Option<Vec<f32>>,
1236}
1237
1238impl Engine {
1239    /// Runs the order model on line boxes (two or more) and decodes the
1240    /// reading order: blocks from the same-block head, lines inside a block
1241    /// by geometry, blocks by the order head.
1242    fn layout(
1243        &self,
1244        boxes: &[[f32; 4]],
1245        texts: &[String],
1246        wh: (u32, u32),
1247    ) -> Result<Layout, OCRError> {
1248        let n = boxes.len();
1249        let feats = order::line_features(boxes, wh, texts);
1250        let (a, cls, same, tab) = run(
1251            &self.order,
1252            ort::inputs![
1253                "feats" => Tensor::from_array(([1, n, order::N_FEATS], feats)).map_err(input_error)?,
1254                "mask" => Tensor::from_array(([1, n], vec![true; n])).map_err(input_error)?,
1255            ],
1256            |o| {
1257                let (_, a) = o["a"].try_extract_tensor::<f32>()?;
1258                let (_, cls) = o["cls"].try_extract_tensor::<f32>()?;
1259                let (_, same) = o["same"].try_extract_tensor::<f32>()?;
1260                let tab = match o.get("tab") {
1261                    Some(tab) if self.table_heads => {
1262                        Some(tab.try_extract_tensor::<f32>()?.1.to_vec())
1263                    }
1264                    _ => None,
1265                };
1266                Ok((a.to_vec(), cls.to_vec(), same.to_vec(), tab))
1267            },
1268        )?;
1269        let order = order::decode_blocks_order(
1270            &order::Square { n, data: &a },
1271            &order::Square { n, data: &same },
1272            boxes,
1273            &cls,
1274        );
1275        Ok(Layout {
1276            order,
1277            cls,
1278            same,
1279            tab,
1280        })
1281    }
1282}
1283
1284/// The first index of the largest value.
1285fn argmax(values: &[f32]) -> usize {
1286    let mut k = 0;
1287    for (i, v) in values.iter().enumerate() {
1288        if *v > values[k] {
1289            k = i;
1290        }
1291    }
1292    k
1293}
1294
1295/// numpy's mean of float32 values: a pairwise sum and a division in
1296/// float32.
1297fn mean_f32(values: &[f32]) -> f64 {
1298    fn pairwise(a: &[f32]) -> f32 {
1299        let n = a.len();
1300        if n < 8 {
1301            a.iter().fold(0.0, |s, &v| s + v)
1302        } else if n <= 128 {
1303            let mut r: [f32; 8] = a[..8].try_into().expect("8 values");
1304            let full = n - n % 8;
1305            for chunk in a[8..full].as_chunks::<8>().0 {
1306                for (r, v) in r.iter_mut().zip(chunk) {
1307                    *r += v;
1308                }
1309            }
1310            let mut res = ((r[0] + r[1]) + (r[2] + r[3])) + ((r[4] + r[5]) + (r[6] + r[7]));
1311            for v in &a[full..] {
1312                res += v;
1313            }
1314            res
1315        } else {
1316            let half = n / 2;
1317            let half = half - half % 8;
1318            pairwise(&a[..half]) + pairwise(&a[half..])
1319        }
1320    }
1321    (pairwise(values) / values.len() as f32) as f64
1322}
1323
1324/// The page shrunk once to at most `PREVIEW_SIDE` pixels a side (area
1325/// averaging): the skew search and the orientation model read it instead
1326/// of the full scan.
1327fn preview(gray: &GrayImage) -> Cow<'_, GrayImage> {
1328    let (w, h) = (gray.width() as f64, gray.height() as f64);
1329    let s = PREVIEW_SIDE / w.max(h);
1330    if s >= 1.0 {
1331        return Cow::Borrowed(gray);
1332    }
1333    let (sw, sh) = (
1334        ((w * s).round_ties_even() as u32).max(1),
1335        ((h * s).round_ties_even() as u32).max(1),
1336    );
1337    Cow::Owned(cv::resize(gray, sw, sh, cv::Interp::Area))
1338}
1339
1340/// The ink map (255 = ink) of the page shrunk to `side` pixels a side.
1341fn ink_thumb(gray: &GrayImage, side: f64) -> GrayImage {
1342    let (w, h) = (gray.width() as f64, gray.height() as f64);
1343    let s = side / w.max(h);
1344    let (sw, sh) = (
1345        ((w * s).round_ties_even() as u32).max(1),
1346        ((h * s).round_ties_even() as u32).max(1),
1347    );
1348    let mut small = cv::resize(gray, sw, sh, cv::Interp::Area);
1349    for v in small.iter_mut() {
1350        *v = 255 - *v;
1351    }
1352    small
1353}
1354
1355/// How sharply the horizontal ink profile changes row to row with the page
1356/// turned by `a` degrees: highest when the text rows are level. Integer row
1357/// sums keep the score exact.
1358fn row_sharpness(ink: &GrayImage, a: f64) -> f64 {
1359    let rows = cv::rotated_row_sums(ink, a);
1360    rows.windows(2).map(|p| (p[1] - p[0]).pow(2)).sum::<i64>() as f64
1361}
1362
1363/// Degrees (counter-clockwise) that straighten the text rows of a page:
1364/// the angle whose horizontal ink profile changes most sharply row to row.
1365/// Searched over ±`DESKEW_MAX` in 2 degree steps on a 400 px thumbnail;
1366/// the two best angles and 0 are then refined in 0.5 and 0.1 degree steps
1367/// on an 800 px thumbnail, where small text rows are still resolved.
1368fn estimate_skew(gray: &GrayImage) -> f64 {
1369    let fine = ink_thumb(gray, 800.0);
1370    // A one-row thumbnail has no row profile, and a blank or flat page has
1371    // no rows to straighten.
1372    let raw = fine.as_raw();
1373    let mean = raw.iter().map(|&v| v as u64).sum::<u64>() as f64 / raw.len() as f64;
1374    if fine.height() < 2 || mean < 0.003 * 255.0 || raw.iter().all(|&v| v == raw[0]) {
1375        return 0.0;
1376    }
1377    let coarse = ink_thumb(gray, 400.0);
1378    let grid = arange(-DESKEW_MAX, DESKEW_MAX + 1e-6, 2.0);
1379    let scores: Vec<f64> = grid
1380        .par_iter()
1381        .map(|&a| row_sharpness(&coarse, a))
1382        .collect();
1383    let mut best: Vec<usize> = (0..grid.len()).collect();
1384    best.sort_by(|&i, &j| scores[i].total_cmp(&scores[j]));
1385    let starts = [0.0, grid[best[best.len() - 2]], grid[best[best.len() - 1]]];
1386    let near: Vec<f64> = starts
1387        .iter()
1388        .flat_map(|&a0| arange(a0 - 2.0, a0 + 2.0 + 1e-6, 0.5))
1389        .collect();
1390    // The 0.5 degree ranges of nearby starts overlap: each angle is scored
1391    // once.
1392    let mut angles: Vec<f64> = Vec::new();
1393    for &a in &near {
1394        if !angles.contains(&a) {
1395            angles.push(a);
1396        }
1397    }
1398    let fine_scores: Vec<f64> = angles
1399        .par_iter()
1400        .map(|&a| row_sharpness(&fine, a))
1401        .collect();
1402    let score = |a: f64| fine_scores[angles.iter().position(|&b| b == a).expect("scored")];
1403    let first_best = |angles: &[f64], score: &dyn Fn(f64) -> f64| {
1404        let mut best = angles[0];
1405        for &a in angles {
1406            if score(a) > score(best) {
1407                best = a;
1408            }
1409        }
1410        best
1411    };
1412    let a1 = first_best(&near, &score);
1413    let last = arange(a1 - 0.5, a1 + 0.5 + 1e-6, 0.1);
1414    let last_scores: Vec<f64> = last.par_iter().map(|&a| row_sharpness(&fine, a)).collect();
1415    last[argmax_f64(&last_scores)]
1416}
1417
1418/// The first index of the largest value.
1419fn argmax_f64(values: &[f64]) -> usize {
1420    let mut k = 0;
1421    for (i, v) in values.iter().enumerate() {
1422        if *v > values[k] {
1423            k = i;
1424        }
1425    }
1426    k
1427}
1428
1429/// `numpy.arange`: values `start + i * delta` with `delta` taken from the
1430/// first two values, so the angles (and the deskew threshold test on them)
1431/// match the reference to the last bit.
1432fn arange(start: f64, stop: f64, step: f64) -> Vec<f64> {
1433    let n = ((stop - start) / step).ceil().max(0.0) as usize;
1434    let delta = (start + step) - start;
1435    (0..n).map(|i| start + i as f64 * delta).collect()
1436}
1437
1438/// A box of the page `rotate_u8(.., degrees)` made, mapped back to the page
1439/// before it: the bounding box of the four mapped corners, clipped to it.
1440fn unrotate_box(b: [f32; 4], degrees: f64, (w, h): (u32, u32)) -> [f32; 4] {
1441    let (w, h) = (w as f64, h as f64);
1442    let m = cv::rotation_matrix(w / 2.0, h / 2.0, -degrees);
1443    let [x0, y0, x1, y1] = b.map(f64::from);
1444    let corners = [(x0, y0), (x1, y0), (x0, y1), (x1, y1)];
1445    let xs = corners.map(|(x, y)| m[0] * x + m[1] * y + m[2]);
1446    let ys = corners.map(|(x, y)| m[3] * x + m[4] * y + m[5]);
1447    let min = |v: [f64; 4]| v.iter().copied().fold(f64::INFINITY, f64::min);
1448    let max = |v: [f64; 4]| v.iter().copied().fold(f64::NEG_INFINITY, f64::max);
1449    let r = |v: f64, hi: f64| ((v.clamp(0.0, hi) * 10.0).round_ties_even() / 10.0) as f32;
1450    [r(min(xs), w), r(min(ys), h), r(max(xs), w), r(max(ys), h)]
1451}
1452
1453/// Loads one model of the folder.
1454fn model(dir: &Path, name: &str, common: &ModelInferenceConfig) -> Result<OrtInfer, OCRError> {
1455    let mut common = common.clone();
1456    common.model_name = Some(format!("indra {name}"));
1457    OrtInfer::from_config(&common, dir.join(name), None)
1458}
1459
1460/// Creates a separate session for another worker without reloading metadata.
1461fn reload_model(original: &OrtInfer, common: &ModelInferenceConfig) -> Result<OrtInfer, OCRError> {
1462    let mut common = common.clone();
1463    common.model_name = Some(original.model_name().to_owned());
1464    OrtInfer::from_config(&common, original.model_path(), Some(original.input_name()))
1465}
1466
1467/// A non-empty JSON array of non-empty strings (a vocabulary).
1468fn string_list(value: &serde_json::Value, what: &str) -> Result<Vec<String>, OCRError> {
1469    let list: Option<Vec<String>> = value.as_array().filter(|a| !a.is_empty()).and_then(|a| {
1470        a.iter()
1471            .map(|v| v.as_str().filter(|s| !s.is_empty()).map(str::to_string))
1472            .collect()
1473    });
1474    list.ok_or_else(|| {
1475        OCRError::config_error(format!(
1476            "{what} must be a non-empty list of non-empty strings"
1477        ))
1478    })
1479}
1480
1481fn read_json(path: &Path) -> Result<serde_json::Value, OCRError> {
1482    let text = std::fs::read_to_string(path).map_err(|e| {
1483        OCRError::model_load_error(path, "cannot read", Some("check the model folder"), Some(e))
1484    })?;
1485    serde_json::from_str(&text)
1486        .map_err(|e| OCRError::config_error(format!("{}: {e}", path.display())))
1487}
1488
1489/// Runs a model on named inputs and reads its outputs while the session is
1490/// held.
1491fn run<R>(
1492    model: &OrtInfer,
1493    inputs: Vec<(Cow<'_, str>, SessionInputValue<'_>)>,
1494    read: impl FnOnce(&SessionOutputs<'_>) -> ort::Result<R>,
1495) -> Result<R, OCRError> {
1496    let mut session = model.get_session(0)?;
1497    // Keep the run options (CUDA arena shrinkage) that `OrtInfer` applies.
1498    let outputs = match model.run_options() {
1499        Some(options) => session.run_with_options(inputs, options),
1500        None => session.run(inputs),
1501    }
1502    .map_err(|e| OCRError::inference_error(model.model_name(), "forward pass", e))?;
1503    read(&outputs).map_err(|e| OCRError::inference_error(model.model_name(), "reading outputs", e))
1504}
1505
1506/// Maps an input tensor that ONNX Runtime rejects to an error.
1507fn input_error(e: ort::Error) -> OCRError {
1508    OCRError::invalid_input(format!("Indra OCR model input: {e}"))
1509}
1510
1511#[cfg(test)]
1512mod tests {
1513    use super::*;
1514
1515    /// numpy's arange makes the third fine angle around 0 read
1516    /// -0.30000000000000004, which passes the deskew threshold; `-0.5 + 0.2`
1517    /// would not.
1518    #[test]
1519    fn deskew_angles_follow_numpy_arange() {
1520        let angles = arange(-0.5, 0.5 + 1e-6, 0.1);
1521        assert_eq!(angles.len(), 11);
1522        assert!(angles[2].abs() >= DESKEW_MIN);
1523    }
1524
1525    #[test]
1526    fn parallel_pages_match_sequential() {
1527        use oar_ocr_core::core::config::OrtExecutionProvider;
1528
1529        let dir = Path::new(env!("CARGO_MANIFEST_DIR")).join("models/datalab-to/indra-ocr-v1.1.0");
1530        if !dir.join("meta.json").is_file() {
1531            eprintln!("skipping Indra page concurrency test: local models are absent");
1532            return;
1533        }
1534        // A small raster page spelling TEST twice; no font or image fixture
1535        // dependency. Unequal sizes, rotation, a blank and a repeated page
1536        // make misplaced or lost results visible in a five-page batch.
1537        let glyphs = [
1538            [31u8, 4, 4, 4, 4, 4, 4],
1539            [31, 16, 16, 30, 16, 16, 31],
1540            [15, 16, 16, 14, 1, 1, 30],
1541            [31, 4, 4, 4, 4, 4, 4],
1542        ];
1543        let mut page = RgbImage::from_pixel(256, 160, image::Rgb([255; 3]));
1544        for y0 in [30, 90] {
1545            for (letter, rows) in glyphs.iter().enumerate() {
1546                for (row, bits) in rows.iter().enumerate() {
1547                    for col in 0..5 {
1548                        if bits & (1 << (4 - col)) != 0 {
1549                            for dy in 0..4 {
1550                                for dx in 0..4 {
1551                                    page.put_pixel(
1552                                        24 + letter as u32 * 28 + col * 4 + dx,
1553                                        y0 + row as u32 * 4 + dy,
1554                                        image::Rgb([0; 3]),
1555                                    );
1556                                }
1557                            }
1558                        }
1559                    }
1560                }
1561            }
1562        }
1563        let images = vec![
1564            page.clone(),
1565            image::imageops::rotate90(&page),
1566            RgbImage::from_pixel(192, 128, image::Rgb([255; 3])),
1567            page.clone(),
1568            image::imageops::rotate180(&page),
1569        ];
1570        let cpu = OrtSessionConfig::new().with_execution_providers(vec![OrtExecutionProvider::CPU]);
1571        let serial = IndraOcrBuilder::new(&dir)
1572            .ort_session(cpu.clone().with_intra_threads(1).with_inter_threads(1))
1573            .build()
1574            .unwrap();
1575        let expected = serial.predict(images.clone()).unwrap();
1576        drop(serial);
1577        let parallel = IndraOcrBuilder::new(&dir)
1578            .workers(3)
1579            .ort_session(cpu)
1580            .build()
1581            .unwrap();
1582        for _ in 0..2 {
1583            let actual = parallel.predict(images.clone()).unwrap();
1584            // Serialize all public fields, including lines and confidences;
1585            // the shared export additionally checks the private table HTML.
1586            assert_eq!(
1587                serde_json::to_value(&actual).unwrap(),
1588                serde_json::to_value(&expected).unwrap()
1589            );
1590            for (index, (a, e)) in actual.iter().zip(&expected).enumerate() {
1591                assert_eq!(a.index, index);
1592                assert_eq!(a.to_json(), e.to_json());
1593            }
1594        }
1595    }
1596}