diff --git a/Cargo.lock b/Cargo.lock index 6710dddeea..9cfe8c9c00 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4540,10 +4540,22 @@ dependencies = [ "pdfcraft-fonts", "png 0.18.1", "rten", + "serde_json", + "sha2 0.11.1", "tempfile", "thiserror 2.0.21", ] +[[package]] +name = "pdfcraft-ocr-worker" +version = "0.18.1" +dependencies = [ + "js-sys", + "pdfcraft-ocr", + "wasm-bindgen", + "web-sys", +] + [[package]] name = "pdfcraft-optimize" version = "0.18.1" diff --git a/crates/ocr/Cargo.toml b/crates/ocr/Cargo.toml index 1b04dbb522..60e66d46c9 100644 --- a/crates/ocr/Cargo.toml +++ b/crates/ocr/Cargo.toml @@ -14,6 +14,8 @@ workspace = true pdfcraft-fonts = { workspace = true } ocrs = { version = "0.13.1", default-features = false, features = ["rten"] } rten = { version = "0.26.0", default-features = false, features = ["onnx_format", "rten_format"] } +serde_json.workspace = true +sha2 = "0.11" thiserror.workspace = true [target.'cfg(not(target_arch = "wasm32"))'.dependencies] diff --git a/crates/ocr/src/lib.rs b/crates/ocr/src/lib.rs index ff5b920dca..aac95f18c6 100644 --- a/crates/ocr/src/lib.rs +++ b/crates/ocr/src/lib.rs @@ -29,7 +29,9 @@ pub mod tesseract; pub use confidence::{ConfidenceBand, band_for, is_suspect, low_confidence}; pub use merge::{MergeStrategy, StrategyOutcome, mean_confidence, recognize_with_strategy, rover_merge}; -pub use ocrs::{ALPHABET, DETECTION_MODEL, LANGUAGES, Models, Ocr, OcrsRecognizer, RECOGNITION_MODEL}; +pub use ocrs::{ + ALPHABET, DETECTION_MODEL, DETECTION_MODEL_SHA256, LANGUAGES, Models, Ocr, OcrsRecognizer, RECOGNITION_MODEL, RECOGNITION_MODEL_SHA256, +}; pub use pdfcraft_fonts::TextPlan; pub use pdfcraft_fonts::helvetica_width; @@ -464,6 +466,58 @@ pub fn text_layer(words: &[PlacedWord], plan: &pdfcraft_fonts::TextPlan) -> Vec< out } +/// The recognised lines as JSON, `[[["word", left, top, right, bottom], …], …]`, with the +/// word's confidence as a sixth item when it has one: how the web build's OCR workers send what +/// they found back to the page. +pub fn lines_to_json(lines: &[Line]) -> String { + let word = |w: &Word| { + let mut items = vec![serde_json::json!(w.text), w.rect[0].into(), w.rect[1].into(), w.rect[2].into(), w.rect[3].into()]; + items.extend(w.confidence.map(serde_json::Value::from)); + serde_json::Value::Array(items) + }; + serde_json::Value::Array(lines.iter().map(|l| serde_json::Value::Array(l.words.iter().map(word).collect())).collect()).to_string() +} + +/// At most this many words come back for one page (a dense page has a few thousand). +const MAX_JSON_WORDS: usize = 100_000; + +/// [`lines_to_json`] read back, the words credited to `source` (the engine the worker ran). +/// Anything else (a word without its box, a box or confidence that isn't a finite number, too +/// many words) is an error rather than a guess. +pub fn lines_from_json(json: &str, source: &str) -> Result, OcrError> { + let bad = |why: &str| OcrError::Recognize(format!("unreadable result: {why}")); + let value: serde_json::Value = serde_json::from_str(json).map_err(|e| bad(&e.to_string()))?; + let lines = value.as_array().ok_or_else(|| bad("not a list of lines"))?; + let number = |v: &serde_json::Value| v.as_f64().map(|v| v as f32).filter(|v| v.is_finite()); + let mut count = 0usize; + let mut out = Vec::with_capacity(lines.len().min(MAX_JSON_WORDS)); + for line in lines { + let mut words = Vec::new(); + for word in line.as_array().ok_or_else(|| bad("a line is not a list of words"))? { + count += 1; + if count > MAX_JSON_WORDS { + return Err(bad("too many words")); + } + let (text, rect, confidence) = match word.as_array().map(Vec::as_slice) { + Some([serde_json::Value::String(text), l, t, r, b]) => (text, [l, t, r, b], None), + Some([serde_json::Value::String(text), l, t, r, b, c]) => (text, [l, t, r, b], Some(c)), + _ => return Err(bad("a word is not [text, left, top, right, bottom, confidence?]")), + }; + let mut r = [0f32; 4]; + for (slot, v) in r.iter_mut().zip(rect) { + *slot = number(v).ok_or_else(|| bad("a word's box is not finite numbers"))?; + } + let mut w = Word::new(text.clone(), r, source); + if let Some(c) = confidence { + w = w.with_confidence(number(c).ok_or_else(|| bad("a word's confidence is not a finite number"))?); + } + words.push(w); + } + out.push(Line { words }); + } + Ok(out) +} + #[cfg(test)] mod tests { use super::*; @@ -854,6 +908,89 @@ mod tests { assert_eq!(ok.gray(-1, -1), Some(255)); } + /// The web build loads the models from downloaded bytes; they read the same text. + #[test] + fn reads_rendered_text_with_models_from_bytes() { + let Some(models) = Models::find().filter(|m| m.detection.ends_with(DETECTION_MODEL)) else { + eprintln!("skipped: ONNX OCR models not installed"); + return; + }; + let ocr = Ocr::from_bytes(std::fs::read(&models.detection).unwrap(), std::fs::read(&models.recognition).unwrap()).unwrap(); + let rec = OcrsRecognizer::from_ocr(ocr); + assert!(rec.available(), "a recogniser loaded from bytes needs no model files"); + let (w, h, px) = test_image::hello(); + let lines = rec.recognize(&OcrImage::new(w, h, px).unwrap(), &RecognizeOptions::default()).unwrap(); + let text: Vec = lines.iter().map(Line::text).collect(); + assert!(text.iter().any(|t| t.contains("HELL") && t.contains("WOR")), "{text:?}"); + } + + /// The web build's workers send lines back as JSON; they arrive unchanged, credited to the + /// engine the worker ran. + #[test] + fn lines_survive_json() { + let lines = vec![ + Line { + words: vec![ + Word::new("Total \"€250\"", [1.5, 2.0, 30.25, 12.0], "ocrs"), + Word::new("due", [32.0, 2.0, 50.0, 12.0], "ocrs").with_confidence(87.5), + ], + }, + Line::default(), + Line { words: vec![Word::new("日本語", [0.0, 20.0, 9.0, 30.0], "ocrs")] }, + ]; + assert_eq!(lines_from_json(&lines_to_json(&lines), "ocrs").unwrap(), lines); + assert_eq!(lines_from_json(&lines_to_json(&[]), "ocrs").unwrap(), Vec::::new()); + let other = lines_from_json(&lines_to_json(&lines), "tesseract").unwrap(); + assert!(other.iter().flat_map(|l| &l.words).all(|w| w.source == "tesseract")); + } + + /// A result that isn't what a worker sends is an error, not a crash or a guess. + #[test] + fn malformed_json_lines_are_an_error() { + for json in [ + "", + "{}", + "[1]", + "[[1]]", + "[[[\"a\", 1, 2, 3]]]", + "[[[\"a\", 1, 2, 3, 4, 5, 6]]]", + "[[[1, 1, 2, 3, 4]]]", + "[[[\"a\", \"1\", 2, 3, 4]]]", + "[[[\"a\", 1e999, 2, 3, 4]]]", + "[[[\"a\", null, 2, 3, 4]]]", + "[[[\"a\", 1, 2, 3, 4, null]]]", + "[[[\"a\", 1, 2, 3, 4, \"90\"]]]", + ] { + assert!(matches!(lines_from_json(json, "ocrs"), Err(OcrError::Recognize(_))), "{json}"); + } + let many = format!("[[{}]]", vec!["[\"a\",0,0,1,1]"; MAX_JSON_WORDS + 1].join(",")); + assert!(lines_from_json(&many, "ocrs").is_err()); + let deep = format!("{}{}", "[".repeat(10_000), "]".repeat(10_000)); + assert!(lines_from_json(&deep, "ocrs").is_err()); + } + + /// A download that isn't a model (an HTML error page, a truncated file) is an error, not a + /// crash or a hang. + #[test] + fn models_from_bad_bytes_are_an_error() { + for (detection, recognition) in [(Vec::new(), Vec::new()), (b"".to_vec(), b"RTEN".to_vec()), (vec![0xff; 64], vec![0; 64])] { + assert!(matches!(Ocr::from_bytes(detection, recognition), Err(OcrError::Load(..)))); + } + } + + /// The hashes [`Ocr::from_bytes`] checks are the ones `cargo xtask models` verifies. + #[test] + fn model_hashes_match_attribution() { + let manifest = include_str!("../../../ATTRIBUTION.toml"); + for (file, sha256) in [(DETECTION_MODEL, DETECTION_MODEL_SHA256), (RECOGNITION_MODEL, RECOGNITION_MODEL_SHA256)] { + let entry = manifest + .split("[[fetched]]") + .find(|e| e.contains(&format!("file = \"{file}\""))) + .unwrap_or_else(|| panic!("{file} is not in ATTRIBUTION.toml")); + assert!(entry.lines().any(|l| l == format!("sha256 = \"{sha256}\"")), "{file}: SHA-256 differs from ATTRIBUTION.toml"); + } + } + mod test_image { /// "HELLO WORLD" drawn with blocky 5×7 letters, 4 pixels per dot, black on white. pub fn hello() -> (u32, u32, Vec) { diff --git a/crates/ocr/src/ocrs.rs b/crates/ocr/src/ocrs.rs index 7f8feb47c4..94f2abfa18 100644 --- a/crates/ocr/src/ocrs.rs +++ b/crates/ocr/src/ocrs.rs @@ -32,6 +32,15 @@ pub const LANGUAGES: &[(&str, &str)] = &[("en", "English")]; /// The model files, as named in ATTRIBUTION.toml. pub const DETECTION_MODEL: &str = "text-detection.onnx"; pub const RECOGNITION_MODEL: &str = "text-recognition.onnx"; +/// Their SHA-256s, as in ATTRIBUTION.toml. [`Ocr::from_bytes`] loads nothing else: rten's ONNX +/// parser can spin forever on bytes that aren't a model (64 × 0xff does). +pub const DETECTION_MODEL_SHA256: &str = "a917b23dbd9524b465df7e922641b3ff2981623df4ded5a0234004ef2fee7cfe"; +pub const RECOGNITION_MODEL_SHA256: &str = "86c145c2edb96c8caed5b1ebb8f44d706408922211451c309c625157dd6061c5"; + +fn sha256_hex(bytes: &[u8]) -> String { + use sha2::{Digest, Sha256}; + Sha256::digest(bytes).iter().map(|b| format!("{b:02x}")).collect() +} pub const ALPHABET: &str = " 0123456789!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~€ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"; @@ -84,9 +93,25 @@ pub struct Ocr { impl Ocr { pub fn load(models: &Models) -> Result { let load = |p: &Path| rten::Model::load_file(p).map_err(|e| OcrError::Load(p.display().to_string(), e.to_string())); + Self::with_models(load(&models.detection)?, load(&models.recognition)?) + } + + /// Load the models from the contents of their files, where there is no file system (the web + /// build downloads them). Bytes that aren't the files ATTRIBUTION.toml names are an error. + pub fn from_bytes(detection: Vec, recognition: Vec) -> Result { + let load = |name: &str, sha256: &str, data: Vec| { + if sha256_hex(&data) != sha256 { + return Err(OcrError::Load(name.into(), "not the expected model file (SHA-256 differs)".into())); + } + rten::Model::load(data).map_err(|e| OcrError::Load(name.into(), e.to_string())) + }; + Self::with_models(load(DETECTION_MODEL, DETECTION_MODEL_SHA256, detection)?, load(RECOGNITION_MODEL, RECOGNITION_MODEL_SHA256, recognition)?) + } + + fn with_models(detection: rten::Model, recognition: rten::Model) -> Result { let params = ocrs::OcrEngineParams { - detection_model: Some(load(&models.detection)?), - recognition_model: Some(load(&models.recognition)?), + detection_model: Some(detection), + recognition_model: Some(recognition), alphabet: Some(ALPHABET.into()), ..Default::default() }; @@ -139,12 +164,21 @@ impl Ocr { /// The ocrs engine as a [`Recognizer`] (the default engine). pub struct OcrsRecognizer { engine: Ocr, + /// Loaded from bytes ([`OcrsRecognizer::from_ocr`]), not from model files: it stays + /// available without them. + in_memory: bool, } impl OcrsRecognizer { /// Load the models from `models`. pub fn load(models: &Models) -> Result { - Ok(OcrsRecognizer { engine: Ocr::load(models)? }) + Ok(OcrsRecognizer { engine: Ocr::load(models)?, in_memory: false }) + } + + /// A recogniser loaded some other way than from model files: the web build downloads the + /// models and loads them with [`Ocr::from_bytes`]. + pub fn from_ocr(engine: Ocr) -> OcrsRecognizer { + OcrsRecognizer { engine, in_memory: true } } /// Load the models found by [`Models::find`]. @@ -159,7 +193,7 @@ impl Recognizer for OcrsRecognizer { } fn available(&self) -> bool { - Models::find().is_some() + self.in_memory || Models::find().is_some() } fn languages(&self) -> Vec<(String, String)> {