From 12d30b43b01ebd602eb7086d7cf6d65e7dc7194c Mon Sep 17 00:00:00 2001 From: Abimael Martell <1450169+abimaelmartell@users.noreply.github.com> Date: Sun, 16 Aug 2026 23:55:51 -0700 Subject: [PATCH] feat(vision): add OAR OCR engine (#357) * feat(vision): add OAR OCR engine * fix(vision): harden OAR runtime loading * refactor(vision): use OCR engine terminology --- Cargo.toml | 6 + docs/rust-api.md | 50 +++++ src/tables/grid.rs | 4 +- src/vision/mod.rs | 11 +- src/vision/oar.rs | 469 +++++++++++++++++++++++++++++++++++++++++++++ tests/ocr_tests.rs | 127 ++++++++++++ 6 files changed, 662 insertions(+), 5 deletions(-) create mode 100644 src/vision/oar.rs create mode 100644 tests/ocr_tests.rs diff --git a/Cargo.toml b/Cargo.toml index 8b11d24..be51de7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -57,6 +57,11 @@ firecrawl-pdfium = { version = "0.1.0", optional = true } dirs = { version = "6.0", optional = true } fs2 = { version = "0.4", optional = true } sha2 = { version = "0.11", optional = true } +# Optional CPU OCR backend. Models and ONNX Runtime stay external: the latter +# is loaded dynamically from ORT_DYLIB_PATH or the platform library search path. +image = { version = "0.25.6", default-features = false, optional = true } +oar-ocr = { version = "0.9.1", default-features = false, features = ["simd"], optional = true } +ort = { version = "=2.0.0-rc.13", default-features = false, features = ["load-dynamic"], optional = true } [target.'cfg(all(windows, not(target_arch = "wasm32")))'.dependencies] windows-sys = { version = "0.61", features = ["Win32_Storage_FileSystem"], optional = true } @@ -75,6 +80,7 @@ default = [] python = ["pyo3"] vision = [] model-cache = ["vision", "dep:dirs", "dep:fs2", "dep:sha2", "dep:windows-sys"] +ocr-oar = ["model-cache", "dep:image", "dep:oar-ocr", "dep:ort"] render-pdfium = ["vision", "dep:firecrawl-pdfium"] [[bin]] diff --git a/docs/rust-api.md b/docs/rust-api.md index 4f614ab..81d5f26 100644 --- a/docs/rust-api.md +++ b/docs/rust-api.md @@ -197,6 +197,56 @@ for page in pages { Browser WASM remains on the default text-only path and does not expose native PDFium rendering. +### Optional OCR engine + +The native-only `ocr-oar` feature adds a CPU PP-OCRv6 Small implementation of +`OcrEngine` backed by OAR and ONNX Runtime. It implies `model-cache`, but does +not enable model auto-download, ONNX Runtime download, or PDF rendering. Model +files remain external, must match the pinned manifest, and are opened only +after `ModelStore` verifies their exact size and SHA-256 digest. Install an +ONNX Runtime shared library separately and set `ORT_DYLIB_PATH` when it is not +available through the platform library search path. The feature currently +requires Rust 1.95 or newer, matching OAR 0.9.1's MSRV. + +```toml +[dependencies] +pdf-inspector = { version = "1", features = ["ocr-oar", "render-pdfium"] } +``` + +Direct engine invocation is intentionally separate from extraction routing and +native/OCR fusion: + +```rust +use pdf_inspector::vision::{ + ModelDownloadPolicy, ModelStore, OarOcrEngine, OcrEngine, OcrMode, + OcrOptions, PdfiumRenderer, RenderOptions, PP_OCR_V6_SMALL, +}; + +let options = OcrOptions::new() + .mode(OcrMode::Force) + .minimum_confidence(0.45) + .model_directory("/opt/firecrawl/models/pp-ocrv6-small") + .model_downloads(ModelDownloadPolicy::Offline); +let models = ModelStore::from_options(&options)?.resolve(&PP_OCR_V6_SMALL)?; +let engine = OarOcrEngine::from_models(&models)?; + +let renderer = PdfiumRenderer::load()?; +let bytes = std::fs::read("scan.pdf")?; +let pages = renderer.render_pages(&bytes, &[1], None, &RenderOptions::new())?; +let ocr_pages = engine.recognize(&pages, &options)?; + +for span in &ocr_pages[0].spans { + println!("{:.3}: {}", span.confidence, span.text); +} +``` + +The engine accepts renderer-neutral RGB, RGBA, and grayscale pages, preserves +OAR's positioned quadrilaterals in bitmap coordinates, filters spans using +`minimum_confidence`, and records the pinned model revision in every `OcrPage`. +`OcrMode::Off` is rejected at the engine boundary so default options cannot run +inference accidentally. Selective routing and OCR/native-text fusion are added +by higher stack layers. + Extract per-page Markdown (one string per page, plus document-wide layout metadata): diff --git a/src/tables/grid.rs b/src/tables/grid.rs index dcf8fdf..de4356b 100644 --- a/src/tables/grid.rs +++ b/src/tables/grid.rs @@ -604,7 +604,7 @@ mod tests { let items: Vec<(usize, &TextItem)> = vec![]; assert_eq!( find_column_boundaries(&items, TableDetectionMode::SmallFont), - vec![] + Vec::::new() ); } @@ -661,7 +661,7 @@ mod tests { #[test] fn test_find_row_boundaries_empty() { let items: Vec<(usize, &TextItem)> = vec![]; - assert_eq!(find_row_boundaries(&items), vec![]); + assert_eq!(find_row_boundaries(&items), Vec::::new()); } #[test] diff --git a/src/vision/mod.rs b/src/vision/mod.rs index 2d489aa..45e9e6f 100644 --- a/src/vision/mod.rs +++ b/src/vision/mod.rs @@ -3,14 +3,17 @@ //! The existing lopdf extractor remains the default path. Native page //! rendering is available only with the `render-pdfium` feature. Engine //! contracts are available with `vision`, while checksum-verified model -//! resolution is a separate `model-cache` feature. These remain separate so -//! browser WASM, text-only consumers, and renderer-only users take on no model -//! management dependencies. +//! resolution is a separate `model-cache` feature. The `ocr-oar` feature adds +//! a CPU PP-OCRv6 Small implementation of [`OcrEngine`]. These remain separate +//! so browser WASM, text-only consumers, and renderer-only users take on no +//! model-management or inference dependencies. #[cfg(all(feature = "vision", not(target_arch = "wasm32")))] mod contracts; #[cfg(all(feature = "model-cache", not(target_arch = "wasm32")))] mod models; +#[cfg(all(feature = "ocr-oar", not(target_arch = "wasm32")))] +mod oar; #[cfg(all(feature = "vision", not(target_arch = "wasm32")))] mod render; @@ -28,6 +31,8 @@ pub use models::{ ModelArtifact, ModelArtifactKind, ModelManifest, ModelPaths, ModelStore, ModelStoreError, PP_OCR_V6_SMALL, }; +#[cfg(all(feature = "ocr-oar", not(target_arch = "wasm32")))] +pub use oar::{OarOcrEngine, OarOcrError, ONNX_RUNTIME_LIBRARY_ENV}; #[cfg(all(feature = "vision", not(target_arch = "wasm32")))] pub use render::{ PagePoint, PageTransform, RenderBufferError, RenderOptions, RenderPixelFormat, RenderedPage, diff --git a/src/vision/oar.rs b/src/vision/oar.rs new file mode 100644 index 0000000..a33f6f9 --- /dev/null +++ b/src/vision/oar.rs @@ -0,0 +1,469 @@ +//! PP-OCRv6 Small implementation backed by OAR and ONNX Runtime. + +use std::path::PathBuf; +use std::time::Instant; + +use image::RgbImage; +use oar_ocr::oarocr::{OAROCRBuilder, OAROCR}; +use oar_ocr::processors::BoundingBox; +use thiserror::Error; + +use super::{ + ImagePoint, ImageQuad, ModelArtifactKind, ModelIdentity, ModelPaths, OcrEngine, OcrMode, + OcrOptions, OcrPage, OcrSpan, RenderPixelFormat, RenderedPage, +}; + +/// Environment variable selecting the ONNX Runtime shared library. +pub const ONNX_RUNTIME_LIBRARY_ENV: &str = "ORT_DYLIB_PATH"; + +/// Failures while constructing or running the OAR OCR backend. +#[derive(Debug, Error)] +#[non_exhaustive] +pub enum OarOcrError { + /// A required file is missing from the resolved model set. + #[error("resolved OCR model set is missing {kind:?}")] + MissingModelArtifact { + /// Missing artifact role. + kind: ModelArtifactKind, + }, + /// OCR was invoked while the caller explicitly disabled it. + #[error("OCR is disabled; select Auto or Force before invoking the engine")] + OcrDisabled, + /// Confidence thresholds must match the normalized engine output range. + #[error("minimum OCR confidence must be finite and between 0 and 1, got {value}")] + InvalidMinimumConfidence { + /// Invalid threshold. + value: f32, + }, + /// Bitmap dimension arithmetic exceeded the host address space. + #[error("rendered page {page} bitmap dimensions overflow the host address space")] + ImageSizeOverflow { + /// 1-indexed page number. + page: u32, + }, + /// A validated renderer buffer could not be represented as an RGB image. + #[error("rendered page {page} could not be converted to an RGB image")] + InvalidImageBuffer { + /// 1-indexed page number. + page: u32, + }, + /// The external ONNX Runtime shared library could not be loaded. + #[error("failed to load ONNX Runtime from {path}: {source}")] + OnnxRuntimeLoad { + /// Requested shared-library path or platform library name. + path: PathBuf, + /// Dynamic-loader failure. + #[source] + source: ort::LoadDynamicError, + }, + /// OAR returned no result for a submitted page. + #[error("OAR returned no result for rendered page {page}")] + MissingPageResult { + /// 1-indexed page number. + page: u32, + }, + /// OAR or ONNX Runtime rejected the models or failed during inference. + #[error(transparent)] + Backend(#[from] oar_ocr::core::OCRError), +} + +/// CPU PP-OCRv6 Small engine using OAR's detection and recognition pipeline. +/// +/// Construction accepts only [`ModelPaths`] that have already passed +/// pdf-inspector's manifest size and SHA-256 verification. OAR's independent +/// model auto-download feature is deliberately not enabled. +#[derive(Debug)] +pub struct OarOcrEngine { + pipeline: OAROCR, + model: ModelIdentity, +} + +impl OarOcrEngine { + /// Loads PP-OCRv6 Small from a resolved, verified model set. + pub fn from_models(models: &ModelPaths) -> Result { + load_onnx_runtime()?; + let detection = required_model(models, ModelArtifactKind::TextDetection)?; + let recognition = required_model(models, ModelArtifactKind::TextRecognition)?; + let dictionary = required_model(models, ModelArtifactKind::CharacterDictionary)?; + + let pipeline = OAROCRBuilder::new(detection, recognition, dictionary).build()?; + let model = ModelIdentity::new(models.manifest_id(), models.revision()); + Ok(Self { pipeline, model }) + } + + fn recognize_page( + &self, + page: &RenderedPage, + options: &OcrOptions, + ) -> Result { + let started = Instant::now(); + let image = rendered_page_to_rgb(page)?; + let result = self + .pipeline + .predict(vec![image])? + .into_iter() + .next() + .ok_or(OarOcrError::MissingPageResult { page: page.page() })?; + + let mut spans = Vec::with_capacity(result.text_regions.len()); + let mut invalid_geometry = 0usize; + let mut missing_recognition = 0usize; + for region in result.text_regions { + let (Some(text), Some(confidence)) = (region.text, region.confidence) else { + missing_recognition += 1; + continue; + }; + if text.trim().is_empty() || !confidence.is_finite() { + missing_recognition += 1; + continue; + } + let confidence = confidence.clamp(0.0, 1.0); + if confidence < options.minimum_confidence { + continue; + } + + let polygon = region.dt_poly.as_ref().unwrap_or(®ion.bounding_box); + let Some(polygon) = bounding_box_to_quad(polygon, page.width(), page.height()) else { + invalid_geometry += 1; + continue; + }; + spans.push(OcrSpan { + text: text.to_string(), + polygon, + confidence, + orientation_degrees: region.orientation_angle, + }); + } + + let mut warnings = Vec::new(); + if !options.languages.is_empty() { + warnings + .push("language hints are not used by the PP-OCRv6 Small OAR backend".to_string()); + } + if missing_recognition > 0 { + warnings.push(format!( + "discarded {missing_recognition} regions without usable recognition output" + )); + } + if invalid_geometry > 0 { + warnings.push(format!( + "discarded {invalid_geometry} recognized regions with invalid geometry" + )); + } + + let mean_confidence = if spans.is_empty() { + None + } else { + Some(spans.iter().map(|span| span.confidence).sum::() / spans.len() as f32) + }; + let processing_time_ms = u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX); + + Ok(OcrPage { + page: page.page(), + spans, + mean_confidence, + model: self.model.clone(), + processing_time_ms, + warnings, + }) + } +} + +fn load_onnx_runtime() -> Result<(), OarOcrError> { + let path = std::env::var_os(ONNX_RUNTIME_LIBRARY_ENV) + .filter(|path| !path.is_empty()) + .map(PathBuf::from) + .unwrap_or_else(default_onnx_runtime_library); + drop( + ort::init_from(&path).map_err(|source| OarOcrError::OnnxRuntimeLoad { + path: path.clone(), + source, + })?, + ); + Ok(()) +} + +fn default_onnx_runtime_library() -> PathBuf { + #[cfg(target_os = "windows")] + const NAME: &str = "onnxruntime.dll"; + #[cfg(any(target_os = "linux", target_os = "android", target_os = "freebsd"))] + const NAME: &str = "libonnxruntime.so"; + #[cfg(any(target_os = "macos", target_os = "ios"))] + const NAME: &str = "libonnxruntime.dylib"; + PathBuf::from(NAME) +} + +impl OcrEngine for OarOcrEngine { + type Error = OarOcrError; + + fn model(&self) -> &ModelIdentity { + &self.model + } + + fn recognize( + &self, + pages: &[RenderedPage], + options: &OcrOptions, + ) -> Result, Self::Error> { + validate_options(options)?; + + pages + .iter() + .map(|page| self.recognize_page(page, options)) + .collect() + } +} + +fn validate_options(options: &OcrOptions) -> Result<(), OarOcrError> { + if options.mode == OcrMode::Off { + return Err(OarOcrError::OcrDisabled); + } + if !options.minimum_confidence.is_finite() || !(0.0..=1.0).contains(&options.minimum_confidence) + { + return Err(OarOcrError::InvalidMinimumConfidence { + value: options.minimum_confidence, + }); + } + Ok(()) +} + +fn required_model( + models: &ModelPaths, + kind: ModelArtifactKind, +) -> Result<&std::path::Path, OarOcrError> { + models + .get(kind) + .ok_or(OarOcrError::MissingModelArtifact { kind }) +} + +fn rendered_page_to_rgb(page: &RenderedPage) -> Result { + let width = usize::try_from(page.width()) + .map_err(|_| OarOcrError::ImageSizeOverflow { page: page.page() })?; + let height = usize::try_from(page.height()) + .map_err(|_| OarOcrError::ImageSizeOverflow { page: page.page() })?; + let output_len = width + .checked_mul(height) + .and_then(|pixels| pixels.checked_mul(3)) + .ok_or(OarOcrError::ImageSizeOverflow { page: page.page() })?; + let input_bpp = page.format().bytes_per_pixel(); + let active_input_row = width + .checked_mul(input_bpp) + .ok_or(OarOcrError::ImageSizeOverflow { page: page.page() })?; + let output_row = width + .checked_mul(3) + .ok_or(OarOcrError::ImageSizeOverflow { page: page.page() })?; + + let mut rgb = vec![0u8; output_len]; + for row in 0..height { + let input_start = row * page.stride(); + let input = &page.pixels()[input_start..input_start + active_input_row]; + let output_start = row * output_row; + let output = &mut rgb[output_start..output_start + output_row]; + match page.format() { + RenderPixelFormat::Rgb8 => output.copy_from_slice(input), + RenderPixelFormat::Rgba8 => { + for (rgba, rgb) in input.chunks_exact(4).zip(output.chunks_exact_mut(3)) { + rgb.copy_from_slice(&rgba[..3]); + } + } + RenderPixelFormat::Gray8 => { + for (&gray, rgb) in input.iter().zip(output.chunks_exact_mut(3)) { + rgb.fill(gray); + } + } + } + } + + RgbImage::from_raw(page.width(), page.height(), rgb) + .ok_or(OarOcrError::InvalidImageBuffer { page: page.page() }) +} + +fn bounding_box_to_quad(bounding_box: &BoundingBox, width: u32, height: u32) -> Option { + let points: Vec = bounding_box + .points + .iter() + .filter(|point| point.x.is_finite() && point.y.is_finite()) + .map(|point| { + ImagePoint::new( + point.x.clamp(0.0, width as f32), + point.y.clamp(0.0, height as f32), + ) + }) + .collect(); + + if bounding_box.points.len() == 4 && points.len() == 4 && is_ordered_convex_quad(&points) { + return Some(ImageQuad::new([points[0], points[1], points[2], points[3]])); + } + if points.len() < 3 { + return None; + } + + let min_x = points + .iter() + .map(|point| point.x) + .fold(f32::INFINITY, f32::min); + let max_x = points + .iter() + .map(|point| point.x) + .fold(f32::NEG_INFINITY, f32::max); + let min_y = points + .iter() + .map(|point| point.y) + .fold(f32::INFINITY, f32::min); + let max_y = points + .iter() + .map(|point| point.y) + .fold(f32::NEG_INFINITY, f32::max); + if max_x <= min_x || max_y <= min_y { + return None; + } + Some(ImageQuad::new([ + ImagePoint::new(min_x, min_y), + ImagePoint::new(max_x, min_y), + ImagePoint::new(max_x, max_y), + ImagePoint::new(min_x, max_y), + ])) +} + +fn is_ordered_convex_quad(points: &[ImagePoint]) -> bool { + if points.len() != 4 { + return false; + } + let mut orientation = 0.0_f32; + for index in 0..4 { + let first = points[index]; + let second = points[(index + 1) % 4]; + let third = points[(index + 2) % 4]; + let cross = (second.x - first.x) * (third.y - second.y) + - (second.y - first.y) * (third.x - second.x); + if cross.abs() <= f32::EPSILON { + return false; + } + if orientation == 0.0 { + orientation = cross.signum(); + } else if cross.signum() != orientation { + return false; + } + } + true +} + +#[cfg(test)] +mod tests { + use oar_ocr::processors::Point; + + use super::*; + use crate::vision::PageTransform; + + fn page(format: RenderPixelFormat, stride: usize, pixels: Vec) -> RenderedPage { + let transform = + PageTransform::from_corners(2, 2, (0.0, 2.0), (2.0, 2.0), (0.0, 0.0)).unwrap(); + RenderedPage::new(1, 2.0, 2.0, 2, 2, stride, format, pixels, transform).unwrap() + } + + #[test] + fn converts_padded_rgb_without_exposing_padding() { + let page = page( + RenderPixelFormat::Rgb8, + 8, + vec![1, 2, 3, 4, 5, 6, 99, 99, 7, 8, 9, 10, 11, 12, 99, 99], + ); + let image = rendered_page_to_rgb(&page).unwrap(); + assert_eq!(image.as_raw(), &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12]); + } + + #[test] + fn converts_rgba_and_gray_to_rgb() { + let rgba = page( + RenderPixelFormat::Rgba8, + 8, + vec![1, 2, 3, 44, 4, 5, 6, 55, 7, 8, 9, 66, 10, 11, 12, 77], + ); + assert_eq!( + rendered_page_to_rgb(&rgba).unwrap().as_raw(), + &[1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12] + ); + + let gray = page(RenderPixelFormat::Gray8, 2, vec![1, 2, 3, 4]); + assert_eq!( + rendered_page_to_rgb(&gray).unwrap().as_raw(), + &[1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4] + ); + } + + #[test] + fn preserves_quads_and_clamps_them_to_the_bitmap() { + let bbox = BoundingBox::new(vec![ + Point::new(-1.0, 2.0), + Point::new(11.0, 2.0), + Point::new(11.0, 9.0), + Point::new(-1.0, 9.0), + ]); + let quad = bounding_box_to_quad(&bbox, 10, 8).unwrap(); + assert_eq!(quad.points[0], ImagePoint::new(0.0, 2.0)); + assert_eq!(quad.points[2], ImagePoint::new(10.0, 8.0)); + } + + #[test] + fn reduces_polygons_to_a_stable_axis_aligned_quad() { + let bbox = BoundingBox::new(vec![ + Point::new(2.0, 1.0), + Point::new(7.0, 2.0), + Point::new(8.0, 6.0), + Point::new(5.0, 9.0), + Point::new(1.0, 5.0), + ]); + let quad = bounding_box_to_quad(&bbox, 10, 10).unwrap(); + assert_eq!(quad.points[0], ImagePoint::new(1.0, 1.0)); + assert_eq!(quad.points[2], ImagePoint::new(8.0, 9.0)); + } + + #[test] + fn normalizes_unordered_or_partially_invalid_quads() { + let unordered = BoundingBox::new(vec![ + Point::new(1.0, 1.0), + Point::new(8.0, 8.0), + Point::new(8.0, 1.0), + Point::new(1.0, 8.0), + ]); + let quad = bounding_box_to_quad(&unordered, 10, 10).unwrap(); + assert_eq!(quad.points[0], ImagePoint::new(1.0, 1.0)); + assert_eq!(quad.points[1], ImagePoint::new(8.0, 1.0)); + assert_eq!(quad.points[2], ImagePoint::new(8.0, 8.0)); + + let partially_invalid = BoundingBox::new(vec![ + Point::new(8.0, 8.0), + Point::new(f32::NAN, 4.0), + Point::new(1.0, 8.0), + Point::new(8.0, 1.0), + Point::new(1.0, 1.0), + ]); + let quad = bounding_box_to_quad(&partially_invalid, 10, 10).unwrap(); + assert_eq!(quad.points[0], ImagePoint::new(1.0, 1.0)); + assert_eq!(quad.points[1], ImagePoint::new(8.0, 1.0)); + assert_eq!(quad.points[2], ImagePoint::new(8.0, 8.0)); + } + + #[test] + fn refuses_disabled_or_invalid_options_before_inference() { + assert!(matches!( + validate_options(&OcrOptions::new()), + Err(OarOcrError::OcrDisabled) + )); + for value in [-0.1, 1.1, f32::NAN, f32::INFINITY] { + let options = OcrOptions::new() + .mode(OcrMode::Force) + .minimum_confidence(value); + assert!(matches!( + validate_options(&options), + Err(OarOcrError::InvalidMinimumConfidence { .. }) + )); + } + assert!(validate_options( + &OcrOptions::new() + .mode(OcrMode::Auto) + .minimum_confidence(1.0) + ) + .is_ok()); + } +} diff --git a/tests/ocr_tests.rs b/tests/ocr_tests.rs new file mode 100644 index 0000000..486a034 --- /dev/null +++ b/tests/ocr_tests.rs @@ -0,0 +1,127 @@ +#![cfg(all(feature = "ocr-oar", not(target_arch = "wasm32")))] + +use pdf_inspector::vision::{ + ModelStore, OarOcrEngine, OcrEngine, OcrMode, OcrOptions, PageTransform, RenderPixelFormat, + RenderedPage, PP_OCR_V6_SMALL, +}; +#[cfg(feature = "render-pdfium")] +use pdf_inspector::vision::{PdfiumRenderer, RenderError, RenderOptions}; + +const MODEL_DIRECTORY_ENV: &str = "PDF_INSPECTOR_OCR_TEST_MODELS"; +const IMAGE_ENV: &str = "PDF_INSPECTOR_OCR_TEST_IMAGE"; +const EXPECTED_TEXT_ENV: &str = "PDF_INSPECTOR_OCR_TEST_EXPECTED"; + +#[cfg(feature = "render-pdfium")] +fn load_renderer() -> Option { + match PdfiumRenderer::load() { + Ok(renderer) => Some(renderer), + Err(RenderError::Pdfium(firecrawl_pdfium::Error::Load( + firecrawl_pdfium::LoadError::LibraryNotFound { .. }, + ))) => { + eprintln!("skipping OCR runtime test because no native PDFium library is installed"); + None + } + Err(error) => panic!("failed to load PDFium: {error}"), + } +} + +#[test] +fn recognizes_an_rgb_image_with_verified_models() { + let Some(model_directory) = std::env::var_os(MODEL_DIRECTORY_ENV) else { + eprintln!("skipping OCR runtime test because {MODEL_DIRECTORY_ENV} is not set"); + return; + }; + let Some(image_path) = std::env::var_os(IMAGE_ENV) else { + eprintln!("skipping OCR runtime test because {IMAGE_ENV} is not set"); + return; + }; + + let image = image::open(image_path).unwrap().into_rgb8(); + let (width, height) = image.dimensions(); + let transform = PageTransform::from_corners( + width, + height, + (0.0, f64::from(height)), + (f64::from(width), f64::from(height)), + (0.0, 0.0), + ) + .unwrap(); + let page = RenderedPage::new( + 1, + width as f32, + height as f32, + width, + height, + width as usize * 3, + RenderPixelFormat::Rgb8, + image.into_raw(), + transform, + ) + .unwrap(); + let results = recognize(&model_directory, &[page]); + assert_usable_result(&results); + + let text = results[0] + .spans + .iter() + .map(|span| span.text.as_str()) + .collect::>() + .join(" "); + eprintln!("recognized: {text}"); + if let Ok(expected) = std::env::var(EXPECTED_TEXT_ENV) { + assert!( + text.to_lowercase().contains(&expected.to_lowercase()), + "expected OCR output to contain {expected:?}, got {text:?}" + ); + } +} + +#[cfg(feature = "render-pdfium")] +#[test] +fn recognizes_a_pdfium_rendered_fixture_with_verified_models() { + let Some(model_directory) = std::env::var_os(MODEL_DIRECTORY_ENV) else { + eprintln!("skipping OCR runtime test because {MODEL_DIRECTORY_ENV} is not set"); + return; + }; + let Some(renderer) = load_renderer() else { + return; + }; + + let bytes = std::fs::read("tests/fixtures/thermo-freon12.pdf").unwrap(); + let pages = renderer + .render_pages( + &bytes, + &[1], + None, + &RenderOptions::new().dpi(150.0).form_fields(false), + ) + .unwrap(); + let results = recognize(&model_directory, &pages); + assert_usable_result(&results); +} + +fn recognize( + model_directory: &std::ffi::OsStr, + pages: &[RenderedPage], +) -> Vec { + let store = ModelStore::new(model_directory).override_root(model_directory); + let models = store.resolve(&PP_OCR_V6_SMALL).unwrap(); + let engine = OarOcrEngine::from_models(&models).unwrap(); + engine + .recognize( + pages, + &OcrOptions::new() + .mode(OcrMode::Force) + .minimum_confidence(0.3), + ) + .unwrap() +} + +fn assert_usable_result(results: &[pdf_inspector::vision::OcrPage]) { + assert_eq!(results.len(), 1); + assert_eq!(results[0].page, 1); + assert_eq!(results[0].model.name, PP_OCR_V6_SMALL.id); + assert_eq!(results[0].model.revision, PP_OCR_V6_SMALL.revision); + assert!(!results[0].spans.is_empty()); + assert!(results[0].spans.iter().all(|span| span.confidence >= 0.3)); +}