Compare commits

...
14 changed files with 3068 additions and 90 deletions
+7
View File
@@ -62,6 +62,9 @@ sha2 = { version = "0.11", optional = true }
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 }
# HTTPS-only streaming downloader for pinned model artifacts. Kept separate
# from model-cache so offline and package-managed deployments avoid HTTP/TLS.
ureq = { version = "3.4", default-features = false, features = ["rustls", "platform-verifier"], optional = true }
[target.'cfg(all(windows, not(target_arch = "wasm32")))'.dependencies]
windows-sys = { version = "0.61", features = ["Win32_Storage_FileSystem"], optional = true }
@@ -80,8 +83,12 @@ default = []
python = ["pyo3"]
vision = []
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"
+169 -7
View File
@@ -137,11 +137,12 @@ pdf-inspector = { version = "1", features = ["vision", "model-cache"] }
The OCR contracts preserve existing behavior by default: OCR is `Off`, learned
layout is disabled, and model resolution is never reached. `ModelStore` itself
does not access the network; a runtime integration can fetch a manifest's
canonical URL only when allowed and pass the stream to `ModelStore::install`.
Offline consumers set an explicit model directory and `ModelDownloadPolicy::Offline`.
Renderer-only consumers do not enable `model-cache` and therefore do not compile
its filesystem, locking, or hashing dependencies.
does not access the network. The optional `model-download` feature provides an
HTTPS downloader that streams pinned artifacts into the checksum-verified
cache only after routing has selected OCR work. Offline consumers set an
explicit model directory and `ModelDownloadPolicy::Offline`. Renderer-only
consumers do not enable `model-cache` or `model-download` and therefore do not
compile their filesystem, hashing, or HTTP dependencies.
```rust
use pdf_inspector::vision::{
@@ -244,8 +245,169 @@ 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.
inference accidentally.
### Selective routing and lazy model acquisition
`route_ocr_pages` applies the existing detector/text-quality recommendations to
the configured mode. `Auto` processes only recommended pages, `Force` processes
all pages (or an explicit page selection), and `Off` always returns an empty
route. `run_ocr_pages` renders only that route, checks that both dependencies
preserve its order, and retains each bitmap's PDF transform for fusion.
```toml
[dependencies]
pdf-inspector = { version = "1", features = [
"render-pdfium",
"ocr-oar",
"model-download",
] }
```
```rust
use pdf_inspector::vision::{
route_ocr_pages, run_ocr_pages, HttpModelDownloader, ModelStore,
OarOcrEngine, OcrMode, OcrOptions, PdfiumRenderer, RenderOptions,
PP_OCR_V6_SMALL,
};
let bytes = std::fs::read("scan.pdf")?;
let extraction = pdf_inspector::extract_pages_markdown_mem(&bytes, None)?;
let options = OcrOptions::new().mode(OcrMode::Auto);
let routed = route_ocr_pages(
options.mode,
extraction.pages.len() as u32,
&extraction.pages_needing_ocr,
None,
)?;
if !routed.is_empty() {
// No HTTP request or model initialization occurs before this point.
let store = ModelStore::from_options(&options)?;
let models = store.resolve_or_download(
&PP_OCR_V6_SMALL,
options.model_downloads,
&HttpModelDownloader::default(),
)?;
let run = run_ocr_pages(
&PdfiumRenderer::load()?,
&OarOcrEngine::from_models(&models)?,
&bytes,
&routed,
None,
&RenderOptions::new(),
&options,
)?;
println!("OCR processed {} pages", run.pages.len());
}
```
The downloader accepts HTTPS only, checks a declared content length, caps the
response stream to the pinned size plus one byte, and delegates final size and
SHA-256 verification to `ModelStore`. The store serializes installation across
processes and publishes completed artifacts atomically. Warm caches make no
network calls; offline mode and explicit model directories never download.
### OCR Markdown assembly and native fusion
`fuse_ocr_pages` maps OCR polygons back into PDF coordinates and sends the
result through pdf-inspector's existing deterministic reading-order, table,
and Markdown pipeline. Pages whose native extraction was rejected use OCR
output. When `Force` runs on a clean native page, normalized duplicate OCR
blocks are removed and only additional image-backed text is retained.
```rust
use pdf_inspector::vision::{fuse_ocr_pages, OcrFusionOptions};
let fused = fuse_ocr_pages(
&extraction.pages,
&run,
extraction.pages.len() as u32,
&OcrFusionOptions::new().render_dpi(150.0),
)?;
for page in &fused.pages {
println!("{}", page.markdown);
if page.provenance.hosted_recommended {
eprintln!("page {} needs the hosted document pipeline", page.page + 1);
}
}
```
Each page carries `Native`, `Ocr`, or `Fused` provenance, the exact OCR model
revision, accepted-page confidence, local stage timings, and non-fatal
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.
The one-call API keeps the most recently used verified OCR engine in process.
Long-lived workers therefore verify the pinned artifacts and build the ONNX
sessions once, then reuse those loaded sessions across documents. The cache is
bounded to one model configuration and keyed by normalized model/runtime paths
plus the pinned manifest revision and artifact digests; switching the model
directory, runtime library, or compiled manifest replaces it. An active engine
owns the model data it already verified, so mutating artifacts in place does
not hot-reload a running process; restart the process when intentionally
replacing files at the same paths. CPU inference uses at most four intra-op
threads per ONNX session so a single small page does not oversubscribe larger
hosts, and recognizes variable-width line crops individually to avoid
padding-heavy CPU batches.
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):
+278 -5
View File
@@ -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<f32>) -> 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::<Vec<_>>()
.join(",");
let recommended = result
.pages_recommended_for_ocr
.iter()
.map(u32::to_string)
.collect::<Vec<_>>()
.join(",");
let hosted = result
.pages_recommending_hosted
.iter()
.map(u32::to_string)
.collect::<Vec<_>>()
.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::<Vec<_>>()
.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::<Vec<_>>()
.join(",");
let table_pages = result
.pages_with_tables
.iter()
.map(u32::to_string)
.collect::<Vec<_>>()
.join(",");
let column_pages = result
.pages_with_columns
.iter()
.map(u32::to_string)
.collect::<Vec<_>>()
.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<Option<&'a str>, 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<f32, String> {
argument_value(args, name)?
.map(|value| {
value
.parse::<f32>()
.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<u32>>,
@@ -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 {
+161 -11
View File
@@ -459,8 +459,35 @@ pub fn extract_pages_markdown_mem(
buffer: &[u8],
pages: Option<&[u32]>,
) -> Result<PagesExtractionResult, PdfError> {
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<u32>;
@@ -538,7 +570,10 @@ pub fn extract_pages_markdown_mem(
let (page_items, page_number_removal_mask): (Vec<TextItem>, Vec<bool>) = 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<u32, f32>,
chart_regions: &HashMap<u32, Vec<(f32, f32, f32, f32)>>,
page_count: u32,
) -> HashSet<HeaderFooterItemKey> {
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`].
+8
View File
@@ -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<crate::types::TextLine>,
page_count: u32,
) -> Vec<crate::types::TextLine> {
preprocess::strip_repeated_lines(lines, page_count)
}
/// Convert positioned text items to markdown with structure detection
pub fn to_markdown_from_items(items: Vec<TextItem>, options: MarkdownOptions) -> String {
to_markdown_from_items_with_rects(items, options, &[])
+17
View File
@@ -442,6 +442,16 @@ fn clean_table_cells(cells: &[Vec<String>]) -> (Vec<Vec<String>>, Vec<String>) {
fn is_footnote_row(text: &str) -> bool {
let trimmed = text.trim();
// Japanese documents commonly use the reference mark followed by an
// ASCII or full-width number (for example `※1` / `※1`). These rows often
// sit immediately below a wide table and must not be merged into its last
// data row as wrapped first-column content.
if let Some(rest) = trimmed.strip_prefix('※') {
return rest.chars().next().is_some_and(|character| {
character.is_ascii_digit() || (''..='').contains(&character)
});
}
// Check for common footnote patterns
// (1), (2), etc.
if trimmed.starts_with('(') && trimmed.len() >= 2 {
@@ -503,6 +513,13 @@ mod tests {
assert!(is_footnote_row("NOTES: uppercase"));
}
#[test]
fn test_is_footnote_row_reference_mark_number() {
assert!(is_footnote_row("※1 explanation"));
assert!(is_footnote_row("※1 説明"));
assert!(!is_footnote_row("※ general marker"));
}
#[test]
fn test_is_footnote_row_plain_text_false() {
assert!(!is_footnote_row("Regular cell text"));
+84
View File
@@ -0,0 +1,84 @@
//! HTTPS acquisition for pinned local model artifacts.
use std::fmt;
use std::io::Read;
use std::time::Duration;
use thiserror::Error;
use super::{ModelArtifact, ModelDownloader};
/// Default end-to-end timeout for one model artifact request.
pub const DEFAULT_MODEL_DOWNLOAD_TIMEOUT: Duration = Duration::from_secs(5 * 60);
/// Streaming HTTPS downloader used by lazy model resolution.
#[derive(Clone)]
pub struct HttpModelDownloader {
agent: ureq::Agent,
}
impl fmt::Debug for HttpModelDownloader {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("HttpModelDownloader")
.finish_non_exhaustive()
}
}
impl Default for HttpModelDownloader {
fn default() -> Self {
Self::new(DEFAULT_MODEL_DOWNLOAD_TIMEOUT)
}
}
impl HttpModelDownloader {
/// Creates an HTTPS-only downloader with an end-to-end request timeout.
pub fn new(timeout: Duration) -> Self {
let config = ureq::Agent::config_builder()
.https_only(true)
.timeout_global(Some(timeout))
.user_agent(concat!("pdf-inspector/", env!("CARGO_PKG_VERSION")))
.build();
Self {
agent: ureq::Agent::new_with_config(config),
}
}
}
impl ModelDownloader for HttpModelDownloader {
type Error = HttpModelDownloadError;
fn open(&self, artifact: &ModelArtifact) -> Result<Box<dyn Read + Send>, Self::Error> {
let response = self.agent.get(artifact.url).call()?;
if let Some(actual) = response.body().content_length() {
if actual != artifact.size {
return Err(HttpModelDownloadError::ContentLength {
expected: artifact.size,
actual,
});
}
}
// One extra byte lets ModelStore report an exact size mismatch while
// preventing a malicious or broken server from filling the disk.
let limit = artifact.size.saturating_add(1);
Ok(Box::new(response.into_body().into_reader().take(limit)))
}
}
/// Failures before a response stream reaches [`super::ModelStore`].
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum HttpModelDownloadError {
/// DNS, TLS, redirect, HTTP status, or response-stream setup failed.
#[error(transparent)]
Request(#[from] ureq::Error),
/// The server declared a size that disagrees with the pinned manifest.
#[error("server declared {actual} bytes; manifest requires {expected}")]
ContentLength {
/// Pinned artifact size.
expected: u64,
/// Server-declared size.
actual: u64,
},
}
+1038
View File
File diff suppressed because it is too large Load Diff
+26 -2
View File
@@ -10,12 +10,20 @@
#[cfg(all(feature = "vision", not(target_arch = "wasm32")))]
mod contracts;
#[cfg(all(feature = "model-download", not(target_arch = "wasm32")))]
mod download;
#[cfg(all(feature = "vision", not(target_arch = "wasm32")))]
mod fusion;
#[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 = "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")))]
mod routing;
#[cfg(all(feature = "render-pdfium", not(target_arch = "wasm32")))]
mod pdfium;
@@ -26,16 +34,32 @@ pub use contracts::{
ModelDownloadPolicy, ModelIdentity, OcrEngine, OcrMode, OcrOptions, OcrPage, OcrProfile,
OcrSpan, PageContentSource, PageProvenance, PageRenderer, VisionTimings,
};
#[cfg(all(feature = "model-download", not(target_arch = "wasm32")))]
pub use download::{HttpModelDownloadError, HttpModelDownloader, DEFAULT_MODEL_DOWNLOAD_TIMEOUT};
#[cfg(all(feature = "vision", not(target_arch = "wasm32")))]
pub use fusion::{
fuse_ocr_pages, ocr_page_to_markdown, FusedPageMarkdown, FusedPages, OcrFusionError,
OcrFusionOptions,
};
#[cfg(all(feature = "model-cache", not(target_arch = "wasm32")))]
pub use models::{
ModelArtifact, ModelArtifactKind, ModelManifest, ModelPaths, ModelStore, ModelStoreError,
PP_OCR_V6_SMALL,
ModelAcquireError, ModelArtifact, ModelArtifactKind, ModelDownloader, 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 = "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,
DEFAULT_RENDER_DPI,
};
#[cfg(all(feature = "vision", not(target_arch = "wasm32")))]
pub use routing::{
route_ocr_pages, run_ocr_pages, OcrRoutingError, OcrRun, OcrRunError, RoutedOcrPage,
};
#[cfg(all(feature = "render-pdfium", not(target_arch = "wasm32")))]
+298 -60
View File
@@ -12,7 +12,7 @@ use fs2::FileExt;
use sha2::{Digest, Sha256};
use thiserror::Error;
use super::OcrOptions;
use super::{ModelDownloadPolicy, OcrOptions};
/// Environment variable overriding the default local model cache.
pub const MODEL_CACHE_ENV: &str = "PDF_INSPECTOR_MODEL_CACHE";
@@ -138,6 +138,19 @@ pub struct ModelStore {
override_root: Option<PathBuf>,
}
/// Streaming source for a pinned model artifact.
///
/// Implementations may use HTTP, an object store, or an application-owned
/// package manager. [`ModelStore`] remains responsible for locking, atomic
/// installation, exact size validation, and SHA-256 verification.
pub trait ModelDownloader: Send + Sync {
/// Downloader-specific failure type.
type Error: std::error::Error + Send + Sync + 'static;
/// Opens a streaming reader for one manifest artifact.
fn open(&self, artifact: &ModelArtifact) -> Result<Box<dyn Read + Send>, Self::Error>;
}
impl ModelStore {
/// Creates a model store rooted at an explicit cache directory.
pub fn new(cache_root: impl Into<PathBuf>) -> Self {
@@ -177,16 +190,17 @@ impl ModelStore {
&self.cache_root
}
/// Effective directory containing one manifest's artifacts.
pub(crate) fn model_root(&self, manifest: &ModelManifest) -> PathBuf {
self.override_root
.clone()
.unwrap_or_else(|| self.manifest_cache_root(manifest))
}
/// Validates and resolves every required artifact.
pub fn resolve(&self, manifest: &ModelManifest) -> Result<ModelPaths, ModelStoreError> {
validate_manifest(manifest)?;
let managed_root;
let root = if let Some(root) = self.override_root.as_deref() {
root
} else {
managed_root = self.manifest_cache_root(manifest);
managed_root.as_path()
};
let root = self.model_root(manifest);
let mut artifacts = BTreeMap::new();
for artifact in manifest.artifacts {
@@ -201,6 +215,54 @@ impl ModelStore {
})
}
/// Resolves a complete model set, fetching only missing or invalid managed
/// cache artifacts when policy permits.
///
/// Explicit [`OcrOptions::model_directory`] overrides are never mutated or
/// supplemented from the network. This method also avoids all downloader
/// calls when the cache is already valid or downloads are offline.
pub fn resolve_or_download<D: ModelDownloader>(
&self,
manifest: &ModelManifest,
policy: ModelDownloadPolicy,
downloader: &D,
) -> Result<ModelPaths, ModelAcquireError<D::Error>> {
validate_manifest(manifest).map_err(ModelAcquireError::Store)?;
match self.resolve(manifest) {
Ok(paths) => return Ok(paths),
Err(source) if self.override_root.is_some() => {
return Err(ModelAcquireError::ExplicitDirectoryIncomplete { source });
}
Err(source) if !source.permits_download_recovery() => {
return Err(ModelAcquireError::Store(source));
}
Err(source) if policy == ModelDownloadPolicy::Offline => {
return Err(ModelAcquireError::DownloadsDisabled { source });
}
Err(_) => {}
}
let root = self.manifest_cache_root(manifest);
for artifact in manifest.artifacts {
let _lock = lock_artifact(&root, artifact).map_err(ModelAcquireError::Store)?;
match verify_artifact(&root.join(artifact.filename), artifact) {
Ok(()) => continue,
Err(source) if source.permits_download_recovery() => {}
Err(source) => return Err(ModelAcquireError::Store(source)),
}
let reader =
downloader
.open(artifact)
.map_err(|source| ModelAcquireError::Download {
kind: artifact.kind,
url: artifact.url,
source,
})?;
install_locked(&root, artifact, reader).map_err(ModelAcquireError::Store)?;
}
self.resolve(manifest).map_err(ModelAcquireError::Store)
}
/// Atomically installs one artifact from a reader after validating its
/// exact size and SHA-256 digest.
///
@@ -219,58 +281,8 @@ impl ModelStore {
.find(|artifact| artifact.kind == kind)
.ok_or(ModelStoreError::ArtifactNotInManifest { kind })?;
let root = self.manifest_cache_root(manifest);
fs::create_dir_all(&root).map_err(|source| ModelStoreError::Io {
path: root.clone(),
source,
})?;
let target = root.join(artifact.filename);
let lock_path = root.join(format!(".{}.lock", artifact.filename));
let lock = OpenOptions::new()
.create(true)
.read(true)
.write(true)
.truncate(false)
.open(&lock_path)
.map_err(|source| ModelStoreError::Io {
path: lock_path.clone(),
source,
})?;
FileExt::lock_exclusive(&lock).map_err(|source| ModelStoreError::Io {
path: lock_path,
source,
})?;
if verify_artifact(&target, artifact).is_ok() {
return Ok(target);
}
sweep_stale_install_files(&root, artifact.filename)?;
let (temporary, mut output) = create_temporary_file(&root, artifact.filename)?;
let result = (|| {
let mut limited = reader.by_ref().take(artifact.size.saturating_add(1));
let (size, digest) =
copy_and_hash(&mut limited, &mut output).map_err(|source| ModelStoreError::Io {
path: temporary.clone(),
source,
})?;
output.sync_all().map_err(|source| ModelStoreError::Io {
path: temporary.clone(),
source,
})?;
validate_size_and_hash(artifact, size, &digest, &temporary)?;
replace_file_atomic(&temporary, &target).map_err(|source| ModelStoreError::Io {
path: target.clone(),
source,
})?;
Ok(target.clone())
})();
if result.is_err() {
let _ = fs::remove_file(&temporary);
}
result
let _lock = lock_artifact(&root, artifact)?;
install_locked(&root, artifact, &mut reader)
}
fn manifest_cache_root(&self, manifest: &ModelManifest) -> PathBuf {
@@ -278,6 +290,43 @@ impl ModelStore {
}
}
/// Failures while resolving or lazily acquiring a model set.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum ModelAcquireError<E>
where
E: std::error::Error + Send + Sync + 'static,
{
/// Manifest validation, cache verification, or installation failed.
#[error(transparent)]
Store(#[from] ModelStoreError),
/// A managed-cache artifact could not be downloaded.
#[error("failed to download {kind:?} from {url}: {source}")]
Download {
/// Artifact role.
kind: ModelArtifactKind,
/// Pinned source URL.
url: &'static str,
/// Downloader failure.
#[source]
source: E,
},
/// Offline policy prevented acquisition of an unavailable artifact.
#[error("model artifacts are unavailable and downloads are disabled: {source}")]
DownloadsDisabled {
/// Original verification failure.
#[source]
source: ModelStoreError,
},
/// An explicit package-managed directory was incomplete or invalid.
#[error("explicit model directory is incomplete or invalid: {source}")]
ExplicitDirectoryIncomplete {
/// Original verification failure.
#[source]
source: ModelStoreError,
},
}
/// Failures while validating or installing model artifacts.
#[derive(Debug, Error)]
#[non_exhaustive]
@@ -333,6 +382,17 @@ pub enum ModelStoreError {
},
}
impl ModelStoreError {
fn permits_download_recovery(&self) -> bool {
matches!(
self,
Self::MissingArtifact { .. }
| Self::SizeMismatch { .. }
| Self::ChecksumMismatch { .. }
)
}
}
fn validate_manifest(manifest: &ModelManifest) -> Result<(), ModelStoreError> {
if manifest.schema_version != 1 {
return Err(ModelStoreError::InvalidManifest(format!(
@@ -403,6 +463,67 @@ fn is_single_normal_path_component(value: &str) -> bool {
matches!(components.next(), Some(Component::Normal(_))) && components.next().is_none()
}
fn lock_artifact(root: &Path, artifact: &ModelArtifact) -> Result<File, ModelStoreError> {
fs::create_dir_all(root).map_err(|source| ModelStoreError::Io {
path: root.to_path_buf(),
source,
})?;
let lock_path = root.join(format!(".{}.lock", artifact.filename));
let lock = OpenOptions::new()
.create(true)
.read(true)
.write(true)
.truncate(false)
.open(&lock_path)
.map_err(|source| ModelStoreError::Io {
path: lock_path.clone(),
source,
})?;
FileExt::lock_exclusive(&lock).map_err(|source| ModelStoreError::Io {
path: lock_path,
source,
})?;
Ok(lock)
}
fn install_locked(
root: &Path,
artifact: &ModelArtifact,
mut reader: impl Read,
) -> Result<PathBuf, ModelStoreError> {
let target = root.join(artifact.filename);
if verify_artifact(&target, artifact).is_ok() {
return Ok(target);
}
sweep_stale_install_files(root, artifact.filename)?;
let (temporary, mut output) = create_temporary_file(root, artifact.filename)?;
let result = (|| {
let mut limited = reader.by_ref().take(artifact.size.saturating_add(1));
let (size, digest) =
copy_and_hash(&mut limited, &mut output).map_err(|source| ModelStoreError::Io {
path: temporary.clone(),
source,
})?;
output.sync_all().map_err(|source| ModelStoreError::Io {
path: temporary.clone(),
source,
})?;
validate_size_and_hash(artifact, size, &digest, &temporary)?;
replace_file_atomic(&temporary, &target).map_err(|source| ModelStoreError::Io {
path: target.clone(),
source,
})?;
Ok(target.clone())
})();
if result.is_err() {
let _ = fs::remove_file(&temporary);
}
result
}
fn sweep_stale_install_files(root: &Path, filename: &str) -> Result<(), ModelStoreError> {
let prefix = format!(".{filename}.");
for entry in fs::read_dir(root).map_err(|source| ModelStoreError::Io {
@@ -601,6 +722,34 @@ mod tests {
artifacts: TEST_ARTIFACTS,
};
#[derive(Debug)]
struct StaticDownloader {
bytes: &'static [u8],
requests: std::sync::Mutex<Vec<ModelArtifactKind>>,
}
impl StaticDownloader {
fn new(bytes: &'static [u8]) -> Self {
Self {
bytes,
requests: std::sync::Mutex::new(Vec::new()),
}
}
fn request_count(&self) -> usize {
self.requests.lock().unwrap().len()
}
}
impl ModelDownloader for StaticDownloader {
type Error = io::Error;
fn open(&self, artifact: &ModelArtifact) -> Result<Box<dyn Read + Send>, Self::Error> {
self.requests.lock().unwrap().push(artifact.kind);
Ok(Box::new(io::Cursor::new(self.bytes)))
}
}
#[test]
fn pinned_pp_ocr_manifest_is_well_formed() {
validate_manifest(&PP_OCR_V6_SMALL).unwrap();
@@ -764,4 +913,93 @@ mod tests {
assert_eq!(first, second);
assert_eq!(fs::read(first).unwrap(), b"hello");
}
#[test]
fn resolve_or_download_fetches_once_then_reuses_verified_cache() {
let temp = tempfile::tempdir().unwrap();
let store = ModelStore::new(temp.path());
let downloader = StaticDownloader::new(b"hello");
let first = store
.resolve_or_download(&TEST_MANIFEST, ModelDownloadPolicy::IfMissing, &downloader)
.unwrap();
let second = store
.resolve_or_download(&TEST_MANIFEST, ModelDownloadPolicy::IfMissing, &downloader)
.unwrap();
assert_eq!(first, second);
assert_eq!(downloader.request_count(), 1);
}
#[test]
fn concurrent_resolve_downloads_once_under_the_artifact_lock() {
let temp = tempfile::tempdir().unwrap();
let store = ModelStore::new(temp.path());
let downloader = std::sync::Arc::new(StaticDownloader::new(b"hello"));
let barrier = std::sync::Arc::new(std::sync::Barrier::new(3));
let handles: Vec<_> = (0..2)
.map(|_| {
let store = store.clone();
let downloader = std::sync::Arc::clone(&downloader);
let barrier = std::sync::Arc::clone(&barrier);
std::thread::spawn(move || {
barrier.wait();
store.resolve_or_download(
&TEST_MANIFEST,
ModelDownloadPolicy::IfMissing,
downloader.as_ref(),
)
})
})
.collect();
barrier.wait();
for handle in handles {
handle.join().unwrap().unwrap();
}
assert_eq!(downloader.request_count(), 1);
}
#[cfg(unix)]
#[test]
fn cache_io_failures_do_not_trigger_downloads() {
use std::os::unix::fs::symlink;
let temp = tempfile::tempdir().unwrap();
let store = ModelStore::new(temp.path());
let root = store.manifest_cache_root(&TEST_MANIFEST);
fs::create_dir_all(&root).unwrap();
symlink("hello.txt", root.join("hello.txt")).unwrap();
let downloader = StaticDownloader::new(b"hello");
assert!(matches!(
store.resolve_or_download(&TEST_MANIFEST, ModelDownloadPolicy::IfMissing, &downloader,),
Err(ModelAcquireError::Store(ModelStoreError::Io { .. }))
));
assert_eq!(downloader.request_count(), 0);
}
#[test]
fn resolve_or_download_honors_offline_and_explicit_directory_boundaries() {
let cache = tempfile::tempdir().unwrap();
let override_dir = tempfile::tempdir().unwrap();
let downloader = StaticDownloader::new(b"hello");
let offline = ModelStore::new(cache.path())
.resolve_or_download(&TEST_MANIFEST, ModelDownloadPolicy::Offline, &downloader)
.unwrap_err();
assert!(matches!(
offline,
ModelAcquireError::DownloadsDisabled { .. }
));
let explicit = ModelStore::new(cache.path())
.override_root(override_dir.path())
.resolve_or_download(&TEST_MANIFEST, ModelDownloadPolicy::IfMissing, &downloader)
.unwrap_err();
assert!(matches!(
explicit,
ModelAcquireError::ExplicitDirectoryIncomplete { .. }
));
assert_eq!(downloader.request_count(), 0);
}
}
+34 -5
View File
@@ -4,6 +4,7 @@ use std::path::PathBuf;
use std::time::Instant;
use image::RgbImage;
use oar_ocr::core::config::onnx::OrtSessionConfig;
use oar_ocr::oarocr::{OAROCRBuilder, OAROCR};
use oar_ocr::processors::BoundingBox;
use thiserror::Error;
@@ -86,7 +87,13 @@ impl OarOcrEngine {
let recognition = required_model(models, ModelArtifactKind::TextRecognition)?;
let dictionary = required_model(models, ModelArtifactKind::CharacterDictionary)?;
let pipeline = OAROCRBuilder::new(detection, recognition, dictionary).build()?;
let pipeline = OAROCRBuilder::new(detection, recognition, dictionary)
.ort_session(ocr_session_config())
// Document line crops often have very different widths. Keeping
// CPU recognition batches at one avoids padding every crop to the
// widest line, reducing both inference work and peak memory.
.region_batch_size(1)
.build()?;
let model = ModelIdentity::new(models.manifest_id(), models.revision());
Ok(Self { pipeline, model })
}
@@ -169,11 +176,18 @@ impl OarOcrEngine {
}
}
fn ocr_session_config() -> OrtSessionConfig {
let available = std::thread::available_parallelism()
.map(std::num::NonZeroUsize::get)
.unwrap_or(1);
OrtSessionConfig::new()
.with_intra_threads(available.min(4))
.with_inter_threads(1)
.with_parallel_execution(false)
}
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);
let path = onnx_runtime_library_path();
drop(
ort::init_from(&path).map_err(|source| OarOcrError::OnnxRuntimeLoad {
path: path.clone(),
@@ -183,6 +197,13 @@ fn load_onnx_runtime() -> Result<(), OarOcrError> {
Ok(())
}
pub(crate) fn onnx_runtime_library_path() -> PathBuf {
std::env::var_os(ONNX_RUNTIME_LIBRARY_ENV)
.filter(|path| !path.is_empty())
.map(PathBuf::from)
.unwrap_or_else(default_onnx_runtime_library)
}
fn default_onnx_runtime_library() -> PathBuf {
#[cfg(target_os = "windows")]
const NAME: &str = "onnxruntime.dll";
@@ -355,6 +376,14 @@ mod tests {
use super::*;
use crate::vision::PageTransform;
#[test]
fn cpu_session_budget_is_bounded_for_small_ocr_models() {
let config = ocr_session_config();
assert!((1..=4).contains(&config.intra_threads.unwrap()));
assert_eq!(config.inter_threads, Some(1));
assert_eq!(config.parallel_execution, Some(false));
}
fn page(format: RenderPixelFormat, stride: usize, pixels: Vec<u8>) -> RenderedPage {
let transform =
PageTransform::from_corners(2, 2, (0.0, 2.0), (2.0, 2.0), (0.0, 0.0)).unwrap();
+533
View File
@@ -0,0 +1,533 @@
//! One-call native extraction and OCR pipeline.
use std::collections::BTreeSet;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::Instant;
use thiserror::Error;
use crate::{MarkdownOptions, PageOcrReasons, PdfError};
use super::oar::onnx_runtime_library_path;
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,
};
#[derive(Debug, Clone, PartialEq, Eq)]
struct OcrEngineCacheKey {
model_root: PathBuf,
runtime_library: PathBuf,
manifest_schema: u32,
manifest_id: &'static str,
manifest_revision: &'static str,
artifact_digests: Vec<&'static str>,
}
#[derive(Debug)]
struct CachedOcrEngine {
key: OcrEngineCacheKey,
engine: Arc<OarOcrEngine>,
}
static OCR_ENGINE_CACHE: OnceLock<Mutex<Option<CachedOcrEngine>>> = OnceLock::new();
/// 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<BTreeSet<u32>>,
/// Password for an encrypted PDF.
pub password: Option<String>,
/// 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<Item = u32>) -> 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<String>) -> 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<FusedPageMarkdown>,
/// 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<u32>,
/// 1-indexed pages actually rendered and recognized.
pub pages_routed_to_ocr: Vec<u32>,
/// 1-indexed pages whose OCR result recommends hosted document parsing.
pub pages_recommending_hosted: Vec<u32>,
/// Original machine-readable OCR reasons for selected pages.
pub ocr_reasons_by_page: Vec<PageOcrReasons>,
/// Selected pages where deterministic table detection found tables.
pub pages_with_tables: Vec<u32>,
/// Selected pages where deterministic layout found multiple columns.
pub pages_with_columns: Vec<u32>,
/// 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<Path>,
options: OcrPdfOptions,
) -> Result<OcrPdfResult, OcrPipelineError> {
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<OcrPdfResult, OcrPipelineError> {
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<Vec<u32>> = options
.page_filter
.as_ref()
.map(|pages| pages.iter().copied().collect());
let selected_pages_zero_indexed: Option<Vec<u32>> = 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 engine = cached_ocr_engine(&options.ocr)?;
run_ocr_pages(
&renderer,
engine.as_ref(),
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 cached_ocr_engine(options: &OcrOptions) -> Result<Arc<OarOcrEngine>, OcrPipelineError> {
let store = ModelStore::from_options(options)?;
let key = OcrEngineCacheKey {
model_root: normalized_cache_path(store.model_root(&PP_OCR_V6_SMALL)),
runtime_library: normalized_cache_path(onnx_runtime_library_path()),
manifest_schema: PP_OCR_V6_SMALL.schema_version,
manifest_id: PP_OCR_V6_SMALL.id,
manifest_revision: PP_OCR_V6_SMALL.revision,
artifact_digests: PP_OCR_V6_SMALL
.artifacts
.iter()
.map(|artifact| artifact.sha256)
.collect(),
};
let cache = OCR_ENGINE_CACHE.get_or_init(|| Mutex::new(None));
{
let cached = cache.lock().unwrap_or_else(|error| error.into_inner());
if let Some(cached) = cached.as_ref().filter(|cached| cached.key == key) {
return Ok(Arc::clone(&cached.engine));
}
}
// The loaded sessions own the verified model data, so a cache hit never
// needs to trust or reopen mutable files on disk. Do the expensive download
// and session construction outside the process-wide cache lock so unrelated
// OCR requests can continue using a warm engine. Concurrent cold misses may
// build redundantly; the second cache check keeps only one shared session.
let models = store.resolve_or_download(
&PP_OCR_V6_SMALL,
options.model_downloads,
&HttpModelDownloader::default(),
)?;
let engine = Arc::new(OarOcrEngine::from_models(&models)?);
let mut cached = cache.lock().unwrap_or_else(|error| error.into_inner());
if let Some(cached) = cached.as_ref().filter(|cached| cached.key == key) {
return Ok(Arc::clone(&cached.engine));
}
*cached = Some(CachedOcrEngine {
key,
engine: Arc::clone(&engine),
});
Ok(engine)
}
fn normalized_cache_path(path: PathBuf) -> PathBuf {
let absolute = if path.is_absolute() {
path
} else {
std::env::current_dir()
.map(|directory| directory.join(&path))
.unwrap_or(path)
};
if let Ok(canonical) = std::fs::canonicalize(&absolute) {
return canonical;
}
// A managed cache often does not exist on the first request. Resolve the
// nearest existing ancestor so a relative path has the same key before
// and after model acquisition creates its final directories.
let mut ancestor = absolute.as_path();
let mut suffix = Vec::new();
while let Some(name) = ancestor.file_name() {
suffix.push(name.to_os_string());
let Some(parent) = ancestor.parent() else {
break;
};
ancestor = parent;
if let Ok(mut canonical) = std::fs::canonicalize(ancestor) {
for component in suffix.iter().rev() {
canonical.push(component);
}
return canonical;
}
}
absolute
}
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!("<!-- Page {} -->\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<HttpModelDownloadError>),
/// 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 01 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 cache_path_is_stable_when_a_relative_directory_is_created() {
let current = std::fs::canonicalize(std::env::current_dir().unwrap()).unwrap();
let temporary = tempfile::tempdir_in(&current).unwrap();
let relative_root = temporary.path().strip_prefix(&current).unwrap();
let relative = relative_root.join("models").join("revision");
let before = normalized_cache_path(relative.clone());
std::fs::create_dir_all(&relative).unwrap();
let after = normalized_cache_path(relative);
assert!(before.is_absolute());
assert_eq!(before, after);
}
#[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("<!-- Page 2 -->"));
}
#[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 { .. }
))
));
}
}
+380
View File
@@ -0,0 +1,380 @@
//! Page routing and renderer/OCR orchestration without Markdown fusion.
use std::collections::BTreeSet;
use std::error::Error;
use std::time::Instant;
use thiserror::Error;
use super::{OcrEngine, OcrMode, OcrOptions, OcrPage, PageRenderer, RenderOptions, RenderedPage};
/// A rendered page paired with OCR output in the same bitmap coordinate space.
#[derive(Debug)]
pub struct RoutedOcrPage {
/// Renderer-owned bitmap and pixel↔PDF transform.
pub rendered: RenderedPage,
/// Positioned OCR spans for the bitmap.
pub ocr: OcrPage,
}
/// Output of one selective OCR invocation.
#[derive(Debug)]
pub struct OcrRun {
/// Pages processed in ascending document order.
pub pages: Vec<RoutedOcrPage>,
/// Total page-rendering wall time.
pub render_time_ms: u64,
/// Total engine wall time.
pub ocr_time_ms: u64,
}
/// Selects 1-indexed pages for OCR.
///
/// `recommended_pages` comes from pdf-inspector's existing detector/text
/// quality signals. `selected_pages` is an optional user page filter. Results
/// are validated, deduplicated, and returned in document order.
pub fn route_ocr_pages(
mode: OcrMode,
page_count: u32,
recommended_pages: &[u32],
selected_pages: Option<&[u32]>,
) -> Result<Vec<u32>, OcrRoutingError> {
match mode {
OcrMode::Off => Ok(Vec::new()),
OcrMode::Auto => {
let mut routed = validated_page_set("recommended", recommended_pages, page_count)?;
if let Some(selected) = selected_pages {
let selected = validated_page_set("selected", selected, page_count)?;
routed.retain(|page| selected.contains(page));
}
Ok(routed.into_iter().collect())
}
OcrMode::Force => {
let routed = selected_pages
.map(|pages| validated_page_set("selected", pages, page_count))
.transpose()?
.unwrap_or_else(|| (1..=page_count).collect());
Ok(routed.into_iter().collect())
}
}
}
/// Renders and recognizes already-routed pages while retaining transforms for
/// the following fusion layer.
///
/// An empty page list returns without calling either dependency, which keeps
/// model resolution and inference lazy when Auto routing finds no OCR work.
#[allow(clippy::too_many_arguments)]
pub fn run_ocr_pages<R, O>(
renderer: &R,
engine: &O,
pdf_bytes: &[u8],
pages: &[u32],
password: Option<&str>,
render_options: &RenderOptions,
ocr_options: &OcrOptions,
) -> Result<OcrRun, OcrRunError>
where
R: PageRenderer,
O: OcrEngine,
{
if pages.is_empty() {
return Ok(OcrRun {
pages: Vec::new(),
render_time_ms: 0,
ocr_time_ms: 0,
});
}
if ocr_options.mode == OcrMode::Off {
return Err(OcrRunError::OcrDisabled);
}
let render_started = Instant::now();
let rendered = renderer
.render_pages(pdf_bytes, pages, password, render_options)
.map_err(|source| OcrRunError::Render {
source: Box::new(source),
})?;
let render_time_ms = elapsed_ms(render_started);
validate_page_order("renderer", pages, rendered.iter().map(RenderedPage::page))?;
let ocr_started = Instant::now();
let recognized =
engine
.recognize(&rendered, ocr_options)
.map_err(|source| OcrRunError::Ocr {
source: Box::new(source),
})?;
let ocr_time_ms = elapsed_ms(ocr_started);
validate_page_order("OCR engine", pages, recognized.iter().map(|page| page.page))?;
Ok(OcrRun {
pages: rendered
.into_iter()
.zip(recognized)
.map(|(rendered, ocr)| RoutedOcrPage { rendered, ocr })
.collect(),
render_time_ms,
ocr_time_ms,
})
}
/// Invalid page routing or renderer/engine contract output.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum OcrRoutingError {
/// A page list contained zero or a page beyond the document.
#[error("{source_name} OCR page {page} is outside the valid range 1..={page_count}")]
InvalidPage {
/// Page-list source.
source_name: &'static str,
/// Invalid 1-indexed page.
page: u32,
/// Document page count.
page_count: u32,
},
}
/// Failures while rendering and recognizing a routed page set.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum OcrRunError {
/// A non-empty route cannot execute with OCR disabled.
#[error("cannot process routed pages while OCR mode is Off")]
OcrDisabled,
/// Page rasterization failed.
#[error("page rendering failed: {source}")]
Render {
/// Renderer-specific failure.
#[source]
source: Box<dyn Error + Send + Sync>,
},
/// OCR inference failed.
#[error("OCR inference failed: {source}")]
Ocr {
/// Engine-specific failure.
#[source]
source: Box<dyn Error + Send + Sync>,
},
/// A dependency returned the wrong count or order.
#[error("{stage} returned pages {actual:?}; expected {expected:?}")]
PageOrderMismatch {
/// Dependency boundary that violated the contract.
stage: &'static str,
/// Requested 1-indexed pages.
expected: Vec<u32>,
/// Returned 1-indexed pages.
actual: Vec<u32>,
},
}
fn validated_page_set(
source_name: &'static str,
pages: &[u32],
page_count: u32,
) -> Result<BTreeSet<u32>, OcrRoutingError> {
let mut result = BTreeSet::new();
for &page in pages {
if page == 0 || page > page_count {
return Err(OcrRoutingError::InvalidPage {
source_name,
page,
page_count,
});
}
result.insert(page);
}
Ok(result)
}
fn validate_page_order(
stage: &'static str,
expected: &[u32],
actual: impl IntoIterator<Item = u32>,
) -> Result<(), OcrRunError> {
let actual: Vec<u32> = actual.into_iter().collect();
if actual == expected {
Ok(())
} else {
Err(OcrRunError::PageOrderMismatch {
stage,
expected: expected.to_vec(),
actual,
})
}
}
fn elapsed_ms(started: Instant) -> u64 {
u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::vision::{
ImagePoint, ImageQuad, ModelIdentity, OcrSpan, PageTransform, RenderBufferError,
RenderPixelFormat,
};
#[derive(Debug, Error)]
#[error("fake failure")]
struct FakeError;
struct FakeRenderer;
impl PageRenderer for FakeRenderer {
type Error = FakeError;
fn render_pages(
&self,
_pdf_bytes: &[u8],
pages: &[u32],
_password: Option<&str>,
_options: &RenderOptions,
) -> Result<Vec<RenderedPage>, Self::Error> {
pages
.iter()
.copied()
.map(rendered_page)
.collect::<Result<_, _>>()
.map_err(|_| FakeError)
}
}
struct FakeEngine {
model: ModelIdentity,
}
impl FakeEngine {
fn new() -> Self {
Self {
model: ModelIdentity::new("fake", "v1"),
}
}
}
impl OcrEngine for FakeEngine {
type Error = FakeError;
fn model(&self) -> &ModelIdentity {
&self.model
}
fn recognize(
&self,
pages: &[RenderedPage],
_options: &OcrOptions,
) -> Result<Vec<OcrPage>, Self::Error> {
Ok(pages
.iter()
.map(|page| OcrPage {
page: page.page(),
spans: vec![OcrSpan {
text: format!("page {}", page.page()),
polygon: ImageQuad::new([
ImagePoint::new(0.0, 0.0),
ImagePoint::new(1.0, 0.0),
ImagePoint::new(1.0, 1.0),
ImagePoint::new(0.0, 1.0),
]),
confidence: 0.9,
orientation_degrees: None,
}],
mean_confidence: Some(0.9),
model: self.model.clone(),
processing_time_ms: 1,
warnings: Vec::new(),
})
.collect())
}
}
fn rendered_page(page: u32) -> Result<RenderedPage, RenderBufferError> {
let transform =
PageTransform::from_corners(1, 1, (0.0, 1.0), (1.0, 1.0), (0.0, 0.0)).unwrap();
RenderedPage::new(
page,
1.0,
1.0,
1,
1,
3,
RenderPixelFormat::Rgb8,
vec![255; 3],
transform,
)
}
#[test]
fn off_auto_and_force_route_expected_pages() {
assert_eq!(
route_ocr_pages(OcrMode::Off, 0, &[99], Some(&[0])).unwrap(),
Vec::<u32>::new()
);
assert_eq!(
route_ocr_pages(OcrMode::Auto, 5, &[5, 3, 3, 1], Some(&[2, 3, 5])).unwrap(),
vec![3, 5]
);
assert_eq!(
route_ocr_pages(OcrMode::Force, 4, &[], None).unwrap(),
vec![1, 2, 3, 4]
);
assert_eq!(
route_ocr_pages(OcrMode::Force, 4, &[], Some(&[4, 2, 2])).unwrap(),
vec![2, 4]
);
}
#[test]
fn routing_rejects_invalid_page_numbers() {
assert!(matches!(
route_ocr_pages(OcrMode::Auto, 2, &[0], None),
Err(OcrRoutingError::InvalidPage { page: 0, .. })
));
assert!(matches!(
route_ocr_pages(OcrMode::Force, 2, &[], Some(&[3])),
Err(OcrRoutingError::InvalidPage { page: 3, .. })
));
}
#[test]
fn run_retains_render_transforms_and_input_order() {
let options = OcrOptions::new().mode(OcrMode::Auto);
let run = run_ocr_pages(
&FakeRenderer,
&FakeEngine::new(),
b"pdf",
&[2, 4],
None,
&RenderOptions::new(),
&options,
)
.unwrap();
assert_eq!(
run.pages
.iter()
.map(|page| page.rendered.page())
.collect::<Vec<_>>(),
vec![2, 4]
);
assert_eq!(run.pages[1].ocr.spans[0].text, "page 4");
assert_eq!(run.pages[1].rendered.pixel_to_page(0.0, 0.0).y, 1.0);
}
#[test]
fn empty_route_is_a_noop_even_when_ocr_is_off() {
let run = run_ocr_pages(
&FakeRenderer,
&FakeEngine::new(),
b"pdf",
&[],
None,
&RenderOptions::new(),
&OcrOptions::new(),
)
.unwrap();
assert!(run.pages.is_empty());
assert_eq!(run.render_time_ms, 0);
assert_eq!(run.ocr_time_ms, 0);
}
}
+35
View File
@@ -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,37 @@ 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 options = OcrPdfOptions::new().ocr(ocr);
let result = process_pdf_with_ocr_mem(&bytes, options.clone()).unwrap();
let repeated = process_pdf_with_ocr_mem(&bytes, options).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
);
assert_eq!(repeated.markdown, result.markdown);
}
fn recognize(
model_directory: &std::ffi::OsStr,
pages: &[RenderedPage],