Compare commits
6
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
e3f5429638 | ||
|
|
f6cbe979f6 | ||
|
|
3f43745313 | ||
|
|
a012cb65a6 | ||
|
|
6567e1ab2d | ||
|
|
3af409d27f |
+14
@@ -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
@@ -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):
|
||||
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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 0–1.
|
||||
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 0–1.
|
||||
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"))
|
||||
);
|
||||
}
|
||||
}
|
||||
@@ -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};
|
||||
@@ -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");
|
||||
}
|
||||
}
|
||||
@@ -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 { .. })
|
||||
));
|
||||
}
|
||||
}
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
@@ -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 { .. })
|
||||
));
|
||||
}
|
||||
Reference in New Issue
Block a user