Compare commits

..
12 changed files with 2645 additions and 78 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"
+156 -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,156 @@ 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.
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, &[])
+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,
},
}
+797
View File
@@ -0,0 +1,797 @@
//! Geometry-aware OCR Markdown assembly and native-text fusion.
use std::collections::{BTreeMap, BTreeSet};
use std::time::Instant;
use thiserror::Error;
use crate::markdown::{to_markdown_from_items_with_rects_and_page_count, MarkdownOptions};
use crate::types::{ItemType, TextItem};
use crate::PageMarkdown;
use super::{
OcrRun, PageContentSource, PageProvenance, RoutedOcrPage, VisionTimings, DEFAULT_RENDER_DPI,
};
/// OCR assembly and hosted-fallback policy.
#[derive(Debug, Clone)]
pub struct OcrFusionOptions {
/// Markdown conversion options applied to positioned OCR spans.
pub markdown: MarkdownOptions,
/// Render resolution recorded in page provenance.
pub render_dpi: f32,
/// Recommend the hosted pipeline below this mean confidence when native
/// extraction already marked the page as requiring OCR.
pub hosted_recommendation_confidence: f32,
}
impl Default for OcrFusionOptions {
fn default() -> Self {
Self {
markdown: MarkdownOptions {
include_page_numbers: false,
strip_headers_footers: false,
..MarkdownOptions::default()
},
render_dpi: DEFAULT_RENDER_DPI,
hosted_recommendation_confidence: 0.5,
}
}
}
impl OcrFusionOptions {
/// Creates OCR fusion options with local defaults.
pub fn new() -> Self {
Self::default()
}
/// Replaces Markdown conversion options.
pub fn markdown(mut self, markdown: MarkdownOptions) -> Self {
self.markdown = markdown;
self
}
/// Records the renderer resolution in provenance.
pub fn render_dpi(mut self, render_dpi: f32) -> Self {
self.render_dpi = render_dpi;
self
}
/// Sets the weak-OCR confidence threshold for recommending hosted parsing.
pub fn hosted_recommendation_confidence(mut self, confidence: f32) -> Self {
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.
#[derive(Debug, Clone, PartialEq)]
pub struct FusedPageMarkdown {
/// 1-indexed document page number, matching OCR and provenance fields.
pub page: u32,
/// Final page Markdown.
pub markdown: String,
/// Native/OCR source, model, timing, and fallback metadata.
pub provenance: PageProvenance,
}
/// Output of fusing a selective OCR run into native page extraction.
#[derive(Debug, Clone, PartialEq)]
pub struct FusedPages {
/// Pages in the same order as the native input.
pub pages: Vec<FusedPageMarkdown>,
/// Batch rendering wall time from the OCR run.
pub render_time_ms: u64,
/// Batch OCR wall time from the OCR run.
pub ocr_time_ms: u64,
}
/// Converts positioned OCR spans to Markdown through pdf-inspector's existing
/// deterministic geometry, reading-order, table, and Markdown pipeline.
///
/// The page number is 1-indexed. `document_page_count` prevents a selected
/// page from being mistaken for a one-page document during page-number logic.
pub fn ocr_page_to_markdown(
page: &RoutedOcrPage,
document_page_count: u32,
options: &MarkdownOptions,
) -> String {
let (items, _) = ocr_text_items(page);
to_markdown_from_items_with_rects_and_page_count(
items,
options.clone(),
&[],
document_page_count,
)
}
/// Fuses a selective OCR run into per-page native Markdown.
///
/// OCR replaces pages whose native extraction was already rejected. On clean
/// native pages (for example in `Force` mode), normalized duplicate OCR blocks
/// are removed and only genuinely additional blocks are appended. Pages that
/// needed OCR but still have no credible local result recommend the hosted
/// document pipeline instead of silently presenting an empty result as final.
pub fn fuse_ocr_pages(
native_pages: &[PageMarkdown],
ocr_run: &OcrRun,
document_page_count: u32,
options: &OcrFusionOptions,
) -> Result<FusedPages, OcrFusionError> {
validate_options(options)?;
let mut native_numbers = BTreeSet::new();
for page in native_pages {
let page_number = page
.page
.checked_add(1)
.ok_or(OcrFusionError::PageOverflow)?;
if !native_numbers.insert(page_number) {
return Err(OcrFusionError::DuplicateNativePage { page: page_number });
}
}
let mut ocr_by_page = BTreeMap::new();
for page in &ocr_run.pages {
let page_number = page.rendered.page();
if !native_numbers.contains(&page_number) {
return Err(OcrFusionError::UnexpectedOcrPage { page: page_number });
}
if ocr_by_page.insert(page_number, page).is_some() {
return Err(OcrFusionError::DuplicateOcrPage { page: page_number });
}
}
let render_shares = equal_time_shares(ocr_run.render_time_ms, ocr_run.pages.len());
let render_by_page: BTreeMap<u32, u64> = ocr_run
.pages
.iter()
.zip(render_shares)
.map(|(page, share)| (page.rendered.page(), share))
.collect();
let mut pages = Vec::with_capacity(native_pages.len());
for native in native_pages {
let page_number = native.page + 1;
let assembly_started = Instant::now();
let mut warnings = Vec::new();
let (markdown, source, ocr_model, ocr_confidence, ocr_ms, hosted_recommended) =
if let Some(local) = ocr_by_page.get(&page_number) {
let (ocr_items, discarded_spans) = ocr_text_items(local);
if discarded_spans > 0 {
warnings.push(format!(
"discarded {discarded_spans} OCR spans with unusable text or geometry"
));
}
warnings.extend(local.ocr.warnings.iter().cloned());
let ocr_markdown = to_markdown_from_items_with_rects_and_page_count(
ocr_items,
options.markdown.clone(),
&[],
document_page_count,
);
let (markdown, source) = if native.markdown.trim().is_empty() || native.needs_ocr {
(ocr_markdown, PageContentSource::Ocr)
} else {
merge_native_and_ocr(&native.markdown, &ocr_markdown)
};
let weak_ocr = local
.ocr
.mean_confidence
.is_none_or(|confidence| confidence < options.hosted_recommendation_confidence);
let recommend_hosted = native.needs_ocr && (markdown.trim().is_empty() || weak_ocr);
if native.needs_ocr && markdown.trim().is_empty() {
warnings.push("OCR produced no usable text".to_string());
}
(
markdown,
source,
Some(local.ocr.model.clone()),
local.ocr.mean_confidence,
local.ocr.processing_time_ms,
recommend_hosted,
)
} else {
if native.needs_ocr {
warnings.push("page was recommended for OCR but was not processed".to_string());
}
(
native.markdown.clone(),
PageContentSource::Native,
None,
None,
0,
native.needs_ocr,
)
};
pages.push(FusedPageMarkdown {
page: page_number,
markdown,
provenance: PageProvenance {
page: page_number,
source,
ocr_model,
layout_model: None,
render_dpi: ocr_by_page
.contains_key(&page_number)
.then_some(options.render_dpi),
ocr_confidence,
timings: VisionTimings {
render_ms: render_by_page.get(&page_number).copied().unwrap_or(0),
ocr_ms,
layout_ms: 0,
assembly_ms: elapsed_ms(assembly_started),
},
warnings,
hosted_recommended,
},
});
}
Ok(FusedPages {
pages,
render_time_ms: ocr_run.render_time_ms,
ocr_time_ms: ocr_run.ocr_time_ms,
})
}
/// Converts recognized line polygons to ordinary PDF-space text items.
fn ocr_text_items(page: &RoutedOcrPage) -> (Vec<TextItem>, usize) {
let mut discarded = 0usize;
let mut items = Vec::with_capacity(page.ocr.spans.len());
for span in &page.ocr.spans {
if span.text.trim().is_empty() {
discarded += 1;
continue;
}
let Some((left, top, right, bottom)) = image_quad_bounds(
&span.polygon.points,
page.rendered.width(),
page.rendered.height(),
) else {
discarded += 1;
continue;
};
let rect = page.rendered.pixel_rect_to_pdf_rect(
f64::from(left),
f64::from(top),
f64::from(right - left),
f64::from(bottom - top),
);
items.push(TextItem {
text: span.text.trim().to_string(),
x: rect.x,
y: rect.y,
width: rect.width,
height: rect.height,
font: "OCR".to_string(),
font_size: rect.height.max(1.0),
page: page.rendered.page(),
is_bold: false,
is_italic: false,
is_underline: false,
is_strikeout: false,
item_type: ItemType::Text,
mcid: None,
});
}
// OCR engines do not share an ordering contract. Geometry gives the
// deterministic top-to-bottom seed expected by the existing layout
// pipeline, which can still replace it with column-aware reading order.
items.sort_by(|first, second| {
first
.page
.cmp(&second.page)
.then(second.y.total_cmp(&first.y))
.then(first.x.total_cmp(&second.x))
});
(items, discarded)
}
fn image_quad_bounds(
points: &[super::ImagePoint; 4],
width: u32,
height: u32,
) -> Option<(f32, f32, f32, f32)> {
if points
.iter()
.any(|point| !point.x.is_finite() || !point.y.is_finite())
{
return None;
}
let left = points
.iter()
.map(|point| point.x)
.fold(f32::INFINITY, f32::min)
.clamp(0.0, width as f32);
let right = points
.iter()
.map(|point| point.x)
.fold(f32::NEG_INFINITY, f32::max)
.clamp(0.0, width as f32);
let top = points
.iter()
.map(|point| point.y)
.fold(f32::INFINITY, f32::min)
.clamp(0.0, height as f32);
let bottom = points
.iter()
.map(|point| point.y)
.fold(f32::NEG_INFINITY, f32::max)
.clamp(0.0, height as f32);
(right > left && bottom > top).then_some((left, top, right, bottom))
}
fn merge_native_and_ocr(native: &str, ocr: &str) -> (String, PageContentSource) {
let native_keys = comparison_units(native);
let mut addition_keys = Vec::new();
let mut additions: Vec<String> = Vec::new();
for block in markdown_blocks(ocr) {
let key = ComparisonKey::new(block);
if key.is_empty()
|| native_keys.iter().any(|native| native.same_content(&key))
|| addition_keys
.iter()
.any(|existing: &ComparisonKey| existing.same_content(&key))
{
continue;
}
let addition = if let Some(native) = best_partial_overlap(&key, &native_keys) {
let Some(novel) = novel_fragments(block, native) else {
continue;
};
novel
} else {
block.trim_end().to_string()
};
let addition_key = ComparisonKey::new(&addition);
if addition_key.is_empty()
|| native_keys
.iter()
.any(|native| native.same_content(&addition_key))
|| addition_keys
.iter()
.any(|existing: &ComparisonKey| existing.same_content(&addition_key))
{
continue;
}
addition_keys.push(addition_key);
additions.push(addition);
}
if additions.is_empty() {
(ensure_trailing_newline(native), PageContentSource::Native)
} else {
let mut result = native.trim_end_matches('\n').to_string();
if !result.is_empty() {
result.push_str("\n\n");
}
result.push_str(&additions.join("\n\n"));
result.push('\n');
(result, PageContentSource::Fused)
}
}
fn markdown_blocks(markdown: &str) -> impl Iterator<Item = &str> {
markdown
.split("\n\n")
.map(str::trim_end)
.filter(|block| !block.trim().is_empty())
}
#[derive(Debug)]
struct ComparisonKey {
tokens: Vec<String>,
punctuation: String,
}
impl ComparisonKey {
fn new(text: &str) -> Self {
let mut tokens = Vec::new();
let mut token = String::new();
let mut punctuation = String::new();
for character in text.chars().flat_map(char::to_lowercase) {
if character.is_alphanumeric() {
token.push(character);
} else {
if !token.is_empty() {
tokens.push(std::mem::take(&mut token));
}
if !character.is_whitespace() {
punctuation.push(character);
}
}
}
if !token.is_empty() {
tokens.push(token);
}
Self {
tokens,
punctuation,
}
}
fn is_empty(&self) -> bool {
self.tokens.is_empty() && self.punctuation.is_empty()
}
fn same_content(&self, other: &Self) -> bool {
if self.tokens.is_empty() || other.tokens.is_empty() {
self.tokens.is_empty()
&& other.tokens.is_empty()
&& self.punctuation == other.punctuation
} else {
token_counts(&self.tokens) == token_counts(&other.tokens)
}
}
}
fn comparison_units(markdown: &str) -> Vec<ComparisonKey> {
let mut units: Vec<_> = markdown_blocks(markdown).map(ComparisonKey::new).collect();
units.extend(
markdown
.lines()
.filter(|line| !line.trim().is_empty())
.map(ComparisonKey::new),
);
units
}
fn token_counts(tokens: &[String]) -> BTreeMap<&str, usize> {
let mut counts = BTreeMap::new();
for token in tokens {
*counts.entry(token.as_str()).or_insert(0) += 1;
}
counts
}
fn best_partial_overlap<'a>(
ocr: &ComparisonKey,
native: &'a [ComparisonKey],
) -> Option<&'a ComparisonKey> {
if ocr.tokens.len() < 2 {
return None;
}
native
.iter()
.filter(|key| !key.tokens.is_empty())
.filter_map(|key| {
let overlap = token_overlap(&ocr.tokens, &key.tokens);
let substantial =
overlap >= 2 && overlap * 2 >= ocr.tokens.len() && overlap * 2 >= key.tokens.len();
substantial.then_some((overlap, key))
})
.max_by_key(|(overlap, _)| *overlap)
.map(|(_, key)| key)
}
fn token_overlap(first: &[String], second: &[String]) -> usize {
let first = token_counts(first);
let second = token_counts(second);
first
.iter()
.map(|(token, count)| (*count).min(second.get(token).copied().unwrap_or(0)))
.sum()
}
fn novel_fragments(block: &str, native: &ComparisonKey) -> Option<String> {
let mut available = token_counts(&native.tokens);
let mut novel = Vec::new();
for fragment in block.split_whitespace() {
let key = ComparisonKey::new(fragment);
let mut adds_content = false;
for token in &key.tokens {
match available.get_mut(token.as_str()) {
Some(count) if *count > 0 => *count -= 1,
_ => adds_content = true,
}
}
if adds_content {
novel.push(fragment);
}
}
(!novel.is_empty()).then(|| novel.join(" "))
}
fn ensure_trailing_newline(markdown: &str) -> String {
let mut result = markdown.trim_end_matches('\n').to_string();
result.push('\n');
result
}
fn equal_time_shares(total: u64, count: usize) -> Vec<u64> {
if count == 0 {
return Vec::new();
}
let count = count as u64;
let base = total / count;
let remainder = total % count;
(0..count)
.map(|index| base + u64::from(index < remainder))
.collect()
}
fn validate_options(options: &OcrFusionOptions) -> Result<(), OcrFusionError> {
if !options.render_dpi.is_finite() || options.render_dpi <= 0.0 {
return Err(OcrFusionError::InvalidRenderDpi {
value: options.render_dpi,
});
}
let confidence = options.hosted_recommendation_confidence;
if !confidence.is_finite() || !(0.0..=1.0).contains(&confidence) {
return Err(OcrFusionError::InvalidHostedConfidence { value: confidence });
}
Ok(())
}
fn elapsed_ms(started: Instant) -> u64 {
u64::try_from(started.elapsed().as_millis()).unwrap_or(u64::MAX)
}
/// Invalid page sets or fusion options.
#[derive(Debug, Error, PartialEq)]
#[non_exhaustive]
pub enum OcrFusionError {
/// A 0-indexed native page could not be converted to 1-indexed form.
#[error("native page number cannot be converted to a 1-indexed page")]
PageOverflow,
/// Native input repeated a page number.
#[error("native page {page} appears more than once")]
DuplicateNativePage {
/// Repeated 1-indexed page number.
page: u32,
},
/// OCR input repeated a page number.
#[error("OCR page {page} appears more than once")]
DuplicateOcrPage {
/// Repeated 1-indexed page number.
page: u32,
},
/// OCR returned a page that was not part of native extraction.
#[error("OCR page {page} is not present in native page extraction")]
UnexpectedOcrPage {
/// Unexpected 1-indexed page number.
page: u32,
},
/// Render resolution is non-finite or non-positive.
#[error("render DPI must be positive and finite, got {value}")]
InvalidRenderDpi {
/// Invalid value.
value: f32,
},
/// Hosted recommendation confidence is outside 01.
#[error("hosted recommendation confidence must be between 0 and 1, got {value}")]
InvalidHostedConfidence {
/// Invalid value.
value: f32,
},
}
#[cfg(test)]
mod tests {
use super::*;
use crate::vision::{
ImagePoint, ImageQuad, ModelIdentity, OcrPage, OcrSpan, PageTransform, RenderPixelFormat,
RenderedPage,
};
fn native(page: u32, markdown: &str, needs_ocr: bool) -> PageMarkdown {
PageMarkdown {
page,
markdown: markdown.to_string(),
needs_ocr,
ocr_reason: needs_ocr.then(|| "scanned".to_string()),
}
}
fn rendered_page(page: u32) -> RenderedPage {
let transform =
PageTransform::from_corners(200, 100, (0.0, 100.0), (200.0, 100.0), (0.0, 0.0))
.unwrap();
RenderedPage::new(
page,
200.0,
100.0,
200,
100,
600,
RenderPixelFormat::Rgb8,
vec![255; 60_000],
transform,
)
.unwrap()
}
fn span(text: &str, top: f32, confidence: f32) -> OcrSpan {
OcrSpan {
text: text.to_string(),
polygon: ImageQuad::new([
ImagePoint::new(10.0, top),
ImagePoint::new(190.0, top),
ImagePoint::new(190.0, top + 10.0),
ImagePoint::new(10.0, top + 10.0),
]),
confidence,
orientation_degrees: None,
}
}
fn routed_page(page: u32, spans: Vec<OcrSpan>, confidence: Option<f32>) -> RoutedOcrPage {
RoutedOcrPage {
rendered: rendered_page(page),
ocr: OcrPage {
page,
spans,
mean_confidence: confidence,
model: ModelIdentity::new("test-ocr", "v1"),
processing_time_ms: 7,
warnings: Vec::new(),
},
}
}
fn run(pages: Vec<RoutedOcrPage>) -> OcrRun {
OcrRun {
pages,
render_time_ms: 5,
ocr_time_ms: 7,
}
}
#[test]
fn scanned_page_uses_geometry_ordered_ocr_and_provenance() {
let native = [native(0, "", true)];
let run = run(vec![routed_page(
1,
vec![
span("Second line", 30.0, 0.9),
span("First line", 10.0, 0.9),
],
Some(0.9),
)]);
let result = fuse_ocr_pages(&native, &run, 1, &OcrFusionOptions::new()).unwrap();
assert!(
result.pages[0].markdown.find("First").unwrap()
< result.pages[0].markdown.find("Second").unwrap()
);
assert_eq!(result.pages[0].provenance.source, PageContentSource::Ocr);
assert_eq!(result.pages[0].page, 1);
assert_eq!(result.pages[0].page, result.pages[0].provenance.page);
assert_eq!(
result.pages[0].provenance.ocr_model.as_ref().unwrap().name,
"test-ocr"
);
assert!(!result.pages[0].provenance.hosted_recommended);
}
#[test]
fn force_mode_deduplicates_native_content() {
let native = [native(0, "Hello, world!\n", false)];
let run = run(vec![routed_page(
1,
vec![span("Hello world", 10.0, 0.9)],
Some(0.9),
)]);
let result = fuse_ocr_pages(&native, &run, 1, &OcrFusionOptions::new()).unwrap();
assert_eq!(result.pages[0].markdown, "Hello, world!\n");
assert_eq!(result.pages[0].provenance.source, PageContentSource::Native);
}
#[test]
fn force_mode_appends_only_additional_ocr_blocks() {
let native = [native(0, "Native title\n", false)];
let run = run(vec![routed_page(
1,
vec![
span("Native title", 10.0, 0.9),
span("Image-only label", 30.0, 0.9),
],
Some(0.9),
)]);
let result = fuse_ocr_pages(&native, &run, 1, &OcrFusionOptions::new()).unwrap();
assert_eq!(result.pages[0].provenance.source, PageContentSource::Fused);
assert_eq!(result.pages[0].markdown.matches("Native title").count(), 1);
assert!(result.pages[0].markdown.contains("Image-only label"));
}
#[test]
fn force_mode_preserves_punctuation_and_short_distinct_blocks() {
let (markdown, source) = merge_native_and_ocr("Figure 1\n", "1\n\n***");
assert_eq!(source, PageContentSource::Fused);
assert_eq!(markdown, "Figure 1\n\n1\n\n***\n");
}
#[test]
fn force_mode_extracts_only_novel_tokens_from_overlapping_blocks() {
let (markdown, source) = merge_native_and_ocr("A C\n", "## A B C");
assert_eq!(source, PageContentSource::Fused);
assert_eq!(markdown, "A C\n\nB\n");
assert_eq!(markdown.matches('A').count(), 1);
assert_eq!(markdown.matches('C').count(), 1);
}
#[test]
fn force_mode_deduplicates_reordered_tokens() {
let native = [native(0, "Alpha Beta Gamma\n", false)];
let run = run(vec![routed_page(
1,
vec![span("Gamma Alpha Beta", 10.0, 0.9)],
Some(0.9),
)]);
let result = fuse_ocr_pages(&native, &run, 1, &OcrFusionOptions::new()).unwrap();
assert_eq!(result.pages[0].provenance.source, PageContentSource::Native);
assert_eq!(result.pages[0].markdown, "Alpha Beta Gamma\n");
}
#[test]
fn fusion_preserves_native_leading_indentation() {
let native = [native(0, " indented code\n", false)];
let duplicate = run(vec![routed_page(
1,
vec![span("indented code", 10.0, 0.9)],
Some(0.9),
)]);
let result = fuse_ocr_pages(&native, &duplicate, 1, &OcrFusionOptions::new()).unwrap();
assert_eq!(result.pages[0].markdown, " indented code\n");
let addition = run(vec![routed_page(
1,
vec![span("image label", 10.0, 0.9)],
Some(0.9),
)]);
let result = fuse_ocr_pages(&native, &addition, 1, &OcrFusionOptions::new()).unwrap();
assert_eq!(
result.pages[0].markdown,
" indented code\n\nimage label\n"
);
}
#[test]
fn missing_or_weak_required_ocr_recommends_hosted() {
let native = [native(0, "", true), native(1, "", true)];
let run = run(vec![routed_page(
2,
vec![span("uncertain", 10.0, 0.3)],
Some(0.3),
)]);
let result = fuse_ocr_pages(&native, &run, 2, &OcrFusionOptions::new()).unwrap();
assert!(result.pages[0].provenance.hosted_recommended);
assert!(result.pages[1].provenance.hosted_recommended);
}
#[test]
fn rejects_ocr_pages_outside_native_selection() {
let error = fuse_ocr_pages(
&[native(0, "text", false)],
&run(vec![routed_page(2, Vec::new(), None)]),
2,
&OcrFusionOptions::new(),
)
.unwrap_err();
assert_eq!(error, OcrFusionError::UnexpectedOcrPage { page: 2 });
}
#[test]
fn time_shares_preserve_batch_total() {
assert_eq!(equal_time_shares(8, 3), vec![3, 3, 2]);
assert_eq!(equal_time_shares(8, 0), Vec::<u64>::new());
}
}
+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")))]
+290 -53
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 {
@@ -201,6 +214,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 +280,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 +289,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 +381,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 +462,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 +721,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 +912,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);
}
}
+426
View File
@@ -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<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 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!("<!-- 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 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);
}
}
+32
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,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],