diff --git a/Cargo.toml b/Cargo.toml index e42da2a..8b11d24 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -52,6 +52,14 @@ 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. @@ -65,7 +73,9 @@ tempfile = "3.3" [features] default = [] python = ["pyo3"] -render-pdfium = ["dep:firecrawl-pdfium"] +vision = [] +model-cache = ["vision", "dep:dirs", "dep:fs2", "dep:sha2", "dep:windows-sys"] +render-pdfium = ["vision", "dep:firecrawl-pdfium"] [[bin]] name = "pdf2md" diff --git a/docs/rust-api.md b/docs/rust-api.md index 8d4d4dc..4f614ab 100644 --- a/docs/rust-api.md +++ b/docs/rust-api.md @@ -117,12 +117,53 @@ 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. +model or change the existing extraction functions. It implies `vision`, +and `PdfiumRenderer` implements the renderer-neutral `PageRenderer` trait. ```toml [dependencies] diff --git a/src/vision/contracts.rs b/src/vision/contracts.rs new file mode 100644 index 0000000..bdff733 --- /dev/null +++ b/src/vision/contracts.rs @@ -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, + /// Optional directory containing an offline model set. + pub model_directory: Option, + /// 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>) -> 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) -> 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, +} + +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) -> 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, revision: impl Into) -> 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, +} + +/// 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, + /// Mean confidence across accepted spans, when available. + pub mean_confidence: Option, + /// 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, +} + +/// 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, +} + +/// 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, + /// 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, +} + +/// 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, + /// Learned layout model, when layout inference ran. + pub layout_model: Option, + /// Render resolution used for local vision. + pub render_dpi: Option, + /// Mean accepted OCR confidence, when available. + pub ocr_confidence: Option, + /// Stage timings. + pub timings: VisionTimings, + /// Non-fatal warnings surfaced to downstream users. + pub warnings: Vec, + /// 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, 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, 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, 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")) + ); + } +} diff --git a/src/vision/mod.rs b/src/vision/mod.rs index 3fa1eea..2d489aa 100644 --- a/src/vision/mod.rs +++ b/src/vision/mod.rs @@ -1,13 +1,37 @@ -//! Optional native vision primitives used by local extraction pipelines. +//! 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 and is kept -//! separate so browser WASM and text-only consumers do not take on PDFium. +//! 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 = "render-pdfium", not(target_arch = "wasm32")))] -pub use pdfium::{ - PagePoint, PdfiumRenderer, RenderError, RenderOptions, RenderPixelFormat, RenderedPage, +#[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}; diff --git a/src/vision/models.rs b/src/vision/models.rs new file mode 100644 index 0000000..4efe3b8 --- /dev/null +++ b/src/vision/models.rs @@ -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, +} + +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 { + 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, +} + +impl ModelStore { + /// Creates a model store rooted at an explicit cache directory. + pub fn new(cache_root: impl Into) -> 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 { + 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) -> 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 { + 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 { + 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 = source.as_os_str().encode_wide().chain(Some(0)).collect(); + let target: Vec = 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 { + 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"); + } +} diff --git a/src/vision/pdfium.rs b/src/vision/pdfium.rs index d9ceeb5..ca4dd75 100644 --- a/src/vision/pdfium.rs +++ b/src/vision/pdfium.rs @@ -1,47 +1,18 @@ -//! PDFium-backed page rendering for OCR. +//! PDFium-backed implementation of the renderer-neutral page contract. use std::path::Path; -use firecrawl_pdfium::{ - PageRect as PdfiumPageRect, Pdfium, PixelFormat, PixelPoint, PixelRect, RenderConfig, -}; +use firecrawl_pdfium::{Pdfium, PixelFormat, PixelPoint, RenderConfig}; 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, -} +use super::{ + PageRenderer, PageTransform, RenderBufferError, RenderOptions, RenderPixelFormat, RenderedPage, +}; 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, - } - } - fn pdfium_format(self) -> PixelFormat { match self { - // PDFium produces BGR directly; `RenderedPage::from_pdfium` + // PDFium produces BGR directly; `rendered_page_from_pdfium` // swaps the red and blue channels in place. Self::Rgb8 => PixelFormat::Bgr8, Self::Rgba8 => PixelFormat::Rgba8, @@ -50,69 +21,7 @@ impl RenderPixelFormat { } } -/// 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 - } - fn pdfium_config(&self) -> RenderConfig { RenderConfig::new() .dpi(self.dpi) @@ -123,147 +32,6 @@ impl RenderOptions { } } -/// 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, -} - -/// One rendered page with owned pixels and its pixel-to-PDF transform. -/// -/// The value contains no live PDFium page or document handles. It can be -/// moved to an OCR worker and retained after [`PdfiumRenderer::render_pages`] -/// 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, - transform: firecrawl_pdfium::PageTransform, -} - -impl RenderedPage { - fn from_pdfium( - page: u32, - page_width: f32, - page_height: f32, - format: RenderPixelFormat, - rendered: firecrawl_pdfium::RenderedPage, - ) -> Self { - let width = rendered.width(); - let height = rendered.height(); - let stride = rendered.stride(); - let transform = *rendered.transform(); - let mut pixels = rendered.into_pixels(); - - if format == RenderPixelFormat::Rgb8 { - bgr_to_rgb_in_place(&mut pixels, width, height, stride); - } - - 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 `/Rotate` entry. - pub fn page_width(&self) -> f32 { - self.page_width - } - - /// Page height in PDF points after applying the page's `/Rotate` entry. - 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 { - self.pixels - } - - /// 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 { - let point = self.transform.pixel_to_page(PixelPoint::new(x, y)); - PagePoint { - x: point.x as f32, - y: point.y as f32, - } - } - - /// 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 rect = self - .transform - .pixel_rect_to_page(PixelRect::new(x, y, width, height)); - PdfRect { - x: rect.left as f32, - y: rect.bottom as f32, - width: rect.width() as f32, - height: rect.height() as f32, - 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 rect = self.transform.page_rect_to_pixel(PdfiumPageRect::new( - f64::from(rect.x), - f64::from(rect.y), - f64::from(rect.x + rect.width), - f64::from(rect.y + rect.height), - )); - (rect.x, rect.y, rect.width, rect.height) - } -} - /// Errors produced by the optional local renderer. #[derive(Debug, Error)] #[non_exhaustive] @@ -282,6 +50,9 @@ pub enum RenderError { /// 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. @@ -317,14 +88,24 @@ impl PdfiumRenderer { /// Renders selected 1-indexed pages in the same order as `pages`. /// - /// The PDF is parsed once for the full batch. Passing an empty page list - /// returns immediately without parsing or allocating. + /// 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, 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, RenderError> { if pages.is_empty() { return Ok(Vec::new()); @@ -351,60 +132,135 @@ impl PdfiumRenderer { let page = document.page(page_number as usize - 1)?; let size = page.size(); let rendered = page.render(&config)?; - rendered_pages.push(RenderedPage::from_pdfium( + rendered_pages.push(rendered_page_from_pdfium( page_number, size.width, size.height, options.pixel_format, rendered, - )); + )?); } Ok(rendered_pages) } } -fn bgr_to_rgb_in_place(pixels: &mut [u8], width: u32, height: u32, stride: usize) { - let row_bytes = width as usize * RenderPixelFormat::Rgb8.bytes_per_pixel(); - assert!(stride >= row_bytes, "pixel stride is shorter than one row"); - assert_eq!( - pixels.len(), - stride * height as usize, - "pixel buffer length does not match stride and height" - ); +impl PageRenderer for PdfiumRenderer { + type Error = RenderError; + + fn render_pages( + &self, + pdf_bytes: &[u8], + pages: &[u32], + password: Option<&str>, + options: &RenderOptions, + ) -> Result, 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 { + 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 defaults_are_ocr_oriented_and_bounded() { - let options = RenderOptions::default(); - assert_eq!(options.dpi, 150.0); - assert_eq!(options.pixel_format, RenderPixelFormat::Rgb8); - assert!(options.annotations); - assert!(options.form_fields); - assert_eq!(options.max_output_bytes_per_page, 256 * 1024 * 1024); - } - #[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); + 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); + 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 { .. }) + )); + } } diff --git a/src/vision/render.rs b/src/vision/render.rs new file mode 100644 index 0000000..7f0fa9a --- /dev/null +++ b/src/vision/render.rs @@ -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 { + 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, + 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, + transform: PageTransform, + ) -> Result { + 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 { + 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); + } +}