Compare commits
15
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
634a29f04d | ||
|
|
0e2e287c04 | ||
|
|
616f9b59fb | ||
|
|
21a436ac1b | ||
|
|
d2d8e35a7b | ||
|
|
30c9dbbc72 | ||
|
|
162c5cbf10 | ||
|
|
4e4cfdc74c | ||
|
|
24e66245cb | ||
|
|
cee7b6c381 | ||
|
|
b5a63d7036 | ||
|
|
6829397bce | ||
|
|
fc16133a58 | ||
|
|
41a6a67c03 | ||
|
|
58fe5a0224 |
@@ -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
@@ -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
@@ -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
@@ -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`].
|
||||
|
||||
@@ -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, &[])
|
||||
|
||||
@@ -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() || ('0'..='9').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"));
|
||||
|
||||
@@ -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,
|
||||
},
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
+26
-2
@@ -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
@@ -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
@@ -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();
|
||||
|
||||
@@ -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 0–1 range.
|
||||
#[error("minimum OCR confidence must be between 0 and 1, got {value}")]
|
||||
InvalidMinimumConfidence {
|
||||
/// Invalid value.
|
||||
value: f32,
|
||||
},
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[test]
|
||||
fn 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(¤t).unwrap();
|
||||
let relative_root = temporary.path().strip_prefix(¤t).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 { .. }
|
||||
))
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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],
|
||||
|
||||
Reference in New Issue
Block a user