1mod 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
69const ADAPTIVE_MIN_H: f32 = 9.0;
72const ADAPTIVE_TARGET_H: f32 = 12.0;
73const ADAPTIVE_MAX_SIDE: usize = 2048;
74const SPLIT_GAP: f64 = 3.0;
77const CELL_GAP: f64 = 1.8;
79const ORIENT_SIDE: usize = 768;
81const ORIENT_MIN: f32 = 0.95;
83const VERTICAL: u32 = 4;
85const PREVIEW_SIDE: f64 = 1600.0;
88const MIN_CONF: f64 = 0.5;
90const DESKEW_MAX: f64 = 45.0;
92const DESKEW_MIN: f64 = 0.3;
94const REC_BATCH: usize = 64;
96const ROUTE_SURE: f32 = 0.95;
100const GROUPS: [&str; 4] = ["base", "cjk", "rtl", "other"];
103const LM_GATE: f64 = 0.9;
106const LM_ALPHA: f64 = 0.5;
107const LM_BETA: f64 = 1.5;
108const LM_BEAM: usize = 8;
109const LM_CAND_FLOOR: f64 = 3e-3;
111
112#[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 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 pub fn math(mut self, enable: bool) -> Self {
141 self.math = enable;
142 self
143 }
144
145 pub fn keep_furniture(mut self, enable: bool) -> Self {
148 self.keep_furniture = enable;
149 self
150 }
151
152 pub fn table_heads(mut self, enable: bool) -> Self {
159 self.table_heads = enable;
160 self
161 }
162
163 pub fn ort_session(mut self, config: OrtSessionConfig) -> Self {
166 self.ort_session = Some(config);
167 self
168 }
169
170 pub fn workers(mut self, workers: usize) -> Self {
181 self.workers = workers;
182 self
183 }
184
185 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 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 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#[derive(Debug)]
328struct Recognizer {
329 model: OrtInfer,
330 vocab: Arc<[String]>,
331 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#[derive(Debug)]
349struct Lm {
350 model: lm::NgramLm,
351 gate: f64,
353 weights: lm::Weights,
354}
355
356struct 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#[derive(Debug)]
373pub struct IndraOcr {
374 engines: Vec<Engine>,
375 pool: Option<rayon::ThreadPool>,
376}
377
378#[derive(Debug)]
380struct Engine {
381 det: OrtInfer,
382 base: Recognizer,
385 order: OrtInfer,
386 orient: Option<OrtInfer>,
387 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 window: usize,
397 keep_furniture: bool,
398 table_heads: bool,
399}
400
401#[derive(Debug, Clone, Serialize)]
403#[non_exhaustive]
404pub struct IndraLine {
405 pub bbox: [f32; 4],
407 pub text: String,
409 pub confidence: f32,
411}
412
413#[derive(Debug, Clone, Serialize)]
415#[non_exhaustive]
416pub struct IndraBlock {
417 pub label: String,
421 pub bbox: [f32; 4],
426 pub text: String,
428 #[serde(skip)]
430 html: Option<String>,
431}
432
433#[derive(Debug, Clone, Serialize)]
435#[non_exhaustive]
436pub struct IndraPage {
437 pub index: usize,
439 pub width: u32,
441 pub height: u32,
443 pub rotation: u32,
446 pub skew: f32,
449 pub vertical: bool,
453 pub lines: Vec<IndraLine>,
455 pub blocks: Vec<IndraBlock>,
458 pub markdown: String,
461}
462
463impl IndraPage {
464 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 "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 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 let mut results: Vec<_> = pool
558 .broadcast(|worker| {
559 let engine = &self.engines[worker.index()];
560 let mut results = Vec::new();
561 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 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 let turns = if vertical { 3 } else { k % 4 };
622 let gray = cv::rot90_cw(&gray, turns);
623 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 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 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 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 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 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 let tall = (self.window > 1).then(|| crop_all(self.window));
772 let n = crops.len();
773 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 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 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 let f = cw / src[j].width as f64;
885 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 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 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 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 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 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 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 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
1227struct Layout {
1229 order: Vec<usize>,
1230 cls: Vec<f32>,
1232 same: Vec<f32>,
1234 tab: Option<Vec<f32>>,
1236}
1237
1238impl Engine {
1239 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
1284fn 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
1295fn 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
1324fn 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
1340fn 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
1355fn 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
1363fn estimate_skew(gray: &GrayImage) -> f64 {
1369 let fine = ink_thumb(gray, 800.0);
1370 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 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
1418fn 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
1429fn 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
1438fn 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
1453fn 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
1460fn 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
1467fn 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
1489fn 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 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
1506fn 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 #[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 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 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}