Compare commits

...
9 changed files with 2197 additions and 1 deletions
+14
View File
@@ -49,6 +49,17 @@ ttf-parser = "0.25"
lopdf = { version = "0.42.0", features = ["rayon"] }
rayon = "1.10"
env_logger = "0.11"
# Optional native page rendering for OCR pipelines. PDFium is loaded at
# runtime, so enabling this feature does not link or download a native library.
firecrawl-pdfium = { version = "0.1.0", optional = true }
# Small support crates used only by the opt-in model cache. Model files remain
# external and are never embedded in pdf-inspector artifacts.
dirs = { version = "6.0", optional = true }
fs2 = { version = "0.4", optional = true }
sha2 = { version = "0.11", optional = true }
[target.'cfg(all(windows, not(target_arch = "wasm32")))'.dependencies]
windows-sys = { version = "0.61", features = ["Win32_Storage_FileSystem"], optional = true }
# Browser builds use JavaScript randomness for encrypted PDFs and embed the
# bundled CMaps because there is no filesystem at runtime.
@@ -62,6 +73,9 @@ tempfile = "3.3"
[features]
default = []
python = ["pyo3"]
vision = []
model-cache = ["vision", "dep:dirs", "dep:fs2", "dep:sha2", "dep:windows-sys"]
render-pdfium = ["vision", "dep:firecrawl-pdfium"]
[[bin]]
name = "pdf2md"
+81 -1
View File
@@ -1,6 +1,6 @@
# pdf-inspector
Fast PDF classification and text extraction. Detects whether a PDF is text-based or scanned, extracts text with position awareness, and converts to clean Markdown — all without OCR. Pure Rust, no ML models, no external services; the only PDF dependency is [lopdf](https://crates.io/crates/lopdf). Also available for [Python](https://pypi.org/project/pdf-inspector/) and [Node.js](https://www.npmjs.com/package/@firecrawl/pdf-inspector).
Fast PDF classification and text extraction. Detects whether a PDF is text-based or scanned, extracts text with position awareness, and converts to clean Markdown — all without OCR. The default build is pure Rust, has no ML models or external services, and uses [lopdf](https://crates.io/crates/lopdf) for PDF parsing. Also available for [Python](https://pypi.org/project/pdf-inspector/) and [Node.js](https://www.npmjs.com/package/@firecrawl/pdf-inspector/).
Built by [Firecrawl](https://firecrawl.dev) to handle text-based PDFs locally in under 200ms, skipping expensive OCR services for the ~54% of PDFs that don't need them.
@@ -117,6 +117,86 @@ let bytes = std::fs::read("document.pdf")?;
let result = process_pdf_mem(&bytes)?;
```
### Vision extension contracts
The native-only `vision` feature exposes the stable seam used by OCR
integrations without selecting or embedding an inference runtime. The
separate `model-cache` feature adds pinned artifact management:
- `PageRenderer`, `OcrEngine`, and `LayoutEngine` traits;
- renderer-neutral owned page buffers and affine pixel↔PDF transforms;
- `OcrOptions` and opt-in `Off`/`Auto`/`Force` routing modes;
- positioned OCR/layout results and per-page provenance types; and
- a versioned PP-OCRv6 Small manifest with checksum-verified, locked, atomic
model-cache installation and explicit offline-directory overrides.
```toml
[dependencies]
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.
```rust
use pdf_inspector::vision::{
ModelDownloadPolicy, ModelStore, OcrMode, OcrOptions, PP_OCR_V6_SMALL,
};
let ocr = OcrOptions::new()
.mode(OcrMode::Auto)
.model_directory("/opt/firecrawl/models/pp-ocrv6-small")
.model_downloads(ModelDownloadPolicy::Offline);
// Verifies exact sizes and SHA-256 digests before an engine opens the files.
let models = ModelStore::from_options(&ocr)?.resolve(&PP_OCR_V6_SMALL)?;
println!("using {} at {}", models.manifest_id(), models.revision());
```
### Optional native page rendering
The `render-pdfium` feature adds a native-only page renderer backed by
[`firecrawl-pdfium`](https://crates.io/crates/firecrawl-pdfium). It is the
rendering boundary for OCR pipelines; enabling it does not include an OCR
model or change the existing extraction functions. It implies `vision`,
and `PdfiumRenderer` implements the renderer-neutral `PageRenderer` trait.
```toml
[dependencies]
pdf-inspector = { version = "1", features = ["render-pdfium"] }
```
PDFium is loaded at runtime. Set `PDFIUM_LIB_PATH`, place its shared library
next to the executable, or use another discovery route supported by
`firecrawl-pdfium`.
```rust
use pdf_inspector::vision::{PdfiumRenderer, RenderOptions};
let renderer = PdfiumRenderer::load()?;
let bytes = std::fs::read("document.pdf")?;
let pages = renderer.render_pages(
&bytes,
&[1, 3], // 1-indexed, matching pages_needing_ocr
None, // optional PDF password
&RenderOptions::new().dpi(150.0),
)?;
for page in pages {
// Owned RGB pixels can leave the PDFium critical section and be sent to
// an OCR worker. OCR pixel boxes can be mapped back to PDF coordinates.
let rect = page.pixel_rect_to_pdf_rect(20.0, 30.0, 100.0, 24.0);
println!("page {}: {}x{}, rect={rect:?}", page.page(), page.width(), page.height());
}
```
Browser WASM remains on the default text-only path and does not expose native
PDFium rendering.
Extract per-page Markdown (one string per page, plus document-wide layout
metadata):
+1
View File
@@ -43,6 +43,7 @@ mod text_quality;
pub mod text_utils;
pub mod tounicode;
pub mod types;
pub mod vision;
pub use detector::{
detect_pdf_type, detect_pdf_type_mem, detect_pdf_type_mem_with_config,
+414
View File
@@ -0,0 +1,414 @@
//! Public contracts between rendering, OCR, layout, and orchestration.
use std::error::Error;
use std::path::PathBuf;
use super::{RenderOptions, RenderedPage};
/// Selects when OCR may run.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
#[non_exhaustive]
pub enum OcrMode {
/// Never run OCR. This is the default and preserves existing behavior.
#[default]
Off,
/// Run OCR only on pages selected by pdf-inspector's OCR routing signals.
Auto,
/// Run OCR on every selected page, including pages with native text.
Force,
}
/// Resource/quality profile for the OCR engine.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
#[non_exhaustive]
pub enum OcrProfile {
/// Lowest latency and memory footprint.
Edge,
/// OCR-oriented balance of quality and CPU cost.
#[default]
Balanced,
/// Highest quality within the lightweight model family.
Quality,
}
/// Controls whether missing model artifacts may be fetched.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
#[non_exhaustive]
pub enum ModelDownloadPolicy {
/// Fetch a pinned artifact only after OCR has actually been selected.
#[default]
IfMissing,
/// Never access the network; require an override or a warm model cache.
Offline,
}
/// OCR engine configuration independent of a particular runtime.
#[derive(Debug, Clone, PartialEq)]
pub struct OcrOptions {
/// Page-level routing behavior.
pub mode: OcrMode,
/// Local quality/resource profile.
pub profile: OcrProfile,
/// Drop recognition spans below this confidence threshold.
pub minimum_confidence: f32,
/// Optional language hints understood by the selected engine.
pub languages: Vec<String>,
/// Optional directory containing an offline model set.
pub model_directory: Option<PathBuf>,
/// Whether a missing pinned artifact may be downloaded.
pub model_downloads: ModelDownloadPolicy,
}
impl Default for OcrOptions {
fn default() -> Self {
Self {
mode: OcrMode::Off,
profile: OcrProfile::Balanced,
minimum_confidence: 0.0,
languages: Vec::new(),
model_directory: None,
model_downloads: ModelDownloadPolicy::IfMissing,
}
}
}
impl OcrOptions {
/// Creates OCR options with OCR disabled.
pub fn new() -> Self {
Self::default()
}
/// Sets page-level OCR routing.
pub fn mode(mut self, mode: OcrMode) -> Self {
self.mode = mode;
self
}
/// Sets the local resource/quality profile.
pub fn profile(mut self, profile: OcrProfile) -> Self {
self.profile = profile;
self
}
/// Sets the minimum accepted recognition confidence.
pub fn minimum_confidence(mut self, minimum_confidence: f32) -> Self {
self.minimum_confidence = minimum_confidence;
self
}
/// Replaces the language hints passed to the OCR engine.
pub fn languages(mut self, languages: impl IntoIterator<Item = impl Into<String>>) -> Self {
self.languages = languages.into_iter().map(Into::into).collect();
self
}
/// Uses an explicit model directory, suitable for offline packaging.
pub fn model_directory(mut self, directory: impl Into<PathBuf>) -> Self {
self.model_directory = Some(directory.into());
self
}
/// Sets the missing-model download policy.
pub fn model_downloads(mut self, policy: ModelDownloadPolicy) -> Self {
self.model_downloads = policy;
self
}
}
/// Configuration for an optional learned layout engine.
///
/// Layout inference is disabled by default. Existing deterministic layout,
/// table, and Markdown logic remains the assembly path when this is disabled.
#[derive(Debug, Clone, PartialEq)]
pub struct LayoutOptions {
/// Whether the learned layout extension may run.
pub enabled: bool,
/// Drop layout regions below this confidence threshold.
pub minimum_confidence: f32,
/// Optional directory containing an offline layout model set.
pub model_directory: Option<PathBuf>,
}
impl Default for LayoutOptions {
fn default() -> Self {
Self {
enabled: false,
minimum_confidence: 0.0,
model_directory: None,
}
}
}
impl LayoutOptions {
/// Creates layout options with learned layout disabled.
pub fn new() -> Self {
Self::default()
}
/// Enables or disables learned layout inference.
pub fn enabled(mut self, enabled: bool) -> Self {
self.enabled = enabled;
self
}
/// Sets the minimum accepted region confidence.
pub fn minimum_confidence(mut self, minimum_confidence: f32) -> Self {
self.minimum_confidence = minimum_confidence;
self
}
/// Uses an explicit layout model directory.
pub fn model_directory(mut self, directory: impl Into<PathBuf>) -> Self {
self.model_directory = Some(directory.into());
self
}
}
/// A point in bitmap space, measured from the top-left in pixels.
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct ImagePoint {
/// Horizontal pixel coordinate.
pub x: f32,
/// Vertical pixel coordinate, increasing downward.
pub y: f32,
}
impl ImagePoint {
/// Creates a bitmap-space point.
pub fn new(x: f32, y: f32) -> Self {
Self { x, y }
}
}
/// Four-point polygon in bitmap coordinates.
#[derive(Debug, Clone, Copy, Default, PartialEq)]
pub struct ImageQuad {
/// Polygon points in engine-provided order.
pub points: [ImagePoint; 4],
}
impl ImageQuad {
/// Creates a four-point bitmap polygon.
pub fn new(points: [ImagePoint; 4]) -> Self {
Self { points }
}
}
/// Stable identity for an inference model used in output provenance.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelIdentity {
/// Model family/name, for example `pp-ocrv6-small`.
pub name: String,
/// Immutable model or artifact-set revision.
pub revision: String,
}
impl ModelIdentity {
/// Creates a model identity.
pub fn new(name: impl Into<String>, revision: impl Into<String>) -> Self {
Self {
name: name.into(),
revision: revision.into(),
}
}
}
/// One positioned OCR recognition result in bitmap coordinates.
#[derive(Debug, Clone, PartialEq)]
pub struct OcrSpan {
/// Recognized text.
pub text: String,
/// Detection polygon in the original rendered page's pixel space.
pub polygon: ImageQuad,
/// Recognition confidence in the inclusive range 01.
pub confidence: f32,
/// Optional text-line orientation in clockwise degrees.
pub orientation_degrees: Option<f32>,
}
/// OCR output for one 1-indexed page.
#[derive(Debug, Clone, PartialEq)]
pub struct OcrPage {
/// 1-indexed PDF page number.
pub page: u32,
/// Positioned recognition spans.
pub spans: Vec<OcrSpan>,
/// Mean confidence across accepted spans, when available.
pub mean_confidence: Option<f32>,
/// Exact model identity used for this result.
pub model: ModelIdentity,
/// OCR wall time for this page.
pub processing_time_ms: u64,
/// Non-fatal engine warnings.
pub warnings: Vec<String>,
}
/// Normalized semantic class emitted by a learned layout engine.
#[derive(Debug, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum LayoutRegionKind {
/// Body or other prose text.
Text,
/// Document heading or title.
Heading,
/// Table region.
Table,
/// Figure/image region.
Figure,
/// Figure or table caption.
Caption,
/// Header/footer/page furniture.
Furniture,
/// Model-specific class retained without changing the common taxonomy.
Other(String),
}
/// One learned layout region in bitmap coordinates.
#[derive(Debug, Clone, PartialEq)]
pub struct LayoutRegion {
/// Normalized semantic class.
pub kind: LayoutRegionKind,
/// Region polygon in the original rendered page's pixel space.
pub polygon: ImageQuad,
/// Model confidence in the inclusive range 01.
pub confidence: f32,
/// Optional model-provided reading-order position.
pub reading_order: Option<u32>,
}
/// Learned layout output for one 1-indexed page.
#[derive(Debug, Clone, PartialEq)]
pub struct LayoutPage {
/// 1-indexed PDF page number.
pub page: u32,
/// Semantic regions.
pub regions: Vec<LayoutRegion>,
/// Exact model identity used for this result.
pub model: ModelIdentity,
/// Layout inference wall time for this page.
pub processing_time_ms: u64,
/// Non-fatal engine warnings.
pub warnings: Vec<String>,
}
/// How final page content was sourced.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[non_exhaustive]
pub enum PageContentSource {
/// Trusted native PDF text only.
Native,
/// OCR output only.
Ocr,
/// Native and OCR spans were fused.
Fused,
}
/// Per-page local processing timings.
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
pub struct VisionTimings {
/// Rasterization wall time.
pub render_ms: u64,
/// OCR wall time.
pub ocr_ms: u64,
/// Optional learned layout wall time.
pub layout_ms: u64,
/// Native/OCR fusion and assembly wall time.
pub assembly_ms: u64,
}
/// Source and model metadata retained for one processed page.
#[derive(Debug, Clone, PartialEq)]
pub struct PageProvenance {
/// 1-indexed PDF page number.
pub page: u32,
/// Final page-content source.
pub source: PageContentSource,
/// OCR model, when OCR ran.
pub ocr_model: Option<ModelIdentity>,
/// Learned layout model, when layout inference ran.
pub layout_model: Option<ModelIdentity>,
/// Render resolution used for local vision.
pub render_dpi: Option<f32>,
/// Mean accepted OCR confidence, when available.
pub ocr_confidence: Option<f32>,
/// Stage timings.
pub timings: VisionTimings,
/// Non-fatal warnings surfaced to downstream users.
pub warnings: Vec<String>,
/// True when this lightweight local path detected a case better suited to
/// Firecrawl's hosted document pipeline.
pub hosted_recommended: bool,
}
/// Converts selected PDF pages into renderer-neutral owned bitmaps.
pub trait PageRenderer: Send + Sync {
/// Renderer-specific failure type.
type Error: Error + Send + Sync + 'static;
/// Renders selected 1-indexed pages in the same order as `pages`.
fn render_pages(
&self,
pdf_bytes: &[u8],
pages: &[u32],
password: Option<&str>,
options: &RenderOptions,
) -> Result<Vec<RenderedPage>, Self::Error>;
}
/// Recognizes positioned text from rendered pages.
pub trait OcrEngine: Send + Sync {
/// Engine-specific failure type.
type Error: Error + Send + Sync + 'static;
/// Exact model identity used by this engine instance.
fn model(&self) -> &ModelIdentity;
/// Recognizes pages in batch and returns results in input order.
fn recognize(
&self,
pages: &[RenderedPage],
options: &OcrOptions,
) -> Result<Vec<OcrPage>, Self::Error>;
}
/// Optional learned semantic layout extension.
pub trait LayoutEngine: Send + Sync {
/// Engine-specific failure type.
type Error: Error + Send + Sync + 'static;
/// Exact model identity used by this engine instance.
fn model(&self) -> &ModelIdentity;
/// Analyzes rendered pages, optionally using their OCR spans.
fn analyze(
&self,
pages: &[RenderedPage],
ocr: &[OcrPage],
options: &LayoutOptions,
) -> Result<Vec<LayoutPage>, Self::Error>;
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn ocr_defaults_never_enable_recognition() {
let options = OcrOptions::default();
assert_eq!(options.mode, OcrMode::Off);
}
#[test]
fn offline_model_override_is_explicit() {
let options = OcrOptions::new()
.mode(OcrMode::Auto)
.model_directory("/models/pp-ocr")
.model_downloads(ModelDownloadPolicy::Offline);
assert_eq!(options.mode, OcrMode::Auto);
assert_eq!(options.model_downloads, ModelDownloadPolicy::Offline);
assert_eq!(
options.model_directory,
Some(PathBuf::from("/models/pp-ocr"))
);
}
}
+37
View File
@@ -0,0 +1,37 @@
//! Optional native vision primitives used by OCR pipelines.
//!
//! The existing lopdf extractor remains the default path. Native page
//! rendering is available only with the `render-pdfium` feature. Engine
//! contracts are available with `vision`, while checksum-verified model
//! resolution is a separate `model-cache` feature. These remain separate so
//! browser WASM, text-only consumers, and renderer-only users take on no model
//! management dependencies.
#[cfg(all(feature = "vision", not(target_arch = "wasm32")))]
mod contracts;
#[cfg(all(feature = "model-cache", not(target_arch = "wasm32")))]
mod models;
#[cfg(all(feature = "vision", not(target_arch = "wasm32")))]
mod render;
#[cfg(all(feature = "render-pdfium", not(target_arch = "wasm32")))]
mod pdfium;
#[cfg(all(feature = "vision", not(target_arch = "wasm32")))]
pub use contracts::{
ImagePoint, ImageQuad, LayoutEngine, LayoutOptions, LayoutPage, LayoutRegion, LayoutRegionKind,
ModelDownloadPolicy, ModelIdentity, OcrEngine, OcrMode, OcrOptions, OcrPage, OcrProfile,
OcrSpan, PageContentSource, PageProvenance, PageRenderer, VisionTimings,
};
#[cfg(all(feature = "model-cache", not(target_arch = "wasm32")))]
pub use models::{
ModelArtifact, ModelArtifactKind, ModelManifest, ModelPaths, ModelStore, ModelStoreError,
PP_OCR_V6_SMALL,
};
#[cfg(all(feature = "vision", not(target_arch = "wasm32")))]
pub use render::{
PagePoint, PageTransform, RenderBufferError, RenderOptions, RenderPixelFormat, RenderedPage,
};
#[cfg(all(feature = "render-pdfium", not(target_arch = "wasm32")))]
pub use pdfium::{PdfiumRenderer, RenderError};
+767
View File
@@ -0,0 +1,767 @@
//! Versioned model manifests and a checksum-verified local cache.
use std::collections::{BTreeMap, BTreeSet};
use std::ffi::OsStr;
use std::fs::{self, File, OpenOptions};
use std::io::{self, Read, Write};
use std::path::{Component, Path, PathBuf};
use std::sync::atomic::{AtomicU64, Ordering};
use std::time::{SystemTime, UNIX_EPOCH};
use fs2::FileExt;
use sha2::{Digest, Sha256};
use thiserror::Error;
use super::OcrOptions;
/// Environment variable overriding the default local model cache.
pub const MODEL_CACHE_ENV: &str = "PDF_INSPECTOR_MODEL_CACHE";
/// Role of an artifact within a local model set.
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
#[non_exhaustive]
pub enum ModelArtifactKind {
/// Text detection ONNX graph.
TextDetection,
/// Text recognition ONNX graph.
TextRecognition,
/// Recognition character dictionary.
CharacterDictionary,
/// Learned document-layout ONNX graph.
Layout,
}
/// One immutable file in a model manifest.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ModelArtifact {
/// Artifact role.
pub kind: ModelArtifactKind,
/// Cache filename without directory components.
pub filename: &'static str,
/// Canonical HTTPS download location.
pub url: &'static str,
/// Lowercase SHA-256 digest.
pub sha256: &'static str,
/// Exact expected file size.
pub size: u64,
}
/// Versioned set of files required by one model configuration.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct ModelManifest {
/// Manifest schema version.
pub schema_version: u32,
/// Stable model-set identifier.
pub id: &'static str,
/// Immutable upstream artifact revision.
pub revision: &'static str,
/// Required artifacts.
pub artifacts: &'static [ModelArtifact],
}
const PP_OCR_V6_SMALL_ARTIFACTS: &[ModelArtifact] = &[
ModelArtifact {
kind: ModelArtifactKind::TextDetection,
filename: "pp-ocrv6_small_det.onnx",
url: "https://github.com/GreatV/oar-ocr/releases/download/v0.7.0/pp-ocrv6_small_det.onnx",
sha256: "d73e0058b7a8086bbd57f3d10b8bcd4ff95363f67e06e2762b5e814fe9c9410e",
size: 9_880_512,
},
ModelArtifact {
kind: ModelArtifactKind::TextRecognition,
filename: "pp-ocrv6_small_rec.onnx",
url: "https://github.com/GreatV/oar-ocr/releases/download/v0.7.0/pp-ocrv6_small_rec.onnx",
sha256: "5435fd747c9e0efe15a96d0b378d5bd157e9492ed8fd80edf08f30d02fa24634",
size: 21_159_378,
},
ModelArtifact {
kind: ModelArtifactKind::CharacterDictionary,
filename: "ppocrv6_dict.txt",
url: "https://github.com/GreatV/oar-ocr/releases/download/v0.7.0/ppocrv6_dict.txt",
sha256: "b5f2bfe2bdd9448429e3e82b51c789775d9b42f2403d082b00662eb77e401c5d",
size: 74_947,
},
];
/// Pinned PP-OCRv6 Small detection/recognition model set.
///
/// The artifact hashes match the registry shipped by `oar-ocr-core` 0.9.1;
/// the revision identifies the upstream release that owns the files.
pub const PP_OCR_V6_SMALL: ModelManifest = ModelManifest {
schema_version: 1,
id: "pp-ocrv6-small",
revision: "oar-ocr-v0.7.0",
artifacts: PP_OCR_V6_SMALL_ARTIFACTS,
};
/// Resolved, verified filesystem paths for one model manifest.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ModelPaths {
manifest_id: String,
revision: String,
artifacts: BTreeMap<ModelArtifactKind, PathBuf>,
}
impl ModelPaths {
/// Stable model-set identifier.
pub fn manifest_id(&self) -> &str {
&self.manifest_id
}
/// Immutable artifact revision.
pub fn revision(&self) -> &str {
&self.revision
}
/// Verified path for an artifact role.
pub fn get(&self, kind: ModelArtifactKind) -> Option<&Path> {
self.artifacts.get(&kind).map(PathBuf::as_path)
}
/// Iterates over verified artifact paths.
pub fn iter(&self) -> impl Iterator<Item = (ModelArtifactKind, &Path)> {
self.artifacts
.iter()
.map(|(&kind, path)| (kind, path.as_path()))
}
}
/// Checksum-verified model cache with an optional offline directory override.
///
/// This type never accesses the network. [`resolve`](Self::resolve) verifies a
/// warm cache or explicit directory, and [`install`](Self::install) atomically
/// installs bytes supplied by a higher-level downloader. Keeping acquisition
/// separate makes offline behavior enforceable and straightforward to test.
#[derive(Debug, Clone)]
pub struct ModelStore {
cache_root: PathBuf,
override_root: Option<PathBuf>,
}
impl ModelStore {
/// Creates a model store rooted at an explicit cache directory.
pub fn new(cache_root: impl Into<PathBuf>) -> Self {
Self {
cache_root: cache_root.into(),
override_root: None,
}
}
/// Builds a store from OCR options and the platform cache directory.
///
/// `PDF_INSPECTOR_MODEL_CACHE` overrides the platform default. An explicit
/// [`OcrOptions::model_directory`] replaces the managed cache at resolve
/// time so offline packaging is deterministic.
pub fn from_options(options: &OcrOptions) -> Result<Self, ModelStoreError> {
let cache_root = match std::env::var_os(MODEL_CACHE_ENV) {
Some(path) if !path.is_empty() => PathBuf::from(path),
_ => dirs::cache_dir()
.ok_or(ModelStoreError::CacheDirectoryUnavailable)?
.join("pdf-inspector")
.join("models"),
};
Ok(Self {
cache_root,
override_root: options.model_directory.clone(),
})
}
/// Checks an explicit offline model directory before the managed cache.
pub fn override_root(mut self, root: impl Into<PathBuf>) -> Self {
self.override_root = Some(root.into());
self
}
/// Managed cache root.
pub fn cache_root(&self) -> &Path {
&self.cache_root
}
/// 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 mut artifacts = BTreeMap::new();
for artifact in manifest.artifacts {
let path = root.join(artifact.filename);
verify_artifact(&path, artifact)?;
artifacts.insert(artifact.kind, path);
}
Ok(ModelPaths {
manifest_id: manifest.id.to_string(),
revision: manifest.revision.to_string(),
artifacts,
})
}
/// Atomically installs one artifact from a reader after validating its
/// exact size and SHA-256 digest.
///
/// A cross-process file lock serializes installs of the same artifact.
/// Already-valid cached files are reused without consuming the reader.
pub fn install(
&self,
manifest: &ModelManifest,
kind: ModelArtifactKind,
mut reader: impl Read,
) -> Result<PathBuf, ModelStoreError> {
validate_manifest(manifest)?;
let artifact = manifest
.artifacts
.iter()
.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
}
fn manifest_cache_root(&self, manifest: &ModelManifest) -> PathBuf {
self.cache_root.join(manifest.id).join(manifest.revision)
}
}
/// Failures while validating or installing model artifacts.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum ModelStoreError {
/// No platform cache location is available.
#[error("no platform cache directory is available; set {MODEL_CACHE_ENV}")]
CacheDirectoryUnavailable,
/// The static/custom manifest is malformed.
#[error("invalid model manifest: {0}")]
InvalidManifest(String),
/// A requested artifact role is not in the manifest.
#[error("artifact role {kind:?} is not present in the model manifest")]
ArtifactNotInManifest {
/// Missing role.
kind: ModelArtifactKind,
},
/// A required artifact is not installed.
#[error("model artifact is missing: {path}")]
MissingArtifact {
/// Expected local path.
path: PathBuf,
/// Canonical download URL for an allowed higher-level fetcher.
download_url: &'static str,
},
/// Artifact byte size differs from the manifest.
#[error("model artifact {path} has {actual} bytes; expected {expected}")]
SizeMismatch {
/// Artifact path.
path: PathBuf,
/// Expected byte count.
expected: u64,
/// Actual byte count.
actual: u64,
},
/// Artifact digest differs from the manifest.
#[error("model artifact checksum mismatch at {path}: expected {expected}, got {actual}")]
ChecksumMismatch {
/// Artifact path.
path: PathBuf,
/// Expected lowercase SHA-256.
expected: &'static str,
/// Actual lowercase SHA-256.
actual: String,
},
/// Filesystem or stream I/O failed.
#[error("model cache I/O failed at {path}: {source}")]
Io {
/// Path involved in the operation.
path: PathBuf,
/// Underlying I/O error.
#[source]
source: io::Error,
},
}
fn validate_manifest(manifest: &ModelManifest) -> Result<(), ModelStoreError> {
if manifest.schema_version != 1 {
return Err(ModelStoreError::InvalidManifest(format!(
"unsupported schema version {}",
manifest.schema_version
)));
}
if manifest.id.is_empty() || manifest.revision.is_empty() || manifest.artifacts.is_empty() {
return Err(ModelStoreError::InvalidManifest(
"id, revision, and artifacts must be non-empty".to_string(),
));
}
for (field, value) in [("id", manifest.id), ("revision", manifest.revision)] {
if !is_single_normal_path_component(value) {
return Err(ModelStoreError::InvalidManifest(format!(
"{field} must be a single path component: {value}"
)));
}
}
let mut kinds = BTreeSet::new();
let mut filenames = BTreeSet::new();
for artifact in manifest.artifacts {
if !is_single_normal_path_component(artifact.filename) {
return Err(ModelStoreError::InvalidManifest(format!(
"artifact filename must be a single path component: {}",
artifact.filename
)));
}
if artifact.size == 0 {
return Err(ModelStoreError::InvalidManifest(format!(
"artifact {} has zero size",
artifact.filename
)));
}
if artifact.sha256.len() != 64
|| !artifact
.sha256
.bytes()
.all(|byte| byte.is_ascii_hexdigit() && !byte.is_ascii_uppercase())
{
return Err(ModelStoreError::InvalidManifest(format!(
"artifact {} has an invalid SHA-256",
artifact.filename
)));
}
if !artifact.url.starts_with("https://") {
return Err(ModelStoreError::InvalidManifest(format!(
"artifact {} must use HTTPS",
artifact.filename
)));
}
if !kinds.insert(artifact.kind) || !filenames.insert(artifact.filename) {
return Err(ModelStoreError::InvalidManifest(format!(
"artifact {} duplicates a kind or filename",
artifact.filename
)));
}
}
Ok(())
}
fn is_single_normal_path_component(value: &str) -> bool {
if value.is_empty() || Path::new(value).file_name() != Some(OsStr::new(value)) {
return false;
}
let mut components = Path::new(value).components();
matches!(components.next(), Some(Component::Normal(_))) && components.next().is_none()
}
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 {
path: root.to_path_buf(),
source,
})? {
let entry = entry.map_err(|source| ModelStoreError::Io {
path: root.to_path_buf(),
source,
})?;
let name = entry.file_name();
let name = name.to_string_lossy();
if name.starts_with(&prefix) && name.ends_with(".part") {
let path = entry.path();
fs::remove_file(&path).map_err(|source| ModelStoreError::Io { path, source })?;
}
}
Ok(())
}
fn create_temporary_file(root: &Path, filename: &str) -> Result<(PathBuf, File), ModelStoreError> {
static TEMP_COUNTER: AtomicU64 = AtomicU64::new(0);
for _ in 0..16 {
let sequence = TEMP_COUNTER.fetch_add(1, Ordering::Relaxed);
let timestamp = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap_or_default()
.as_nanos();
let path = root.join(format!(
".{filename}.{}.{}.{}.part",
std::process::id(),
timestamp,
sequence
));
match OpenOptions::new().create_new(true).write(true).open(&path) {
Ok(file) => return Ok((path, file)),
Err(source) if source.kind() == io::ErrorKind::AlreadyExists => continue,
Err(source) => return Err(ModelStoreError::Io { path, source }),
}
}
let path = root.join(format!(".{filename}.part"));
Err(ModelStoreError::Io {
path,
source: io::Error::new(
io::ErrorKind::AlreadyExists,
"could not allocate a unique model install file",
),
})
}
#[cfg(not(windows))]
fn replace_file_atomic(source: &Path, target: &Path) -> io::Result<()> {
fs::rename(source, target)
}
#[cfg(windows)]
fn replace_file_atomic(source: &Path, target: &Path) -> io::Result<()> {
use std::os::windows::ffi::OsStrExt;
use windows_sys::Win32::Storage::FileSystem::{
MoveFileExW, MOVEFILE_REPLACE_EXISTING, MOVEFILE_WRITE_THROUGH,
};
let source: Vec<u16> = source.as_os_str().encode_wide().chain(Some(0)).collect();
let target: Vec<u16> = target.as_os_str().encode_wide().chain(Some(0)).collect();
let result = unsafe {
MoveFileExW(
source.as_ptr(),
target.as_ptr(),
MOVEFILE_REPLACE_EXISTING | MOVEFILE_WRITE_THROUGH,
)
};
if result == 0 {
Err(io::Error::last_os_error())
} else {
Ok(())
}
}
fn verify_artifact(path: &Path, artifact: &ModelArtifact) -> Result<(), ModelStoreError> {
let metadata = match fs::metadata(path) {
Ok(metadata) => metadata,
Err(source) if source.kind() == io::ErrorKind::NotFound => {
return Err(ModelStoreError::MissingArtifact {
path: path.to_path_buf(),
download_url: artifact.url,
});
}
Err(source) => {
return Err(ModelStoreError::Io {
path: path.to_path_buf(),
source,
});
}
};
if metadata.len() != artifact.size {
return Err(ModelStoreError::SizeMismatch {
path: path.to_path_buf(),
expected: artifact.size,
actual: metadata.len(),
});
}
let mut file = File::open(path).map_err(|source| ModelStoreError::Io {
path: path.to_path_buf(),
source,
})?;
let mut hasher = Sha256::new();
io::copy(&mut file, &mut DigestWriter(&mut hasher)).map_err(|source| ModelStoreError::Io {
path: path.to_path_buf(),
source,
})?;
let digest = digest_hex(hasher.finalize());
validate_size_and_hash(artifact, metadata.len(), &digest, path)
}
fn copy_and_hash(reader: &mut impl Read, writer: &mut impl Write) -> io::Result<(u64, String)> {
let mut hasher = Sha256::new();
let mut buffer = [0_u8; 64 * 1024];
let mut size = 0_u64;
loop {
let read = reader.read(&mut buffer)?;
if read == 0 {
break;
}
writer.write_all(&buffer[..read])?;
hasher.update(&buffer[..read]);
size = size
.checked_add(read as u64)
.ok_or_else(|| io::Error::other("model artifact size overflow"))?;
}
Ok((size, digest_hex(hasher.finalize())))
}
fn digest_hex(digest: impl AsRef<[u8]>) -> String {
const HEX: &[u8; 16] = b"0123456789abcdef";
let bytes = digest.as_ref();
let mut encoded = String::with_capacity(bytes.len() * 2);
for &byte in bytes {
encoded.push(HEX[(byte >> 4) as usize] as char);
encoded.push(HEX[(byte & 0x0f) as usize] as char);
}
encoded
}
fn validate_size_and_hash(
artifact: &ModelArtifact,
size: u64,
digest: &str,
path: &Path,
) -> Result<(), ModelStoreError> {
if size != artifact.size {
return Err(ModelStoreError::SizeMismatch {
path: path.to_path_buf(),
expected: artifact.size,
actual: size,
});
}
if digest != artifact.sha256 {
return Err(ModelStoreError::ChecksumMismatch {
path: path.to_path_buf(),
expected: artifact.sha256,
actual: digest.to_string(),
});
}
Ok(())
}
struct DigestWriter<'a>(&'a mut Sha256);
impl Write for DigestWriter<'_> {
fn write(&mut self, buffer: &[u8]) -> io::Result<usize> {
self.0.update(buffer);
Ok(buffer.len())
}
fn flush(&mut self) -> io::Result<()> {
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
const TEST_ARTIFACTS: &[ModelArtifact] = &[ModelArtifact {
kind: ModelArtifactKind::CharacterDictionary,
filename: "hello.txt",
url: "https://example.com/hello.txt",
sha256: "2cf24dba5fb0a30e26e83b2ac5b9e29e1b161e5c1fa7425e73043362938b9824",
size: 5,
}];
const TEST_MANIFEST: ModelManifest = ModelManifest {
schema_version: 1,
id: "test-model",
revision: "v1",
artifacts: TEST_ARTIFACTS,
};
#[test]
fn pinned_pp_ocr_manifest_is_well_formed() {
validate_manifest(&PP_OCR_V6_SMALL).unwrap();
assert_eq!(PP_OCR_V6_SMALL.artifacts.len(), 3);
}
#[test]
fn installs_and_resolves_verified_artifact() {
let temp = tempfile::tempdir().unwrap();
let store = ModelStore::new(temp.path());
let installed = store
.install(
&TEST_MANIFEST,
ModelArtifactKind::CharacterDictionary,
&b"hello"[..],
)
.unwrap();
assert!(installed.ends_with("hello.txt"));
let resolved = store.resolve(&TEST_MANIFEST).unwrap();
assert_eq!(
resolved.get(ModelArtifactKind::CharacterDictionary),
Some(installed.as_path())
);
}
#[test]
fn rejected_install_does_not_poison_cache() {
let temp = tempfile::tempdir().unwrap();
let store = ModelStore::new(temp.path());
assert!(matches!(
store.install(
&TEST_MANIFEST,
ModelArtifactKind::CharacterDictionary,
&b"HELLO"[..],
),
Err(ModelStoreError::ChecksumMismatch { .. })
));
assert!(matches!(
store.resolve(&TEST_MANIFEST),
Err(ModelStoreError::MissingArtifact { .. })
));
}
#[test]
fn oversized_install_stops_after_one_extra_byte() {
let temp = tempfile::tempdir().unwrap();
let store = ModelStore::new(temp.path());
assert!(matches!(
store.install(
&TEST_MANIFEST,
ModelArtifactKind::CharacterDictionary,
&b"hello and far too much data"[..],
),
Err(ModelStoreError::SizeMismatch {
expected: 5,
actual: 6,
..
})
));
}
#[test]
fn install_replaces_invalid_cache_atomically() {
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();
fs::write(root.join("hello.txt"), b"HELLO").unwrap();
store
.install(
&TEST_MANIFEST,
ModelArtifactKind::CharacterDictionary,
&b"hello"[..],
)
.unwrap();
assert_eq!(fs::read(root.join("hello.txt")).unwrap(), b"hello");
}
#[test]
fn stale_partial_installs_are_swept_under_the_lock() {
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();
let stale = root.join(".hello.txt.123.0.part");
fs::write(&stale, b"stale").unwrap();
store
.install(
&TEST_MANIFEST,
ModelArtifactKind::CharacterDictionary,
&b"hello"[..],
)
.unwrap();
assert!(!stale.exists());
}
#[test]
fn manifest_paths_cannot_escape_the_cache() {
const BAD_ID: ModelManifest = ModelManifest {
id: "../escape",
..TEST_MANIFEST
};
const BAD_REVISION: ModelManifest = ModelManifest {
revision: "nested/revision",
..TEST_MANIFEST
};
const BAD_FILENAME_ARTIFACTS: &[ModelArtifact] = &[ModelArtifact {
filename: "hello.txt/",
..TEST_ARTIFACTS[0]
}];
const BAD_FILENAME: ModelManifest = ModelManifest {
artifacts: BAD_FILENAME_ARTIFACTS,
..TEST_MANIFEST
};
for manifest in [&BAD_ID, &BAD_REVISION, &BAD_FILENAME] {
assert!(matches!(
validate_manifest(manifest),
Err(ModelStoreError::InvalidManifest(_))
));
}
}
#[test]
fn explicit_override_is_verified_without_copying() {
let cache = tempfile::tempdir().unwrap();
let override_dir = tempfile::tempdir().unwrap();
fs::write(override_dir.path().join("hello.txt"), b"hello").unwrap();
let store = ModelStore::new(cache.path()).override_root(override_dir.path());
let resolved = store.resolve(&TEST_MANIFEST).unwrap();
assert_eq!(
resolved.get(ModelArtifactKind::CharacterDictionary),
Some(override_dir.path().join("hello.txt").as_path())
);
}
#[test]
fn concurrent_installs_converge_on_one_verified_file() {
let temp = tempfile::tempdir().unwrap();
let first_store = ModelStore::new(temp.path());
let second_store = first_store.clone();
let first = std::thread::spawn(move || {
first_store.install(
&TEST_MANIFEST,
ModelArtifactKind::CharacterDictionary,
&b"hello"[..],
)
});
let second = std::thread::spawn(move || {
second_store.install(
&TEST_MANIFEST,
ModelArtifactKind::CharacterDictionary,
&b"hello"[..],
)
});
let first = first.join().unwrap().unwrap();
let second = second.join().unwrap().unwrap();
assert_eq!(first, second);
assert_eq!(fs::read(first).unwrap(), b"hello");
}
}
+266
View File
@@ -0,0 +1,266 @@
//! PDFium-backed implementation of the renderer-neutral page contract.
use std::path::Path;
use firecrawl_pdfium::{Pdfium, PixelFormat, PixelPoint, RenderConfig};
use thiserror::Error;
use super::{
PageRenderer, PageTransform, RenderBufferError, RenderOptions, RenderPixelFormat, RenderedPage,
};
impl RenderPixelFormat {
fn pdfium_format(self) -> PixelFormat {
match self {
// PDFium produces BGR directly; `rendered_page_from_pdfium`
// swaps the red and blue channels in place.
Self::Rgb8 => PixelFormat::Bgr8,
Self::Rgba8 => PixelFormat::Rgba8,
Self::Gray8 => PixelFormat::Gray8,
}
}
}
impl RenderOptions {
fn pdfium_config(&self) -> RenderConfig {
RenderConfig::new()
.dpi(self.dpi)
.pixel_format(self.pixel_format.pdfium_format())
.annotations(self.annotations)
.form_fields(self.form_fields)
.max_output_bytes(self.max_output_bytes_per_page)
}
}
/// Errors produced by the optional local renderer.
#[derive(Debug, Error)]
#[non_exhaustive]
pub enum RenderError {
/// Page numbers in pdf-inspector APIs are 1-indexed, so zero is invalid.
#[error("page numbers are 1-indexed; page 0 is invalid")]
InvalidPageNumber,
/// The requested 1-indexed page is not present in the document.
#[error("page {page} is out of bounds for a {page_count}-page document")]
PageOutOfBounds {
/// Requested 1-indexed page.
page: u32,
/// Number of pages in the document.
page_count: usize,
},
/// PDFium loading, document parsing, form setup, or rendering failed.
#[error(transparent)]
Pdfium(#[from] firecrawl_pdfium::Error),
/// PDFium returned an internally inconsistent bitmap or transform.
#[error(transparent)]
Buffer(#[from] RenderBufferError),
}
/// Loaded PDFium renderer used to prepare pages for OCR.
///
/// PDFium calls are safe from concurrent threads but serialize inside the
/// underlying binding. Returned [`RenderedPage`] values are ordinary owned
/// data and can be processed concurrently after rendering.
#[derive(Debug, Clone, Copy)]
pub struct PdfiumRenderer {
pdfium: Pdfium,
}
impl PdfiumRenderer {
/// Loads PDFium using `firecrawl-pdfium`'s documented discovery chain.
pub fn load() -> Result<Self, RenderError> {
Ok(Self {
pdfium: Pdfium::load()?,
})
}
/// Loads PDFium from an explicit native library path.
pub fn load_from_path(path: impl AsRef<Path>) -> Result<Self, RenderError> {
Ok(Self {
pdfium: Pdfium::load_from_path(path)?,
})
}
/// Path of the active PDFium library, if it was loaded from a concrete
/// file rather than through the system loader.
pub fn loaded_from(&self) -> Option<&Path> {
self.pdfium.loaded_from()
}
/// Renders selected 1-indexed pages in the same order as `pages`.
///
/// This inherent method mirrors [`PageRenderer`] so existing callers do
/// not need to import the trait.
pub fn render_pages(
&self,
pdf_bytes: &[u8],
pages: &[u32],
password: Option<&str>,
options: &RenderOptions,
) -> Result<Vec<RenderedPage>, RenderError> {
self.render_pages_impl(pdf_bytes, pages, password, options)
}
fn render_pages_impl(
&self,
pdf_bytes: &[u8],
pages: &[u32],
password: Option<&str>,
options: &RenderOptions,
) -> Result<Vec<RenderedPage>, RenderError> {
if pages.is_empty() {
return Ok(Vec::new());
}
if pages.contains(&0) {
return Err(RenderError::InvalidPageNumber);
}
let document = self.pdfium.load_document(pdf_bytes.to_vec(), password)?;
let page_count = document.page_count();
if let Some(&page) = pages.iter().find(|&&page| page as usize > page_count) {
return Err(RenderError::PageOutOfBounds { page, page_count });
}
if options.form_fields {
document.enable_form_rendering()?;
}
let config = options.pdfium_config();
let mut rendered_pages = Vec::with_capacity(pages.len());
for &page_number in pages {
let page = document.page(page_number as usize - 1)?;
let size = page.size();
let rendered = page.render(&config)?;
rendered_pages.push(rendered_page_from_pdfium(
page_number,
size.width,
size.height,
options.pixel_format,
rendered,
)?);
}
Ok(rendered_pages)
}
}
impl PageRenderer for PdfiumRenderer {
type Error = RenderError;
fn render_pages(
&self,
pdf_bytes: &[u8],
pages: &[u32],
password: Option<&str>,
options: &RenderOptions,
) -> Result<Vec<RenderedPage>, Self::Error> {
self.render_pages_impl(pdf_bytes, pages, password, options)
}
}
fn rendered_page_from_pdfium(
page: u32,
page_width: f32,
page_height: f32,
format: RenderPixelFormat,
rendered: firecrawl_pdfium::RenderedPage,
) -> Result<RenderedPage, RenderBufferError> {
let width = rendered.width();
let height = rendered.height();
let stride = rendered.stride();
let pdfium_transform = *rendered.transform();
let corner = |x, y| {
let point = pdfium_transform.pixel_to_page(PixelPoint::new(x, y));
(point.x, point.y)
};
let transform = PageTransform::from_corners(
width,
height,
corner(0.0, 0.0),
corner(f64::from(width), 0.0),
corner(0.0, f64::from(height)),
)
.ok_or(RenderBufferError::InvalidTransform)?;
let mut pixels = rendered.into_pixels();
if format == RenderPixelFormat::Rgb8 {
bgr_to_rgb_in_place(&mut pixels, width, height, stride)?;
}
RenderedPage::new(
page,
page_width,
page_height,
width,
height,
stride,
format,
pixels,
transform,
)
}
fn bgr_to_rgb_in_place(
pixels: &mut [u8],
width: u32,
height: u32,
stride: usize,
) -> Result<(), RenderBufferError> {
let row_bytes = (width as usize)
.checked_mul(RenderPixelFormat::Rgb8.bytes_per_pixel())
.ok_or(RenderBufferError::SizeOverflow)?;
if stride < row_bytes {
return Err(RenderBufferError::InvalidStride {
stride,
minimum: row_bytes,
});
}
let expected = stride
.checked_mul(height as usize)
.ok_or(RenderBufferError::SizeOverflow)?;
if pixels.len() != expected {
return Err(RenderBufferError::InvalidBufferLength {
actual: pixels.len(),
expected,
});
}
for row in pixels.chunks_exact_mut(stride) {
for pixel in row[..row_bytes].chunks_exact_mut(3) {
pixel.swap(0, 2);
}
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn bgr_pixels_are_converted_to_rgb_in_place() {
let mut pixels = vec![1, 2, 3, 4, 5, 6];
bgr_to_rgb_in_place(&mut pixels, 2, 1, 6).unwrap();
assert_eq!(pixels, [3, 2, 1, 6, 5, 4]);
}
#[test]
fn bgr_conversion_skips_row_padding() {
let mut pixels = vec![1, 2, 3, 9, 7, 8, 9, 6];
bgr_to_rgb_in_place(&mut pixels, 1, 2, 4).unwrap();
assert_eq!(pixels, [3, 2, 1, 9, 9, 8, 7, 6]);
}
#[test]
fn malformed_bgr_buffers_return_errors() {
assert!(matches!(
bgr_to_rgb_in_place(&mut [0; 6], 2, 1, 5),
Err(RenderBufferError::InvalidStride { .. })
));
assert!(matches!(
bgr_to_rgb_in_place(&mut [0; 5], 1, 2, 3),
Err(RenderBufferError::InvalidBufferLength { .. })
));
}
}
+552
View File
@@ -0,0 +1,552 @@
//! Renderer-neutral page bitmap and coordinate types.
use thiserror::Error;
use crate::PdfRect;
/// Default rendering resolution for OCR.
pub const DEFAULT_RENDER_DPI: f32 = 150.0;
/// Default maximum size of one rendered page: 256 MiB.
pub const DEFAULT_MAX_OUTPUT_BYTES: u64 = 256 * 1024 * 1024;
/// Pixel layout returned by [`RenderedPage`].
#[derive(Debug, Clone, Copy, Default, PartialEq, Eq)]
#[non_exhaustive]
pub enum RenderPixelFormat {
/// Three bytes per pixel in red, green, blue order. This is the default
/// because OCR preprocessors conventionally consume RGB images.
#[default]
Rgb8,
/// Four bytes per pixel in red, green, blue, alpha order.
Rgba8,
/// One luminance byte per pixel.
Gray8,
}
impl RenderPixelFormat {
/// Number of bytes used by one pixel.
pub fn bytes_per_pixel(self) -> usize {
match self {
Self::Rgb8 => 3,
Self::Rgba8 => 4,
Self::Gray8 => 1,
}
}
}
/// Configuration for pages rendered as input to a local vision pipeline.
#[derive(Debug, Clone, PartialEq)]
pub struct RenderOptions {
/// Output resolution. Defaults to 150 DPI.
pub dpi: f32,
/// Pixel layout. Defaults to three-channel RGB.
pub pixel_format: RenderPixelFormat,
/// Include PDF annotations in the rendered bitmap.
pub annotations: bool,
/// Include visible static AcroForm field appearances.
pub form_fields: bool,
/// Maximum allocation for each rendered page.
pub max_output_bytes_per_page: u64,
}
impl Default for RenderOptions {
fn default() -> Self {
Self {
dpi: DEFAULT_RENDER_DPI,
pixel_format: RenderPixelFormat::Rgb8,
annotations: true,
form_fields: true,
max_output_bytes_per_page: DEFAULT_MAX_OUTPUT_BYTES,
}
}
}
impl RenderOptions {
/// Creates local-rendering options with OCR-oriented defaults.
pub fn new() -> Self {
Self::default()
}
/// Sets the output resolution in dots per inch.
pub fn dpi(mut self, dpi: f32) -> Self {
self.dpi = dpi;
self
}
/// Sets the output pixel layout.
pub fn pixel_format(mut self, pixel_format: RenderPixelFormat) -> Self {
self.pixel_format = pixel_format;
self
}
/// Toggles annotation rendering.
pub fn annotations(mut self, annotations: bool) -> Self {
self.annotations = annotations;
self
}
/// Toggles visible static form-field rendering.
pub fn form_fields(mut self, form_fields: bool) -> Self {
self.form_fields = form_fields;
self
}
/// Sets the maximum allocation for each rendered page.
pub fn max_output_bytes_per_page(mut self, bytes: u64) -> Self {
self.max_output_bytes_per_page = bytes;
self
}
}
/// A point in PDF page space, measured in points from the bottom-left.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct PagePoint {
/// Horizontal position in PDF points.
pub x: f32,
/// Vertical position in PDF points, increasing upward.
pub y: f32,
}
/// Affine transform between top-left pixel space and PDF page space.
///
/// Renderers create this from the page-space images of the bitmap corners.
/// Keeping the coefficients in pdf-inspector makes [`RenderedPage`] neutral
/// to the renderer implementation that produced it.
#[derive(Debug, Clone, Copy, PartialEq)]
pub struct PageTransform {
forward: [f64; 6],
inverse: [f64; 6],
pixel_width: u32,
pixel_height: u32,
}
impl PageTransform {
/// Builds a transform from the PDF-space images of device corners
/// `(0, 0)`, `(pixel_width, 0)`, and `(0, pixel_height)`.
pub fn from_corners(
pixel_width: u32,
pixel_height: u32,
origin: (f64, f64),
x_axis: (f64, f64),
y_axis: (f64, f64),
) -> Option<Self> {
if pixel_width == 0 || pixel_height == 0 {
return None;
}
let values = [origin.0, origin.1, x_axis.0, x_axis.1, y_axis.0, y_axis.1];
if values.iter().any(|value| !value.is_finite()) {
return None;
}
let width = f64::from(pixel_width);
let height = f64::from(pixel_height);
let a = (x_axis.0 - origin.0) / width;
let c = (x_axis.1 - origin.1) / width;
let b = (y_axis.0 - origin.0) / height;
let d = (y_axis.1 - origin.1) / height;
let (e, f) = origin;
let forward = [a, b, c, d, e, f];
if forward.iter().any(|coefficient| !coefficient.is_finite()) {
return None;
}
let determinant = a * d - b * c;
if determinant == 0.0 || !determinant.is_finite() {
return None;
}
let inverse_a = d / determinant;
let inverse_b = -b / determinant;
let inverse_c = -c / determinant;
let inverse_d = a / determinant;
let inverse_e = -(inverse_a * e + inverse_b * f);
let inverse_f = -(inverse_c * e + inverse_d * f);
let inverse = [
inverse_a, inverse_b, inverse_c, inverse_d, inverse_e, inverse_f,
];
if inverse.iter().any(|coefficient| !coefficient.is_finite()) {
return None;
}
Some(Self {
forward,
inverse,
pixel_width,
pixel_height,
})
}
/// Width of the bitmap this transform describes.
pub fn pixel_width(&self) -> u32 {
self.pixel_width
}
/// Height of the bitmap this transform describes.
pub fn pixel_height(&self) -> u32 {
self.pixel_height
}
/// Converts a bitmap point to PDF page space.
pub fn pixel_to_page(&self, x: f64, y: f64) -> PagePoint {
let [a, b, c, d, e, f] = self.forward;
PagePoint {
x: (a * x + b * y + e) as f32,
y: (c * x + d * y + f) as f32,
}
}
/// Converts a PDF page-space point to bitmap coordinates.
pub fn page_to_pixel(&self, x: f64, y: f64) -> (f64, f64) {
let [a, b, c, d, e, f] = self.inverse;
(a * x + b * y + e, c * x + d * y + f)
}
}
/// Invalid renderer output rejected by [`RenderedPage::new`].
#[derive(Debug, Error, Clone, PartialEq, Eq)]
#[non_exhaustive]
pub enum RenderBufferError {
/// Page numbers are 1-indexed.
#[error("rendered page number must be at least 1")]
InvalidPageNumber,
/// Bitmap dimensions must be non-zero.
#[error("rendered bitmap dimensions must be non-zero")]
InvalidDimensions,
/// Page dimensions must be positive finite numbers.
#[error("rendered PDF page dimensions must be positive and finite")]
InvalidPageDimensions,
/// Transform dimensions must match the bitmap dimensions.
#[error("coordinate transform dimensions do not match the rendered bitmap")]
TransformDimensions,
/// Renderer did not provide an invertible finite coordinate transform.
#[error("renderer returned an invalid coordinate transform")]
InvalidTransform,
/// The stride cannot hold one active row of pixels.
#[error("pixel stride {stride} is shorter than the active row size {minimum}")]
InvalidStride {
/// Supplied bytes per row.
stride: usize,
/// Minimum bytes required for one row.
minimum: usize,
},
/// Pixel buffer size is inconsistent with height and stride.
#[error("pixel buffer has {actual} bytes; expected {expected}")]
InvalidBufferLength {
/// Actual byte count.
actual: usize,
/// Required byte count.
expected: usize,
},
/// Dimension arithmetic overflowed the host address space.
#[error("rendered bitmap dimensions overflow the host address space")]
SizeOverflow,
}
/// One rendered page with owned pixels and its pixel-to-PDF transform.
///
/// The value contains no live renderer, page, or document handles. It can be
/// moved to an OCR worker and retained after rendering returns.
#[derive(Debug, Clone)]
pub struct RenderedPage {
page: u32,
page_width: f32,
page_height: f32,
width: u32,
height: u32,
stride: usize,
format: RenderPixelFormat,
pixels: Vec<u8>,
transform: PageTransform,
}
impl RenderedPage {
/// Creates a renderer-neutral owned page after validating its buffer.
#[allow(clippy::too_many_arguments)]
pub fn new(
page: u32,
page_width: f32,
page_height: f32,
width: u32,
height: u32,
stride: usize,
format: RenderPixelFormat,
pixels: Vec<u8>,
transform: PageTransform,
) -> Result<Self, RenderBufferError> {
if page == 0 {
return Err(RenderBufferError::InvalidPageNumber);
}
if width == 0 || height == 0 {
return Err(RenderBufferError::InvalidDimensions);
}
if page_width <= 0.0
|| page_height <= 0.0
|| !page_width.is_finite()
|| !page_height.is_finite()
{
return Err(RenderBufferError::InvalidPageDimensions);
}
if transform.pixel_width() != width || transform.pixel_height() != height {
return Err(RenderBufferError::TransformDimensions);
}
let row_bytes = (width as usize)
.checked_mul(format.bytes_per_pixel())
.ok_or(RenderBufferError::SizeOverflow)?;
if stride < row_bytes {
return Err(RenderBufferError::InvalidStride {
stride,
minimum: row_bytes,
});
}
let expected = stride
.checked_mul(height as usize)
.ok_or(RenderBufferError::SizeOverflow)?;
if pixels.len() != expected {
return Err(RenderBufferError::InvalidBufferLength {
actual: pixels.len(),
expected,
});
}
Ok(Self {
page,
page_width,
page_height,
width,
height,
stride,
format,
pixels,
transform,
})
}
/// 1-indexed page number.
pub fn page(&self) -> u32 {
self.page
}
/// Page width in PDF points after applying the page's rotation.
pub fn page_width(&self) -> f32 {
self.page_width
}
/// Page height in PDF points after applying the page's rotation.
pub fn page_height(&self) -> f32 {
self.page_height
}
/// Bitmap width in pixels.
pub fn width(&self) -> u32 {
self.width
}
/// Bitmap height in pixels.
pub fn height(&self) -> u32 {
self.height
}
/// Number of bytes between adjacent bitmap rows.
pub fn stride(&self) -> usize {
self.stride
}
/// Pixel layout of [`pixels`](Self::pixels).
pub fn format(&self) -> RenderPixelFormat {
self.format
}
/// Owned bitmap bytes, with rows ordered top-to-bottom.
pub fn pixels(&self) -> &[u8] {
&self.pixels
}
/// Consumes the page and returns its pixel buffer.
pub fn into_pixels(self) -> Vec<u8> {
self.pixels
}
/// Coordinate transform associated with the rendered page.
pub fn transform(&self) -> PageTransform {
self.transform
}
/// Converts a bitmap point (top-left origin, y-down) to PDF page space
/// (bottom-left origin, y-up).
pub fn pixel_to_page(&self, x: f64, y: f64) -> PagePoint {
self.transform.pixel_to_page(x, y)
}
/// Converts a bitmap rectangle to the repository's existing PDF-space
/// rectangle type. The returned page number remains 1-indexed.
pub fn pixel_rect_to_pdf_rect(&self, x: f64, y: f64, width: f64, height: f64) -> PdfRect {
let points = [
self.transform.pixel_to_page(x, y),
self.transform.pixel_to_page(x + width, y),
self.transform.pixel_to_page(x, y + height),
self.transform.pixel_to_page(x + width, y + height),
];
let left = points
.iter()
.map(|point| point.x)
.fold(f32::INFINITY, f32::min);
let right = points
.iter()
.map(|point| point.x)
.fold(f32::NEG_INFINITY, f32::max);
let bottom = points
.iter()
.map(|point| point.y)
.fold(f32::INFINITY, f32::min);
let top = points
.iter()
.map(|point| point.y)
.fold(f32::NEG_INFINITY, f32::max);
PdfRect {
x: left,
y: bottom,
width: right - left,
height: top - bottom,
page: self.page,
}
}
/// Converts a PDF-space rectangle to bitmap coordinates
/// `(x, y, width, height)` with a top-left origin.
pub fn pdf_rect_to_pixel(&self, rect: &PdfRect) -> (f64, f64, f64, f64) {
let left = f64::from(rect.x);
let right = f64::from(rect.x + rect.width);
let bottom = f64::from(rect.y);
let top = f64::from(rect.y + rect.height);
let points = [
self.transform.page_to_pixel(left, bottom),
self.transform.page_to_pixel(right, bottom),
self.transform.page_to_pixel(left, top),
self.transform.page_to_pixel(right, top),
];
let min_x = points
.iter()
.map(|point| point.0)
.fold(f64::INFINITY, f64::min);
let max_x = points
.iter()
.map(|point| point.0)
.fold(f64::NEG_INFINITY, f64::max);
let min_y = points
.iter()
.map(|point| point.1)
.fold(f64::INFINITY, f64::min);
let max_y = points
.iter()
.map(|point| point.1)
.fold(f64::NEG_INFINITY, f64::max);
(min_x, min_y, max_x - min_x, max_y - min_y)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn transform() -> PageTransform {
PageTransform::from_corners(400, 200, (0.0, 100.0), (200.0, 100.0), (0.0, 0.0)).unwrap()
}
#[test]
fn transform_maps_both_directions_at_non_identity_scale() {
let transform = transform();
let point = transform.pixel_to_page(100.0, 50.0);
assert!((point.x - 50.0).abs() < 1e-6);
assert!((point.y - 75.0).abs() < 1e-6);
let pixel = transform.page_to_pixel(f64::from(point.x), f64::from(point.y));
assert!((pixel.0 - 100.0).abs() < 1e-6);
assert!((pixel.1 - 50.0).abs() < 1e-6);
}
#[test]
fn rendered_page_accepts_padding_and_validates_length() {
let page = RenderedPage::new(
1,
200.0,
100.0,
400,
200,
1_204,
RenderPixelFormat::Rgb8,
vec![0; 1_204 * 200],
transform(),
)
.unwrap();
assert_eq!(page.stride(), 1_204);
assert!(matches!(
RenderedPage::new(
1,
200.0,
100.0,
400,
200,
1_204,
RenderPixelFormat::Rgb8,
vec![0; 5],
transform(),
),
Err(RenderBufferError::InvalidBufferLength { .. })
));
}
#[test]
fn rotated_transform_round_trips_rectangles() {
let transform =
PageTransform::from_corners(100, 200, (0.0, 0.0), (0.0, 100.0), (200.0, 0.0)).unwrap();
let page = RenderedPage::new(
1,
200.0,
100.0,
100,
200,
300,
RenderPixelFormat::Rgb8,
vec![0; 300 * 200],
transform,
)
.unwrap();
let pdf = page.pixel_rect_to_pdf_rect(10.0, 20.0, 30.0, 40.0);
let pixel = page.pdf_rect_to_pixel(&pdf);
assert!((pixel.0 - 10.0).abs() < 1e-5);
assert!((pixel.1 - 20.0).abs() < 1e-5);
assert!((pixel.2 - 30.0).abs() < 1e-5);
assert!((pixel.3 - 40.0).abs() < 1e-5);
}
#[test]
fn skewed_transform_bounds_all_rectangle_corners() {
let transform =
PageTransform::from_corners(100, 100, (0.0, 100.0), (100.0, 125.0), (25.0, 0.0))
.unwrap();
let page = RenderedPage::new(
1,
125.0,
125.0,
100,
100,
300,
RenderPixelFormat::Rgb8,
vec![0; 30_000],
transform,
)
.unwrap();
let pdf = page.pixel_rect_to_pdf_rect(10.0, 20.0, 30.0, 40.0);
assert!((pdf.x - 15.0).abs() < 1e-5);
assert!((pdf.y - 42.5).abs() < 1e-5);
assert!((pdf.width - 40.0).abs() < 1e-5);
assert!((pdf.height - 47.5).abs() < 1e-5);
let pixels = page.pdf_rect_to_pixel(&pdf);
assert!(pixels.0 <= 10.0 && pixels.1 <= 20.0);
assert!(pixels.0 + pixels.2 >= 40.0 && pixels.1 + pixels.3 >= 60.0);
}
}
+65
View File
@@ -0,0 +1,65 @@
#![cfg(all(feature = "render-pdfium", not(target_arch = "wasm32")))]
use pdf_inspector::vision::{PdfiumRenderer, RenderError, RenderOptions, RenderPixelFormat};
fn load_renderer() -> Option<PdfiumRenderer> {
match PdfiumRenderer::load() {
Ok(renderer) => Some(renderer),
Err(RenderError::Pdfium(firecrawl_pdfium::Error::Load(
firecrawl_pdfium::LoadError::LibraryNotFound { .. },
))) => {
eprintln!("skipping PDFium runtime test because no native library is installed");
None
}
Err(error) => panic!("failed to load PDFium: {error}"),
}
}
#[test]
fn renders_owned_rgb_page_and_round_trips_coordinates() {
let Some(renderer) = load_renderer() else {
return;
};
let bytes = std::fs::read("tests/fixtures/thermo-freon12.pdf").unwrap();
let pages = renderer
.render_pages(
&bytes,
&[1],
None,
&RenderOptions::new().dpi(150.0).form_fields(false),
)
.unwrap();
assert_eq!(pages.len(), 1);
let page = &pages[0];
assert_eq!(page.page(), 1);
assert_eq!(page.format(), RenderPixelFormat::Rgb8);
assert_eq!(page.stride(), page.width() as usize * 3);
assert_eq!(page.pixels().len(), page.stride() * page.height() as usize);
assert!((page.width() as f32 - page.page_width()).abs() > 1.0);
let pdf_rect = page.pixel_rect_to_pdf_rect(10.0, 10.0, 20.0, 12.0);
let pixel_rect = page.pdf_rect_to_pixel(&pdf_rect);
assert!((pixel_rect.0 - 10.0).abs() < 0.01);
assert!((pixel_rect.1 - 10.0).abs() < 0.01);
assert!((pixel_rect.2 - 20.0).abs() < 0.01);
assert!((pixel_rect.3 - 12.0).abs() < 0.01);
}
#[test]
fn rejects_zero_and_out_of_range_page_numbers() {
let Some(renderer) = load_renderer() else {
return;
};
let bytes = std::fs::read("tests/fixtures/thermo-freon12.pdf").unwrap();
assert!(matches!(
renderer.render_pages(&bytes, &[0], None, &RenderOptions::new()),
Err(RenderError::InvalidPageNumber)
));
assert!(matches!(
renderer.render_pages(&bytes, &[u32::MAX], None, &RenderOptions::new()),
Err(RenderError::PageOutOfBounds { .. })
));
}