Compare commits

...
6 changed files with 192 additions and 22 deletions
+13
View File
@@ -380,6 +380,19 @@ 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
+17
View File
@@ -442,6 +442,16 @@ fn clean_table_cells(cells: &[Vec<String>]) -> (Vec<Vec<String>>, Vec<String>) {
fn is_footnote_row(text: &str) -> bool {
let trimmed = text.trim();
// Japanese documents commonly use the reference mark followed by an
// ASCII or full-width number (for example `※1` / `※1`). These rows often
// sit immediately below a wide table and must not be merged into its last
// data row as wrapped first-column content.
if let Some(rest) = trimmed.strip_prefix('※') {
return rest.chars().next().is_some_and(|character| {
character.is_ascii_digit() || (''..='').contains(&character)
});
}
// Check for common footnote patterns
// (1), (2), etc.
if trimmed.starts_with('(') && trimmed.len() >= 2 {
@@ -503,6 +513,13 @@ mod tests {
assert!(is_footnote_row("NOTES: uppercase"));
}
#[test]
fn test_is_footnote_row_reference_mark_number() {
assert!(is_footnote_row("※1 explanation"));
assert!(is_footnote_row("※1 説明"));
assert!(!is_footnote_row("※ general marker"));
}
#[test]
fn test_is_footnote_row_plain_text_false() {
assert!(!is_footnote_row("Regular cell text"));
+8 -7
View File
@@ -190,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 {
+34 -5
View File
@@ -4,6 +4,7 @@ use std::path::PathBuf;
use std::time::Instant;
use image::RgbImage;
use oar_ocr::core::config::onnx::OrtSessionConfig;
use oar_ocr::oarocr::{OAROCRBuilder, OAROCR};
use oar_ocr::processors::BoundingBox;
use thiserror::Error;
@@ -86,7 +87,13 @@ impl OarOcrEngine {
let recognition = required_model(models, ModelArtifactKind::TextRecognition)?;
let dictionary = required_model(models, ModelArtifactKind::CharacterDictionary)?;
let pipeline = OAROCRBuilder::new(detection, recognition, dictionary).build()?;
let pipeline = OAROCRBuilder::new(detection, recognition, dictionary)
.ort_session(ocr_session_config())
// Document line crops often have very different widths. Keeping
// CPU recognition batches at one avoids padding every crop to the
// widest line, reducing both inference work and peak memory.
.region_batch_size(1)
.build()?;
let model = ModelIdentity::new(models.manifest_id(), models.revision());
Ok(Self { pipeline, model })
}
@@ -169,11 +176,18 @@ impl OarOcrEngine {
}
}
fn ocr_session_config() -> OrtSessionConfig {
let available = std::thread::available_parallelism()
.map(std::num::NonZeroUsize::get)
.unwrap_or(1);
OrtSessionConfig::new()
.with_intra_threads(available.min(4))
.with_inter_threads(1)
.with_parallel_execution(false)
}
fn load_onnx_runtime() -> Result<(), OarOcrError> {
let path = std::env::var_os(ONNX_RUNTIME_LIBRARY_ENV)
.filter(|path| !path.is_empty())
.map(PathBuf::from)
.unwrap_or_else(default_onnx_runtime_library);
let path = onnx_runtime_library_path();
drop(
ort::init_from(&path).map_err(|source| OarOcrError::OnnxRuntimeLoad {
path: path.clone(),
@@ -183,6 +197,13 @@ fn load_onnx_runtime() -> Result<(), OarOcrError> {
Ok(())
}
pub(crate) fn onnx_runtime_library_path() -> PathBuf {
std::env::var_os(ONNX_RUNTIME_LIBRARY_ENV)
.filter(|path| !path.is_empty())
.map(PathBuf::from)
.unwrap_or_else(default_onnx_runtime_library)
}
fn default_onnx_runtime_library() -> PathBuf {
#[cfg(target_os = "windows")]
const NAME: &str = "onnxruntime.dll";
@@ -355,6 +376,14 @@ mod tests {
use super::*;
use crate::vision::PageTransform;
#[test]
fn cpu_session_budget_is_bounded_for_small_ocr_models() {
let config = ocr_session_config();
assert!((1..=4).contains(&config.intra_threads.unwrap()));
assert_eq!(config.inter_threads, Some(1));
assert_eq!(config.parallel_execution, Some(false));
}
fn page(format: RenderPixelFormat, stride: usize, pixels: Vec<u8>) -> RenderedPage {
let transform =
PageTransform::from_corners(2, 2, (0.0, 2.0), (2.0, 2.0), (0.0, 0.0)).unwrap();
+116 -9
View File
@@ -1,13 +1,15 @@
//! One-call native extraction and OCR pipeline.
use std::collections::BTreeSet;
use std::path::Path;
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,
@@ -15,6 +17,24 @@ use super::{
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 {
@@ -220,16 +240,10 @@ pub fn process_pdf_with_ocr_mem(
// 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)?;
let engine = cached_ocr_engine(&options.ocr)?;
run_ocr_pages(
&renderer,
&engine,
engine.as_ref(),
buffer,
&routed,
options.password.as_deref(),
@@ -268,6 +282,84 @@ pub fn process_pdf_with_ocr_mem(
})
}
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() {
@@ -335,6 +427,21 @@ pub enum OcrPipelineError {
mod tests {
use super::*;
#[test]
fn cache_path_is_stable_when_a_relative_directory_is_created() {
let current = std::fs::canonicalize(std::env::current_dir().unwrap()).unwrap();
let temporary = tempfile::tempdir_in(&current).unwrap();
let relative_root = temporary.path().strip_prefix(&current).unwrap();
let relative = relative_root.join("models").join("revision");
let before = normalized_cache_path(relative.clone());
std::fs::create_dir_all(&relative).unwrap();
let after = normalized_cache_path(relative);
assert!(before.is_absolute());
assert_eq!(before, after);
}
#[test]
fn off_mode_extracts_native_text_without_runtime_side_effects() {
let bytes = std::fs::read("tests/fixtures/thermo-freon12.pdf").unwrap();
+4 -1
View File
@@ -121,7 +121,9 @@ fn complete_ocr_pipeline_routes_and_assembles_a_scanned_fixture() {
.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();
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());
@@ -130,6 +132,7 @@ fn complete_ocr_pipeline_routes_and_assembles_a_scanned_fixture() {
result.pages[0].provenance.ocr_model.as_ref().unwrap().name,
PP_OCR_V6_SMALL.id
);
assert_eq!(repeated.markdown, result.markdown);
}
fn recognize(