From 2cebb3c95ff0b068d2cdebcf40661eb34995e597 Mon Sep 17 00:00:00 2001 From: Abimael Martell <1450169+abimaelmartell@users.noreply.github.com> Date: Sun, 16 Aug 2026 23:55:53 -0700 Subject: [PATCH] feat(vision): expose OCR pipeline (#360) * feat(vision): expose OCR pipeline * fix(vision): harden OCR API * refactor(vision): expose OCR API --- Cargo.toml | 3 + docs/rust-api.md | 56 ++++++ src/bin/pdf2md.rs | 283 ++++++++++++++++++++++++++- src/lib.rs | 172 +++++++++++++++-- src/markdown/mod.rs | 8 + src/vision/fusion.rs | 5 + src/vision/mod.rs | 6 + src/vision/pipeline.rs | 426 +++++++++++++++++++++++++++++++++++++++++ tests/ocr_tests.rs | 32 ++++ 9 files changed, 975 insertions(+), 16 deletions(-) create mode 100644 src/vision/pipeline.rs diff --git a/Cargo.toml b/Cargo.toml index 415f235..326b33b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -86,6 +86,9 @@ model-cache = ["vision", "dep:dirs", "dep:fs2", "dep:sha2", "dep:windows-sys"] model-download = ["model-cache", "dep:ureq"] ocr-oar = ["model-cache", "dep:image", "dep:oar-ocr", "dep:ort"] render-pdfium = ["vision", "dep:firecrawl-pdfium"] +# Complete native OCR path. This remains opt-in so default library, +# renderer-only, and browser consumers do not inherit inference or HTTP/TLS. +ocr = ["render-pdfium", "ocr-oar", "model-download"] [[bin]] name = "pdf2md" diff --git a/docs/rust-api.md b/docs/rust-api.md index 6def665..73ef49f 100644 --- a/docs/rust-api.md +++ b/docs/rust-api.md @@ -340,6 +340,62 @@ warnings. A page that required OCR recommends the hosted pipeline when local OCR is missing, empty, or below the configurable page-confidence threshold. This keeps the lightweight path explicit about cases it cannot finish well. +### Complete OCR API + +The `ocr` convenience feature enables the renderer, OCR engine, verified +model acquisition, routing, and fusion layers together. It is the intended +downstream application integration boundary; lower-level features remain +available for consumers that bring their own renderer, model package manager, +or engine. + +```toml +[dependencies] +pdf-inspector = { version = "1", features = ["ocr"] } +``` + +```rust +use pdf_inspector::vision::{ + process_pdf_with_ocr, OcrMode, OcrPdfOptions, +}; + +let result = process_pdf_with_ocr( + "document.pdf", + OcrPdfOptions::new() + .mode(OcrMode::Auto) + .pages([1, 2, 3]), +)?; + +println!("{}", result.markdown); +println!("OCR pages: {:?}", result.pages_routed_to_ocr); +println!( + "Hosted fallback pages: {:?}", + result.pages_recommending_hosted, +); +``` + +Native extraction always runs first. In `Auto`, a clean PDF returns before +PDFium loading, model-cache access, HTTP, or OAR initialization. Model files +remain external and the default crate feature set remains unchanged. `Off` +provides the same native-only behavior through the OCR result/provenance +shape; `Force` renders every selected page. Learned layout intentionally +returns an explicit unsupported error in this lightweight pipeline. + +Build the CLI with the same opt-in feature: + +```bash +cargo build --release --features ocr --bin pdf2md +pdf2md document.pdf --ocr auto --raw +pdf2md document.pdf --ocr auto --json +pdf2md document.pdf --ocr auto --ocr-offline --ocr-model-dir /opt/models/pp-ocrv6-small +``` + +CLI controls include `--ocr-dpi`, `--ocr-min-confidence`, +`--ocr-hosted-threshold`, `--select-pages`, and the existing encrypted-PDF +`--password` option. JSON output includes per-page Markdown, source/model +provenance, confidence, timings, warnings, routed pages, and hosted-fallback +recommendations. Page numbers in `OcrPdfResult` and its per-page provenance +are 1-indexed, matching the PDF page numbers accepted by `OcrPdfOptions::pages`. + Extract per-page Markdown (one string per page, plus document-wide layout metadata): diff --git a/src/bin/pdf2md.rs b/src/bin/pdf2md.rs index 1d5e82a..7599e1f 100644 --- a/src/bin/pdf2md.rs +++ b/src/bin/pdf2md.rs @@ -1,6 +1,11 @@ //! CLI tool for PDF to Markdown conversion use pdf_inspector::extractor::ItemType; +#[cfg(all(feature = "ocr", not(target_arch = "wasm32")))] +use pdf_inspector::vision::{ + process_pdf_with_ocr, ModelDownloadPolicy, OcrMode, OcrOptions, OcrPdfOptions, OcrPdfResult, + PageContentSource, RenderOptions, +}; use pdf_inspector::{ extract_text_with_positions_pages_with_password, process_pdf_with_options, LayoutComplexity, PdfOptions, PdfType, ProcessMode, TextItem, @@ -103,6 +108,134 @@ fn format_items_json(items: &[TextItem]) -> String { ) } +#[cfg(all(feature = "ocr", not(target_arch = "wasm32")))] +fn optional_json_number(value: Option) -> String { + value + .filter(|value| value.is_finite()) + .map(|value| format!("{value:.4}")) + .unwrap_or_else(|| "null".to_string()) +} + +#[cfg(all(feature = "ocr", not(target_arch = "wasm32")))] +fn format_ocr_json(result: &OcrPdfResult) -> String { + let routed = result + .pages_routed_to_ocr + .iter() + .map(u32::to_string) + .collect::>() + .join(","); + let recommended = result + .pages_recommended_for_ocr + .iter() + .map(u32::to_string) + .collect::>() + .join(","); + let hosted = result + .pages_recommending_hosted + .iter() + .map(u32::to_string) + .collect::>() + .join(","); + let pages = result + .pages + .iter() + .map(|page| { + let provenance = &page.provenance; + let source = match provenance.source { + PageContentSource::Native => "native", + PageContentSource::Ocr => "ocr", + PageContentSource::Fused => "fused", + _ => "unknown", + }; + let model = provenance + .ocr_model + .as_ref() + .map(|model| { + format!( + r#"{{"name":"{}","revision":"{}"}}"#, + json_escape(&model.name), + json_escape(&model.revision) + ) + }) + .unwrap_or_else(|| "null".to_string()); + let warnings = provenance + .warnings + .iter() + .map(|warning| format!(r#""{}""#, json_escape(warning))) + .collect::>() + .join(","); + format!( + r#"{{"page":{},"source":"{}","markdown":"{}","ocr_model":{},"render_dpi":{},"ocr_confidence":{},"hosted_recommended":{},"timings":{{"render_ms":{},"ocr_ms":{},"layout_ms":{},"assembly_ms":{}}},"warnings":[{}]}}"#, + provenance.page, + source, + json_escape(&page.markdown), + model, + optional_json_number(provenance.render_dpi), + optional_json_number(provenance.ocr_confidence), + provenance.hosted_recommended, + provenance.timings.render_ms, + provenance.timings.ocr_ms, + provenance.timings.layout_ms, + provenance.timings.assembly_ms, + warnings, + ) + }) + .collect::>() + .join(","); + let table_pages = result + .pages_with_tables + .iter() + .map(u32::to_string) + .collect::>() + .join(","); + let column_pages = result + .pages_with_columns + .iter() + .map(u32::to_string) + .collect::>() + .join(","); + let ocr_reasons = format_ocr_reasons_by_page(&result.ocr_reasons_by_page); + format!( + r#"{{"page_count":{},"processing_time_ms":{},"render_time_ms":{},"ocr_time_ms":{},"pages_recommended_for_ocr":[{}],"pages_routed_to_ocr":[{}],"pages_recommending_hosted":[{}],"ocr_reasons_by_page":[{}],"is_complex":{},"pages_with_tables":[{}],"pages_with_columns":[{}],"pages":[{}],"markdown":"{}"}}"#, + result.page_count, + result.processing_time_ms, + result.render_time_ms, + result.ocr_time_ms, + recommended, + routed, + hosted, + ocr_reasons, + result.is_complex, + table_pages, + column_pages, + pages, + json_escape(&result.markdown), + ) +} + +fn argument_value<'a>(args: &'a [String], name: &str) -> Result, String> { + args.iter() + .position(|argument| argument == name) + .map(|index| { + args.get(index + 1) + .map(String::as_str) + .ok_or_else(|| format!("{name} requires a value")) + }) + .transpose() +} + +#[cfg(all(feature = "ocr", not(target_arch = "wasm32")))] +fn float_argument(args: &[String], name: &str, default: f32) -> Result { + argument_value(args, name)? + .map(|value| { + value + .parse::() + .map_err(|_| format!("{name} requires a number, got {value:?}")) + }) + .transpose() + .map(|value| value.unwrap_or(default)) +} + fn extract_items_json( pdf_path: &str, page_filter: Option<&HashSet>, @@ -242,6 +375,12 @@ fn main() { eprintln!(" --password PW Password for an encrypted PDF"); eprintln!(" --detect-only Only detect PDF type (no extraction)"); eprintln!(" --analyze Detect + extract + layout analysis (no markdown)"); + eprintln!(" --ocr MODE OCR mode: off, auto, or force (requires feature `ocr`)"); + eprintln!(" --ocr-dpi N OCR render resolution (default: 150)"); + eprintln!(" --ocr-min-confidence N Drop OCR spans below N (default: 0)"); + eprintln!(" --ocr-hosted-threshold N Recommend hosted parsing below N (default: 0.5)"); + eprintln!(" --ocr-model-dir DIR Use a package-managed local model directory"); + eprintln!(" --ocr-offline Never download missing OCR models"); process::exit(1); } @@ -253,6 +392,10 @@ fn main() { let page_numbers = args.iter().any(|a| a == "--pages"); let detect_only = args.iter().any(|a| a == "--detect-only"); let analyze = args.iter().any(|a| a == "--analyze"); + let ocr_mode_argument = argument_value(&args, "--ocr").unwrap_or_else(|error| { + eprintln!("Error: {error}"); + process::exit(1); + }); // Parse --password value let password = args.iter().position(|a| a == "--password").map(|i| { @@ -283,6 +426,141 @@ fn main() { }) }); + let output_file = args + .get(2) + .filter(|a| !a.starts_with("--")) + .map(|s| s.as_str()); + + let has_ocr_only_option = [ + "--ocr-dpi", + "--ocr-min-confidence", + "--ocr-hosted-threshold", + "--ocr-model-dir", + "--ocr-offline", + ] + .iter() + .any(|option| args.iter().any(|argument| argument == option)); + if ocr_mode_argument.is_none() && has_ocr_only_option { + eprintln!("Error: OCR options require --ocr off, --ocr auto, or --ocr force"); + process::exit(1); + } + + if let Some(mode) = ocr_mode_argument { + if items_json_output || detect_only || analyze { + eprintln!( + "Error: --ocr cannot be combined with --items-json, --detect-only, or --analyze" + ); + process::exit(1); + } + + #[cfg(not(all(feature = "ocr", not(target_arch = "wasm32"))))] + { + let _ = mode; + eprintln!("Error: this pdf2md build does not include OCR; rebuild with --features ocr"); + process::exit(1); + } + + #[cfg(all(feature = "ocr", not(target_arch = "wasm32")))] + { + let mode = match mode { + "off" => OcrMode::Off, + "auto" => OcrMode::Auto, + "force" => OcrMode::Force, + value => { + eprintln!("Error: invalid --ocr mode {value:?}; expected off, auto, or force"); + process::exit(1); + } + }; + let dpi = float_argument(&args, "--ocr-dpi", 150.0).unwrap_or_else(|error| { + eprintln!("Error: {error}"); + process::exit(1); + }); + let minimum_confidence = float_argument(&args, "--ocr-min-confidence", 0.0) + .unwrap_or_else(|error| { + eprintln!("Error: {error}"); + process::exit(1); + }); + let hosted_threshold = float_argument(&args, "--ocr-hosted-threshold", 0.5) + .unwrap_or_else(|error| { + eprintln!("Error: {error}"); + process::exit(1); + }); + let model_directory = + argument_value(&args, "--ocr-model-dir").unwrap_or_else(|error| { + eprintln!("Error: {error}"); + process::exit(1); + }); + + let mut ocr = OcrOptions::new() + .mode(mode) + .minimum_confidence(minimum_confidence); + if let Some(directory) = model_directory { + ocr = ocr.model_directory(directory); + } + if args.iter().any(|argument| argument == "--ocr-offline") { + ocr = ocr.model_downloads(ModelDownloadPolicy::Offline); + } + let mut markdown = pdf_inspector::MarkdownOptions::default(); + if compact_output { + markdown.profile = pdf_inspector::MarkdownProfile::Compact; + } + markdown.include_page_numbers = page_numbers; + let mut pdf_options = OcrPdfOptions::new() + .render(RenderOptions::new().dpi(dpi)) + .ocr(ocr) + .markdown(markdown) + .hosted_recommendation_confidence(hosted_threshold); + if let Some(pages) = page_filter.clone() { + pdf_options = pdf_options.pages(pages); + } + if let Some(password) = password.clone() { + pdf_options = pdf_options.password(password); + } + + match process_pdf_with_ocr(pdf_path, pdf_options) { + Ok(result) => { + if json_output { + println!("{}", format_ocr_json(&result)); + } else if raw_output { + print!("{}", result.markdown); + } else { + eprintln!("PDF to Markdown Conversion (OCR)"); + eprintln!("======================================"); + eprintln!("File: {pdf_path}"); + eprintln!("Pages: {}", result.page_count); + eprintln!("Pages routed to OCR: {:?}", result.pages_routed_to_ocr); + if !result.pages_recommending_hosted.is_empty() { + eprintln!( + "Hosted parsing recommended for pages: {:?}", + result.pages_recommending_hosted + ); + } + eprintln!("Processing time: {}ms", result.processing_time_ms); + if let Some(output) = output_file { + fs::write(output, &result.markdown) + .expect("Failed to write output file"); + eprintln!("Markdown written to: {output}"); + } else { + eprintln!(); + eprintln!("--- Markdown Output ---"); + eprintln!(); + print!("{}", result.markdown); + } + } + } + Err(error) => { + if json_output { + println!(r#"{{"error":"{}"}}"#, json_escape(&error.to_string())); + } else { + eprintln!("Error: {error}"); + } + process::exit(1); + } + } + return; + } + } + if items_json_output { match extract_items_json(pdf_path, page_filter.as_ref(), password.as_deref()) { Ok(json) => println!("{}", json), @@ -294,11 +572,6 @@ fn main() { return; } - let output_file = args - .get(2) - .filter(|a| !a.starts_with("--")) - .map(|s| s.as_str()); - let process_mode = if detect_only { ProcessMode::DetectOnly } else if analyze { diff --git a/src/lib.rs b/src/lib.rs index 70d875d..215798d 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -459,8 +459,35 @@ pub fn extract_pages_markdown_mem( buffer: &[u8], pages: Option<&[u32]>, ) -> Result { + extract_pages_markdown_mem_impl(buffer, pages, None, &MarkdownOptions::default(), false) + .map(|(result, _)| result) +} + +#[cfg(all(feature = "ocr", not(target_arch = "wasm32")))] +pub(crate) fn extract_pages_markdown_mem_for_ocr( + buffer: &[u8], + pages: Option<&[u32]>, + password: Option<&str>, + markdown_options: &MarkdownOptions, +) -> Result<(PagesExtractionResult, u32), PdfError> { + extract_pages_markdown_mem_impl( + buffer, + pages, + password, + markdown_options, + markdown_options.strip_headers_footers, + ) +} + +fn extract_pages_markdown_mem_impl( + buffer: &[u8], + pages: Option<&[u32]>, + password: Option<&str>, + markdown_options: &MarkdownOptions, + strip_repeated_headers_footers: bool, +) -> Result<(PagesExtractionResult, u32), PdfError> { validate_pdf_bytes(buffer)?; - let (doc, page_count) = load_document_from_mem(buffer)?; + let (doc, page_count) = load_document_from_mem_with_password(buffer, password)?; let font_cmaps = FontCMaps::from_doc(&doc); // Extract ALL pages to get accurate, document-wide font stats. A malformed @@ -503,6 +530,11 @@ pub fn extract_pages_markdown_mem( // Compute font stats from full document (cross-page consistency). let font_stats = markdown::analysis::calculate_font_stats_from_items(&filtered_items); + let repeated_header_footer_items = if strip_repeated_headers_footers { + repeated_header_footer_item_keys(&all_items, &page_thresholds, &chart_regions, page_count) + } else { + HashSet::new() + }; // When caller doesn't specify pages, return every page in document order. let all_pages: Vec; @@ -538,7 +570,10 @@ pub fn extract_pages_markdown_mem( let (page_items, page_number_removal_mask): (Vec, Vec) = all_items .iter() .zip(&page_number_removal_mask) - .filter(|(item, _)| item.page == page_1idx) + .filter(|(item, _)| { + item.page == page_1idx + && !repeated_header_footer_items.contains(&HeaderFooterItemKey::from(*item)) + }) .map(|(item, remove)| (item.clone(), *remove)) .unzip(); @@ -575,7 +610,7 @@ pub fn extract_pages_markdown_mem( base_font_size: Some(font_stats.most_common_size), include_page_numbers: false, strip_headers_footers: false, - ..MarkdownOptions::default() + ..markdown_options.clone() }; let md = if has_text_quality_issue { @@ -634,14 +669,129 @@ pub fn extract_pages_markdown_mem( }); } - Ok(PagesExtractionResult { - pages: results, - pages_with_tables: complexity.pages_with_tables, - pages_with_columns: complexity.pages_with_columns, - pages_needing_ocr, - ocr_reasons_by_page: page_ocr_reasons_vec(ocr_reasons_by_page), - is_complex: complexity.is_complex, - }) + Ok(( + PagesExtractionResult { + pages: results, + pages_with_tables: complexity.pages_with_tables, + pages_with_columns: complexity.pages_with_columns, + pages_needing_ocr, + ocr_reasons_by_page: page_ocr_reasons_vec(ocr_reasons_by_page), + is_complex: complexity.is_complex, + }, + page_count, + )) +} + +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +struct HeaderFooterItemKey { + page: u32, + x: u32, + y: u32, + text: String, +} + +impl From<&TextItem> for HeaderFooterItemKey { + fn from(item: &TextItem) -> Self { + Self { + page: item.page, + x: item.x.to_bits(), + y: item.y.to_bits(), + text: item.text.clone(), + } + } +} + +fn repeated_header_footer_item_keys( + items: &[TextItem], + page_thresholds: &HashMap, + chart_regions: &HashMap>, + page_count: u32, +) -> HashSet { + let candidates = items + .iter() + .filter(|item| { + matches!( + item.item_type, + types::ItemType::Text | types::ItemType::FormField + ) + }) + .cloned() + .collect(); + let lines = extractor::group_prefiltered_items_into_lines_with_thresholds_and_charts( + candidates, + page_thresholds, + &HashSet::new(), + chart_regions, + ); + let all_items: HashSet<_> = lines + .iter() + .flat_map(|line| line.items.iter().map(HeaderFooterItemKey::from)) + .collect(); + let kept = markdown::strip_repeated_header_footer_lines(lines, page_count); + let kept_items: HashSet<_> = kept + .iter() + .flat_map(|line| line.items.iter().map(HeaderFooterItemKey::from)) + .collect(); + all_items.difference(&kept_items).cloned().collect() +} + +#[cfg(all(test, feature = "ocr", not(target_arch = "wasm32")))] +mod ocr_header_footer_tests { + use super::*; + + fn item(page: u32, text: &str, y: f32) -> TextItem { + TextItem { + text: text.to_string(), + x: 10.0, + y, + width: 120.0, + height: 10.0, + font: "Test".to_string(), + font_size: 10.0, + page, + is_bold: false, + is_italic: false, + is_underline: false, + is_strikeout: false, + item_type: types::ItemType::Text, + mcid: None, + } + } + + #[test] + fn local_pipeline_prefilters_document_wide_repeated_headers() { + let mut items = Vec::new(); + let mut thresholds = HashMap::new(); + for page in 1..=3 { + items.push(item(page, "Repeated report header", 800.0)); + for line in 0..12 { + items.push(item( + page, + &format!("Page {page} paragraph {line} unique content"), + 700.0 - line as f32 * 40.0, + )); + } + thresholds.insert(page, 0.1); + } + + let removed = repeated_header_footer_item_keys(&items, &thresholds, &HashMap::new(), 3); + assert_eq!(removed.len(), 2); + for page in 1..=3 { + assert_eq!( + removed.contains(&HeaderFooterItemKey::from(&item( + page, + "Repeated report header", + 800.0, + ))), + page > 1, + ); + assert!(!removed.contains(&HeaderFooterItemKey::from(&item( + page, + &format!("Page {page} paragraph 5 unique content"), + 500.0, + )))); + } + } } /// Path-based wrapper for [`extract_pages_markdown_mem`]. diff --git a/src/markdown/mod.rs b/src/markdown/mod.rs index ba21108..acba84d 100644 --- a/src/markdown/mod.rs +++ b/src/markdown/mod.rs @@ -1144,6 +1144,14 @@ pub fn to_markdown(text: &str, options: MarkdownOptions) -> String { output } +/// Applies the document-wide repeated header/footer classifier to grouped lines. +pub(crate) fn strip_repeated_header_footer_lines( + lines: Vec, + page_count: u32, +) -> Vec { + preprocess::strip_repeated_lines(lines, page_count) +} + /// Convert positioned text items to markdown with structure detection pub fn to_markdown_from_items(items: Vec, options: MarkdownOptions) -> String { to_markdown_from_items_with_rects(items, options, &[]) diff --git a/src/vision/fusion.rs b/src/vision/fusion.rs index a3ae58b..f0213ec 100644 --- a/src/vision/fusion.rs +++ b/src/vision/fusion.rs @@ -62,6 +62,11 @@ impl OcrFusionOptions { self.hosted_recommendation_confidence = confidence; self } + + /// Validates rendering and hosted-fallback thresholds without doing work. + pub fn validate(&self) -> Result<(), OcrFusionError> { + validate_options(self) + } } /// Final Markdown and provenance for one page. diff --git a/src/vision/mod.rs b/src/vision/mod.rs index 5e3ca2a..7657f05 100644 --- a/src/vision/mod.rs +++ b/src/vision/mod.rs @@ -18,6 +18,8 @@ mod fusion; mod models; #[cfg(all(feature = "ocr-oar", not(target_arch = "wasm32")))] mod oar; +#[cfg(all(feature = "ocr", not(target_arch = "wasm32")))] +mod pipeline; #[cfg(all(feature = "vision", not(target_arch = "wasm32")))] mod render; #[cfg(all(feature = "vision", not(target_arch = "wasm32")))] @@ -46,6 +48,10 @@ pub use models::{ }; #[cfg(all(feature = "ocr-oar", not(target_arch = "wasm32")))] pub use oar::{OarOcrEngine, OarOcrError, ONNX_RUNTIME_LIBRARY_ENV}; +#[cfg(all(feature = "ocr", not(target_arch = "wasm32")))] +pub use pipeline::{ + process_pdf_with_ocr, process_pdf_with_ocr_mem, OcrPdfOptions, OcrPdfResult, OcrPipelineError, +}; #[cfg(all(feature = "vision", not(target_arch = "wasm32")))] pub use render::{ PagePoint, PageTransform, RenderBufferError, RenderOptions, RenderPixelFormat, RenderedPage, diff --git a/src/vision/pipeline.rs b/src/vision/pipeline.rs new file mode 100644 index 0000000..a00d5c2 --- /dev/null +++ b/src/vision/pipeline.rs @@ -0,0 +1,426 @@ +//! One-call native extraction and OCR pipeline. + +use std::collections::BTreeSet; +use std::path::Path; +use std::time::Instant; + +use thiserror::Error; + +use crate::{MarkdownOptions, PageOcrReasons, PdfError}; + +use super::{ + fuse_ocr_pages, route_ocr_pages, run_ocr_pages, FusedPageMarkdown, HttpModelDownloadError, + HttpModelDownloader, ModelAcquireError, ModelStore, ModelStoreError, OarOcrEngine, OarOcrError, + OcrFusionError, OcrFusionOptions, OcrMode, OcrOptions, OcrRoutingError, OcrRun, OcrRunError, + PdfiumRenderer, RenderError, RenderOptions, PP_OCR_V6_SMALL, +}; + +/// Options for native extraction with optional OCR. +#[derive(Clone)] +pub struct OcrPdfOptions { + /// Page rasterization settings used when OCR is routed. + pub render: RenderOptions, + /// OCR routing, model, and recognition settings. + pub ocr: OcrOptions, + /// Markdown formatting shared by native and OCR assembly. + pub markdown: MarkdownOptions, + /// Optional 1-indexed page selection. `None` processes the full document. + pub page_filter: Option>, + /// Password for an encrypted PDF. + pub password: Option, + /// Weak OCR threshold for recommending the hosted pipeline. + pub hosted_recommendation_confidence: f32, +} + +impl Default for OcrPdfOptions { + fn default() -> Self { + Self { + render: RenderOptions::default(), + ocr: OcrOptions::default(), + markdown: MarkdownOptions::default(), + page_filter: None, + password: None, + hosted_recommendation_confidence: 0.5, + } + } +} + +impl std::fmt::Debug for OcrPdfOptions { + fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + formatter + .debug_struct("OcrPdfOptions") + .field("render", &self.render) + .field("ocr", &self.ocr) + .field("markdown", &self.markdown) + .field("page_filter", &self.page_filter) + .field("password", &self.password.as_ref().map(|_| "[REDACTED]")) + .field( + "hosted_recommendation_confidence", + &self.hosted_recommendation_confidence, + ) + .finish() + } +} + +impl OcrPdfOptions { + /// Creates options with OCR disabled, preserving the native-only path. + pub fn new() -> Self { + Self::default() + } + + /// Replaces page rasterization settings. + pub fn render(mut self, render: RenderOptions) -> Self { + self.render = render; + self + } + + /// Replaces OCR routing and recognition settings. + pub fn ocr(mut self, ocr: OcrOptions) -> Self { + self.ocr = ocr; + self + } + + /// Sets OCR routing without changing the remaining OCR settings. + pub fn mode(mut self, mode: OcrMode) -> Self { + self.ocr.mode = mode; + self + } + + /// Replaces Markdown formatting options. + pub fn markdown(mut self, markdown: MarkdownOptions) -> Self { + self.markdown = markdown; + self + } + + /// Restricts processing to 1-indexed pages in ascending order. + pub fn pages(mut self, pages: impl IntoIterator) -> Self { + self.page_filter = Some(pages.into_iter().collect()); + self + } + + /// Sets the password used to decrypt the PDF. + pub fn password(mut self, password: impl Into) -> Self { + self.password = Some(password.into()); + self + } + + /// Sets the weak-OCR threshold for recommending hosted document parsing. + pub fn hosted_recommendation_confidence(mut self, confidence: f32) -> Self { + self.hosted_recommendation_confidence = confidence; + self + } +} + +/// Complete native/OCR Markdown output for a PDF request. +#[derive(Debug, Clone, PartialEq)] +pub struct OcrPdfResult { + /// Final document Markdown in selected-page order. + pub markdown: String, + /// Final per-page Markdown and provenance, using 1-indexed page numbers. + pub pages: Vec, + /// Total pages in the PDF, independent of page selection. + pub page_count: u32, + /// 1-indexed selected pages recommended for OCR by native extraction. + pub pages_recommended_for_ocr: Vec, + /// 1-indexed pages actually rendered and recognized. + pub pages_routed_to_ocr: Vec, + /// 1-indexed pages whose OCR result recommends hosted document parsing. + pub pages_recommending_hosted: Vec, + /// Original machine-readable OCR reasons for selected pages. + pub ocr_reasons_by_page: Vec, + /// Selected pages where deterministic table detection found tables. + pub pages_with_tables: Vec, + /// Selected pages where deterministic layout found multiple columns. + pub pages_with_columns: Vec, + /// Whether deterministic extraction found tables or columns. + pub is_complex: bool, + /// End-to-end processing time. + pub processing_time_ms: u64, + /// Batch page-rendering time; zero when no OCR work was routed. + pub render_time_ms: u64, + /// Batch OCR time; zero when no OCR work was routed. + pub ocr_time_ms: u64, +} + +/// Processes a PDF file through native extraction and selective OCR. +pub fn process_pdf_with_ocr( + path: impl AsRef, + options: OcrPdfOptions, +) -> Result { + let bytes = std::fs::read(path).map_err(PdfError::from)?; + process_pdf_with_ocr_mem(&bytes, options) +} + +/// Processes PDF bytes through native extraction and selective OCR. +/// +/// Native extraction always runs first. `Auto` initializes PDFium, downloads +/// models, and starts OAR only if the detector selected at least one page. +/// `Off` therefore has no renderer, model-cache, network, or inference side +/// effects even though the complete feature is compiled into the application. +pub fn process_pdf_with_ocr_mem( + buffer: &[u8], + options: OcrPdfOptions, +) -> Result { + OcrFusionOptions::new() + .render_dpi(options.render.dpi) + .hosted_recommendation_confidence(options.hosted_recommendation_confidence) + .validate()?; + let minimum_confidence = options.ocr.minimum_confidence; + if !minimum_confidence.is_finite() || !(0.0..=1.0).contains(&minimum_confidence) { + return Err(OcrPipelineError::InvalidMinimumConfidence { + value: minimum_confidence, + }); + } + if options + .page_filter + .as_ref() + .is_some_and(|pages| pages.contains(&0)) + { + return Err(OcrPipelineError::InvalidSelectedPage { page: 0 }); + } + + let started = Instant::now(); + let selected_pages: Option> = options + .page_filter + .as_ref() + .map(|pages| pages.iter().copied().collect()); + let selected_pages_zero_indexed: Option> = selected_pages + .as_ref() + .map(|pages| pages.iter().map(|page| page - 1).collect()); + + let mut page_markdown_options = options.markdown.clone(); + page_markdown_options.include_page_numbers = false; + let (native, page_count) = crate::extract_pages_markdown_mem_for_ocr( + buffer, + selected_pages_zero_indexed.as_deref(), + options.password.as_deref(), + &page_markdown_options, + )?; + if let Some(invalid) = selected_pages + .as_ref() + .and_then(|pages| pages.iter().copied().find(|page| *page > page_count)) + { + return Err(OcrPipelineError::InvalidSelectedPage { page: invalid }); + } + + let routed = route_ocr_pages( + options.ocr.mode, + page_count, + &native.pages_needing_ocr, + selected_pages.as_deref(), + )?; + + let ocr_run = if routed.is_empty() { + OcrRun { + pages: Vec::new(), + render_time_ms: 0, + ocr_time_ms: 0, + } + } else { + // Resolve the native renderer before any network request so a missing + // PDFium installation cannot trigger a model download it cannot use. + let renderer = PdfiumRenderer::load()?; + let store = ModelStore::from_options(&options.ocr)?; + let models = store.resolve_or_download( + &PP_OCR_V6_SMALL, + options.ocr.model_downloads, + &HttpModelDownloader::default(), + )?; + let engine = OarOcrEngine::from_models(&models)?; + run_ocr_pages( + &renderer, + &engine, + buffer, + &routed, + options.password.as_deref(), + &options.render, + &options.ocr, + )? + }; + + let fusion_options = OcrFusionOptions::new() + .markdown(page_markdown_options) + .render_dpi(options.render.dpi) + .hosted_recommendation_confidence(options.hosted_recommendation_confidence); + let fused = fuse_ocr_pages(&native.pages, &ocr_run, page_count, &fusion_options)?; + let pages_recommending_hosted = fused + .pages + .iter() + .filter(|page| page.provenance.hosted_recommended) + .map(|page| page.provenance.page) + .collect(); + let markdown = assemble_document_markdown(&fused.pages, options.markdown.include_page_numbers); + + Ok(OcrPdfResult { + markdown, + pages: fused.pages, + page_count, + pages_recommended_for_ocr: native.pages_needing_ocr, + pages_routed_to_ocr: routed, + pages_recommending_hosted, + ocr_reasons_by_page: native.ocr_reasons_by_page, + pages_with_tables: native.pages_with_tables, + pages_with_columns: native.pages_with_columns, + is_complex: native.is_complex, + processing_time_ms: elapsed_ms(started), + render_time_ms: fused.render_time_ms, + ocr_time_ms: fused.ocr_time_ms, + }) +} + +fn assemble_document_markdown(pages: &[FusedPageMarkdown], include_page_numbers: bool) -> String { + let mut document = String::new(); + for (index, page) in pages.iter().enumerate() { + if index > 0 { + document.push_str("\n\n"); + } + if include_page_numbers { + document.push_str(&format!("\n\n", page.page)); + } + document.push_str(page.markdown.trim()); + } + if !document.is_empty() { + document.push('\n'); + } + document +} + +fn elapsed_ms(started: Instant) -> u64 { + u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX) +} + +/// Failures from the complete OCR pipeline. +#[derive(Debug, Error)] +#[non_exhaustive] +pub enum OcrPipelineError { + /// PDF loading or native extraction failed. + #[error(transparent)] + Pdf(#[from] PdfError), + /// Page routing rejected an invalid request. + #[error(transparent)] + Routing(#[from] OcrRoutingError), + /// The model cache could not be located or initialized. + #[error(transparent)] + ModelStore(#[from] ModelStoreError), + /// A pinned model set could not be resolved or acquired. + #[error(transparent)] + ModelAcquire(#[from] ModelAcquireError), + /// PDFium could not load or rasterize the request. + #[error(transparent)] + Render(#[from] RenderError), + /// The OAR engine could not initialize. + #[error(transparent)] + Oar(#[from] OarOcrError), + /// Selective rendering or OCR execution failed. + #[error(transparent)] + Run(#[from] OcrRunError), + /// OCR/native Markdown fusion failed. + #[error(transparent)] + Fusion(#[from] OcrFusionError), + /// Page zero is invalid because public page selections are 1-indexed. + #[error("selected page {page} is invalid; page numbers are 1-indexed")] + InvalidSelectedPage { + /// Invalid page number. + page: u32, + }, + /// OCR span confidence is outside the inclusive 0–1 range. + #[error("minimum OCR confidence must be between 0 and 1, got {value}")] + InvalidMinimumConfidence { + /// Invalid value. + value: f32, + }, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn off_mode_extracts_native_text_without_runtime_side_effects() { + let bytes = std::fs::read("tests/fixtures/thermo-freon12.pdf").unwrap(); + let result = process_pdf_with_ocr_mem(&bytes, OcrPdfOptions::new()).unwrap(); + + assert_eq!(result.page_count, 3); + assert_eq!(result.pages.len(), 3); + assert!(result.markdown.contains("Thermodynamic Properties")); + assert!(result.pages_recommended_for_ocr.is_empty()); + assert!(result.pages_routed_to_ocr.is_empty()); + assert!(result.pages_recommending_hosted.is_empty()); + assert_eq!(result.render_time_ms, 0); + assert_eq!(result.ocr_time_ms, 0); + } + + #[test] + fn auto_mode_does_not_load_pdfium_or_models_for_clean_pdf() { + let bytes = std::fs::read("tests/fixtures/thermo-freon12.pdf").unwrap(); + let result = + process_pdf_with_ocr_mem(&bytes, OcrPdfOptions::new().mode(OcrMode::Auto)).unwrap(); + + assert!(result.pages_routed_to_ocr.is_empty()); + assert!(result.markdown.contains("Freon 12")); + } + + #[test] + fn selection_and_page_markers_use_public_one_indexed_pages() { + let bytes = std::fs::read("tests/fixtures/thermo-freon12.pdf").unwrap(); + let mut markdown = MarkdownOptions::default(); + markdown.include_page_numbers = true; + let result = + process_pdf_with_ocr_mem(&bytes, OcrPdfOptions::new().pages([2]).markdown(markdown)) + .unwrap(); + + assert_eq!(result.pages.len(), 1); + assert_eq!(result.pages[0].page, 2); + assert_eq!(result.pages[0].page, result.pages[0].provenance.page); + assert!(result.markdown.starts_with("")); + } + + #[test] + fn off_mode_marks_unprocessed_scan_for_hosted_fallback() { + let bytes = std::fs::read("tests/fixtures/scan_with_native_header_text.pdf").unwrap(); + let result = process_pdf_with_ocr_mem(&bytes, OcrPdfOptions::new()).unwrap(); + + assert!(result.pages_routed_to_ocr.is_empty()); + assert_eq!(result.pages_recommending_hosted, vec![1]); + } + + #[test] + fn password_is_redacted_and_used_for_native_extraction() { + let options = OcrPdfOptions::new().password("secret123"); + assert!(!format!("{options:?}").contains("secret123")); + + let bytes = std::fs::read("tests/fixtures/encrypted-secret123.pdf").unwrap(); + let result = process_pdf_with_ocr_mem(&bytes, options).unwrap(); + assert!(result.markdown.contains("Procurement")); + } + + #[test] + fn rejects_out_of_range_selection_even_with_ocr_off() { + let bytes = std::fs::read("tests/fixtures/thermo-freon12.pdf").unwrap(); + let error = process_pdf_with_ocr_mem(&bytes, OcrPdfOptions::new().pages([4])).unwrap_err(); + assert!(matches!( + error, + OcrPipelineError::InvalidSelectedPage { page: 4 } + )); + } + + #[test] + fn invalid_expensive_options_fail_before_pdf_or_runtime_access() { + let mut invalid_dpi = OcrPdfOptions::new(); + invalid_dpi.render.dpi = f32::NAN; + assert!(matches!( + process_pdf_with_ocr_mem(b"not a PDF", invalid_dpi), + Err(OcrPipelineError::Fusion( + OcrFusionError::InvalidRenderDpi { .. } + )) + )); + + let invalid_hosted = OcrPdfOptions::new().hosted_recommendation_confidence(1.1); + assert!(matches!( + process_pdf_with_ocr_mem(b"not a PDF", invalid_hosted), + Err(OcrPipelineError::Fusion( + OcrFusionError::InvalidHostedConfidence { .. } + )) + )); + } +} diff --git a/tests/ocr_tests.rs b/tests/ocr_tests.rs index 486a034..10c2e59 100644 --- a/tests/ocr_tests.rs +++ b/tests/ocr_tests.rs @@ -1,5 +1,9 @@ #![cfg(all(feature = "ocr-oar", not(target_arch = "wasm32")))] +#[cfg(feature = "ocr")] +use pdf_inspector::vision::{ + process_pdf_with_ocr_mem, ModelDownloadPolicy, OcrPdfOptions, PageContentSource, +}; use pdf_inspector::vision::{ ModelStore, OarOcrEngine, OcrEngine, OcrMode, OcrOptions, PageTransform, RenderPixelFormat, RenderedPage, PP_OCR_V6_SMALL, @@ -100,6 +104,34 @@ fn recognizes_a_pdfium_rendered_fixture_with_verified_models() { assert_usable_result(&results); } +#[cfg(all(feature = "ocr", feature = "render-pdfium"))] +#[test] +fn complete_ocr_pipeline_routes_and_assembles_a_scanned_fixture() { + 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/scan_with_native_header_text.pdf").unwrap(); + let ocr = OcrOptions::new() + .mode(OcrMode::Auto) + .minimum_confidence(0.3) + .model_directory(model_directory) + .model_downloads(ModelDownloadPolicy::Offline); + let result = process_pdf_with_ocr_mem(&bytes, OcrPdfOptions::new().ocr(ocr)).unwrap(); + + assert_eq!(result.pages_routed_to_ocr, vec![1]); + assert!(!result.markdown.trim().is_empty()); + assert_eq!(result.pages[0].provenance.source, PageContentSource::Ocr); + assert_eq!( + result.pages[0].provenance.ocr_model.as_ref().unwrap().name, + PP_OCR_V6_SMALL.id + ); +} + fn recognize( model_directory: &std::ffi::OsStr, pages: &[RenderedPage],