diff --git a/.gitignore b/.gitignore index 6b903898..70fbb2cc 100644 --- a/.gitignore +++ b/.gitignore @@ -8,6 +8,8 @@ onnxruntime_profile_*.json # wasm-bindgen output for the docling-wasm demo (built from source, see its README) crates/docling-wasm/www/pkg/ +# ML models the scan demo fetches same-origin (download_dependencies.sh, run in www/) +crates/docling-wasm/www/models/ # Secrets — never commit (holds CARGO_REGISTRY_TOKEN) .env diff --git a/Cargo.lock b/Cargo.lock index 654b7a7b..e4645db1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1529,8 +1529,13 @@ version = "0.49.0" dependencies = [ "console_error_panic_hook", "docling", + "docling-core", + "docling-pdf", + "image", + "js-sys", "serde_json", "wasm-bindgen", + "wasm-bindgen-futures", ] [[package]] diff --git a/crates/docling-pdf/Cargo.toml b/crates/docling-pdf/Cargo.toml index 4a837aa2..8cde9c3f 100644 --- a/crates/docling-pdf/Cargo.toml +++ b/crates/docling-pdf/Cargo.toml @@ -18,7 +18,11 @@ default = ["ml"] # (`convert_text_layer` — lopdf content-stream parsing + the docling-parse # line sanitizer, the same extraction `--no-ocr` uses), which compiles for # wasm32-unknown-unknown. +# The ONNX-free OCR pre/post half (`ocr_prep`) alone — what the wasm build +# (#79 phase 2) pulls: pure Rust + the `image` crate, no ort/pdfium. +ocr-prep = ["dep:image"] ml = [ + "ocr-prep", "dep:pdfium-render", "dep:ort", "dep:image", diff --git a/crates/docling-pdf/src/layout.rs b/crates/docling-pdf/src/layout.rs index 782a9684..895ef7f5 100644 --- a/crates/docling-pdf/src/layout.rs +++ b/crates/docling-pdf/src/layout.rs @@ -47,13 +47,12 @@ pub struct Region { pub b: f32, } -#[cfg(feature = "ml")] /// Base confidence threshold (docling-ibm-models `base_threshold`): the raw /// RT-DETR floor before docling's `LayoutPostprocessor` applies its stricter /// per-label thresholds ([`label_threshold`]). const THRESHOLD: f32 = 0.3; -#[cfg(feature = "ml")] -const SIDE: u32 = 640; +/// RT-DETR's fixed square input side. +pub const SIDE: u32 = 640; /// Per-label confidence threshold, ported from docling's /// `LayoutPostprocessor.CONFIDENCE_THRESHOLDS`. The raw predictor keeps every @@ -222,46 +221,81 @@ impl LayoutModel { let logits = &logits[p * num_queries * num_classes..(p + 1) * num_queries * num_classes]; let boxes = &boxes[p * num_queries * 4..(p + 1) * num_queries * 4]; - - // sigmoid over every (query, class); take the top `num_queries` scores. - let mut scored: Vec<(f32, usize)> = (0..num_queries * num_classes) - .map(|idx| (sigmoid(logits[idx]), idx)) - .collect(); - scored.sort_unstable_by(|a, b| b.0.total_cmp(&a.0)); - scored.truncate(num_queries); - - let mut regions = Vec::new(); - for (score, idx) in scored { - if score <= THRESHOLD { - continue; - } - let label_id = idx % num_classes; - let q = idx / num_classes; - let cx = boxes[q * 4]; - let cy = boxes[q * 4 + 1]; - let w = boxes[q * 4 + 2]; - let h = boxes[q * 4 + 3]; - // center_to_corners, then scale normalized coords to page points. - let l = (cx - w / 2.0) * page_w; - let t = (cy - h / 2.0) * page_h; - let r = (cx + w / 2.0) * page_w; - let b = (cy + h / 2.0) * page_h; - regions.push(Region { - label: LABELS.get(label_id).copied().unwrap_or("text"), - score, - l, - t, - r, - b, - }); - } - all.push(regions); + all.push(decode_layout( + logits, + boxes, + num_queries, + num_classes, + *page_w, + *page_h, + )); } Ok(all) } } -#[cfg(feature = "ml")] fn sigmoid(x: f32) -> f32 { 1.0 / (1.0 + (-x).exp()) } + +/// Pack one page image into the model's `(1, 3, SIDE, SIDE)` input: resize +/// (aspect ignored, RT-DETR convention), rescale to `[0,1]`, CHW. Shared +/// with the browser build (#157), which delegates only the session call. +#[cfg(feature = "ocr-prep")] +pub fn layout_input(img: &image::RgbImage) -> Vec { + let n = (SIDE * SIDE) as usize; + let mut data = vec![0f32; 3 * n]; + let resized = image::imageops::resize(img, SIDE, SIDE, image::imageops::FilterType::Triangle); + for (i, px) in resized.pixels().enumerate() { + data[i] = px[0] as f32 / 255.0; + data[n + i] = px[1] as f32 / 255.0; + data[2 * n + i] = px[2] as f32 / 255.0; + } + data +} + +/// Decode one page's raw RT-DETR outputs into scored [`Region`]s in page +/// points — sigmoid over every (query, class), top-`num_queries` kept, boxes +/// converted center→corners and scaled. Shared with the browser build; the +/// native batch path calls it per page, so both decode identically. +pub fn decode_layout( + logits: &[f32], + boxes: &[f32], + num_queries: usize, + num_classes: usize, + page_w: f32, + page_h: f32, +) -> Vec { + let mut scored: Vec<(f32, usize)> = (0..num_queries * num_classes) + .map(|idx| (sigmoid(logits[idx]), idx)) + .collect(); + scored.sort_unstable_by(|a, b| b.0.total_cmp(&a.0)); + scored.truncate(num_queries); + + let mut regions = Vec::new(); + for (score, idx) in scored { + if score <= THRESHOLD { + continue; + } + let label_id = idx % num_classes; + let q = idx / num_classes; + let cx = boxes[q * 4]; + let cy = boxes[q * 4 + 1]; + let w = boxes[q * 4 + 2]; + let h = boxes[q * 4 + 3]; + // center_to_corners, then scale normalized coords to page points. + let l = (cx - w / 2.0) * page_w; + let t = (cy - h / 2.0) * page_h; + let r = (cx + w / 2.0) * page_w; + let b = (cy + h / 2.0) * page_h; + regions.push(Region { + label: LABELS.get(label_id).copied().unwrap_or("text"), + score, + l, + t, + r, + b, + }); + } + regions +} diff --git a/crates/docling-pdf/src/lib.rs b/crates/docling-pdf/src/lib.rs index acaf23ef..ddc4397b 100644 --- a/crates/docling-pdf/src/lib.rs +++ b/crates/docling-pdf/src/lib.rs @@ -27,15 +27,25 @@ pub mod layout; mod mets; #[cfg(feature = "ml")] mod ocr; +#[cfg(feature = "ocr-prep")] +pub mod ocr_prep; pub mod pdfium_backend; mod reading_order; -#[cfg(feature = "ml")] +// Pure-Rust region resampling (page→1024px box-average, crop→448 bilinear) — +// available to the browser TableFormer path (#157 stage 3), not just `ml`. +#[cfg(feature = "ocr-prep")] pub mod resample; +#[cfg(feature = "ocr-prep")] +pub mod scanned; #[cfg(feature = "ml")] pub mod tableformer; pub mod textparse; -#[cfg(feature = "ml")] -mod tf_match; +#[cfg(feature = "ocr-prep")] +pub mod tf_core; +// docling's TableFormer cell matcher — pure Rust, shared with the browser +// TableFormer path (#157 stage 3). +#[cfg(feature = "ocr-prep")] +pub mod tf_match; pub mod timing; #[cfg(feature = "ml")] diff --git a/crates/docling-pdf/src/ocr.rs b/crates/docling-pdf/src/ocr.rs index 83b9058c..0a2b3f91 100644 --- a/crates/docling-pdf/src/ocr.rs +++ b/crates/docling-pdf/src/ocr.rs @@ -5,103 +5,24 @@ //! is recognised and decoded with CTC — producing [`TextCell`]s the normal //! layout assembly then consumes. This avoids a separate text-detection model. -use std::collections::BTreeMap; - -use image::{imageops, imageops::FilterType, Rgb, RgbImage}; +use image::RgbImage; use ort::session::Session; use ort::value::Tensor; use crate::layout::Region; +// The ONNX-free half (line prep, batching, CTC decode) lives in `ocr_prep` +// so the wasm build shares it verbatim (#79 phase 2). +use crate::ocr_prep::{ + batch_input, decode_row, dict_chars, prep_region_lines, width_batches, PrepLine, REC_HEIGHT, +}; use crate::pdfium_backend::TextCell; -const REC_HEIGHT: u32 = 48; - -/// Cap on lines per recognition run: bounds peak input-tensor memory -/// (16 × 3 × 48 × 2400 px ≈ 22 MB f32) without costing measurable batching -/// benefit — same-width groups are rarely larger. -const REC_BATCH: usize = 16; - -/// A text-line crop prepared for recognition: resized to the fixed model -/// height, normalised to `[-1, 1]`, laid out CHW. -struct PrepLine { - /// Width after the aspect-preserving resize to `REC_HEIGHT`. - w: usize, - /// `3 * REC_HEIGHT * w` values. - data: Vec, -} - -/// Prepare one line crop, or `None` for a degenerate (zero-sized) crop. -fn prep_line(line: &RgbImage) -> Option { - let (w, h) = line.dimensions(); - if w == 0 || h == 0 { - return None; - } - let new_w = ((w as f32) * REC_HEIGHT as f32 / h as f32) - .round() - .clamp(8.0, 2400.0) as u32; - let resized = imageops::resize(line, new_w, REC_HEIGHT, FilterType::Triangle); - let n = (REC_HEIGHT * new_w) as usize; - // Normalise to [-1, 1]: (x/255 - 0.5) / 0.5. - let mut data = vec![0f32; 3 * n]; - for (i, px) in resized.pixels().enumerate() { - data[i] = px[0] as f32 / 127.5 - 1.0; - data[n + i] = px[1] as f32 / 127.5 - 1.0; - data[2 * n + i] = px[2] as f32 / 127.5 - 1.0; - } - Some(PrepLine { - w: new_w as usize, - data, - }) -} - pub struct OcrModel { rec: Session, /// CTC classes: index 0 = blank, 1..=6623 = dictionary, 6624 = space. chars: Vec, } -/// Greedy CTC decode of one row's `(T, C)` probabilities. -fn decode_row(chars: &[String], probs: &[f32], nc: usize) -> String { - let mut out = String::new(); - let mut prev = 0usize; - for row in probs.chunks_exact(nc) { - let mut best = 0usize; - let mut bestv = row[0]; - for (c, &v) in row.iter().enumerate().skip(1) { - if v > bestv { - bestv = v; - best = c; - } - } - if best != prev && best != 0 { - if let Some(ch) = chars.get(best) { - out.push_str(ch); - } - } - prev = best; - } - out -} - -fn luma(p: &Rgb) -> f32 { - 0.299 * p[0] as f32 + 0.587 * p[1] as f32 + 0.114 * p[2] as f32 -} - -/// Layout labels whose content is recognised as running text. -fn is_text_label(label: &str) -> bool { - matches!( - label, - "text" - | "title" - | "section_header" - | "list_item" - | "caption" - | "footnote" - | "code" - | "formula" - ) -} - /// OCR recognition language: which PP-OCRv3 model + dictionary pair runs. /// /// The default is **English** (`models/ocr_rec_en.onnx` + `models/en_dict.txt`): @@ -194,10 +115,10 @@ impl OcrModel { .map_err(|e| format!("ocr: load {rec_path}: {e}"))?; let dict = std::fs::read_to_string(&dict_path) .map_err(|e| format!("ocr: read dict {dict_path}: {e}"))?; - let mut chars = vec![String::new()]; // blank at 0 - chars.extend(dict.lines().map(|s| s.to_string())); - chars.push(" ".to_string()); - Ok(Self { rec, chars }) + Ok(Self { + rec, + chars: dict_chars(&dict), + }) } /// Recognise a batch of prepared *same-width* lines in one session run. @@ -208,13 +129,14 @@ impl OcrModel { /// the scanned corpus), whereas width-padding leaks into the real /// timesteps through the model's global-attention blocks and measurably /// changes low-confidence characters. - fn recognize_batch(&mut self, w: usize, batch: &[&PrepLine]) -> Result, String> { - let n = batch.len(); - let hw = REC_HEIGHT as usize * w; - let mut data = vec![0f32; n * 3 * hw]; - for (i, pl) in batch.iter().enumerate() { - data[i * 3 * hw..(i + 1) * 3 * hw].copy_from_slice(&pl.data); - } + fn recognize_batch( + &mut self, + w: usize, + chunk: &[usize], + lines: &[PrepLine], + ) -> Result, String> { + let n = chunk.len(); + let data = batch_input(w, chunk, lines); let input = Tensor::from_array(([n, 3, REC_HEIGHT as usize, w], data)) .map_err(|e| format!("ocr: input tensor: {e}"))?; let outputs = self @@ -245,51 +167,16 @@ impl OcrModel { regions: &[Region], scale: f32, ) -> Result, String> { - let (iw, ih) = img.dimensions(); - // Gather every line crop on the page first, so equal-width lines can - // share a recognition run regardless of which region they came from. - let mut bboxes: Vec<(f32, f32, f32, f32)> = Vec::new(); - let mut lines: Vec = Vec::new(); - for region in regions { - if !is_text_label(region.label) { - continue; - } - let l = (region.l * scale).max(0.0) as u32; - let t = (region.t * scale).max(0.0) as u32; - let r = ((region.r * scale).max(0.0) as u32).min(iw); - let b = ((region.b * scale).max(0.0) as u32).min(ih); - if r <= l || b <= t { - continue; - } - let crop = imageops::crop_imm(img, l, t, r - l, b - t).to_image(); - for (lx, ly, rx, ry) in segment_lines(&crop) { - let line = imageops::crop_imm(&crop, lx, ly, rx - lx, ry - ly).to_image(); - let Some(pl) = prep_line(&line) else { - continue; - }; - bboxes.push(( - (l + lx) as f32 / scale, - (t + ly) as f32 / scale, - (l + rx) as f32 / scale, - (t + ry) as f32 / scale, - )); - lines.push(pl); - } - } + // Gather every line crop on the page first (shared with the browser + // path), so equal-width lines can share a recognition run regardless + // of which region they came from. + let (bboxes, lines) = prep_region_lines(img, regions, scale); - // Group page-order line indices by exact width (BTreeMap: run order is - // deterministic) and recognise each group batched. - let mut by_width: BTreeMap> = BTreeMap::new(); - for (ix, pl) in lines.iter().enumerate() { - by_width.entry(pl.w).or_default().push(ix); - } + // Deterministic width-batching (shared with the wasm path). let mut texts = vec![String::new(); lines.len()]; - for (w, ixs) in by_width { - for chunk in ixs.chunks(REC_BATCH) { - let batch: Vec<&PrepLine> = chunk.iter().map(|&i| &lines[i]).collect(); - for (&i, text) in chunk.iter().zip(self.recognize_batch(w, &batch)?) { - texts[i] = text; - } + for (w, chunk) in width_batches(&lines) { + for (&i, text) in chunk.iter().zip(self.recognize_batch(w, &chunk, &lines)?) { + texts[i] = text; } } @@ -305,67 +192,3 @@ impl OcrModel { Ok(cells) } } - -/// Split a region crop into text lines via a horizontal ink-projection profile. -/// Returns tight `(l, t, r, b)` boxes in crop pixels. -fn segment_lines(crop: &RgbImage) -> Vec<(u32, u32, u32, u32)> { - let (w, h) = crop.dimensions(); - if w == 0 || h == 0 { - return Vec::new(); - } - let mean: f32 = crop.pixels().map(luma).sum::() / (w * h) as f32; - let thresh = mean * 0.7; // ink = noticeably darker than the page average - let min_ink = ((w as f32) * 0.005).max(1.0) as u32; - - let mut profile = vec![0u32; h as usize]; - for y in 0..h { - let mut row = 0u32; - for x in 0..w { - if luma(crop.get_pixel(x, y)) < thresh { - row += 1; - } - } - profile[y as usize] = row; - } - - // Maximal runs of text rows, separated by (near-)blank rows. - let mut runs: Vec<(u32, u32)> = Vec::new(); - let mut start: Option = None; - for y in 0..h { - let text = profile[y as usize] >= min_ink; - if text && start.is_none() { - start = Some(y); - } else if !text { - if let Some(s) = start.take() { - if y - s >= 4 { - runs.push((s, y)); - } - } - } - } - if let Some(s) = start { - if h - s >= 4 { - runs.push((s, h)); - } - } - - // Tighten each line to its horizontal ink bounds. - runs.into_iter() - .map(|(t, b)| { - let (mut l, mut r) = (w, 0u32); - for y in t..b { - for x in 0..w { - if luma(crop.get_pixel(x, y)) < thresh { - l = l.min(x); - r = r.max(x + 1); - } - } - } - if l >= r { - (0, t, w, b) - } else { - (l, t, r, b) - } - }) - .collect() -} diff --git a/crates/docling-pdf/src/ocr_prep.rs b/crates/docling-pdf/src/ocr_prep.rs new file mode 100644 index 00000000..1b1da9b1 --- /dev/null +++ b/crates/docling-pdf/src/ocr_prep.rs @@ -0,0 +1,437 @@ +//! ONNX-free half of the PP-OCRv3 recognition pipeline: everything before +//! and after the `session.run` call — line segmentation, crop preparation, +//! width-batching, CTC decoding, dictionary handling. +//! +//! Split out of `ocr.rs` (which keeps the `ort` session) so the browser build +//! can reuse it (issue #79 phase 2): `docling-wasm` runs these exact +//! functions and delegates only the inference call to ONNX Runtime Web on +//! the JS side. Keeping one implementation is what makes the wasm output +//! byte-comparable to the native CPU path — any drift then comes from the +//! runtime, not from pre/post-processing. + +use image::{imageops, imageops::FilterType, Rgb, RgbImage}; + +/// PP-OCRv3's fixed input height. +pub const REC_HEIGHT: u32 = 48; + +/// Cap on lines per recognition run: bounds peak input-tensor memory +/// (16 × 3 × 48 × 2400 px ≈ 22 MB f32) without costing measurable batching +/// benefit — same-width groups are rarely larger. +pub const REC_BATCH: usize = 16; + +/// A text-line crop prepared for recognition: resized to the fixed model +/// height, normalised to `[-1, 1]`, laid out CHW. +pub struct PrepLine { + /// Width after the aspect-preserving resize to [`REC_HEIGHT`]. + pub w: usize, + /// `3 * REC_HEIGHT * w` values. + pub data: Vec, +} + +/// Prepare one line crop, or `None` for a degenerate (zero-sized) crop. +pub fn prep_line(line: &RgbImage) -> Option { + let (w, h) = line.dimensions(); + if w == 0 || h == 0 { + return None; + } + let new_w = ((w as f32) * REC_HEIGHT as f32 / h as f32) + .round() + .clamp(8.0, 2400.0) as u32; + let resized = imageops::resize(line, new_w, REC_HEIGHT, FilterType::Triangle); + let n = (REC_HEIGHT * new_w) as usize; + // Normalise to [-1, 1]: (x/255 - 0.5) / 0.5. + let mut data = vec![0f32; 3 * n]; + for (i, px) in resized.pixels().enumerate() { + data[i] = px[0] as f32 / 127.5 - 1.0; + data[n + i] = px[1] as f32 / 127.5 - 1.0; + data[2 * n + i] = px[2] as f32 / 127.5 - 1.0; + } + Some(PrepLine { + w: new_w as usize, + data, + }) +} + +/// The CTC class table for a recognition dictionary file: index 0 = blank, +/// then one class per dictionary line, then the space class. +pub fn dict_chars(dict: &str) -> Vec { + let mut chars = vec![String::new()]; // blank at 0 + chars.extend(dict.lines().map(|s| s.to_string())); + chars.push(" ".to_string()); + chars +} + +/// Greedy CTC decode of one row's `(T, C)` probabilities. +pub fn decode_row(chars: &[String], probs: &[f32], nc: usize) -> String { + let mut out = String::new(); + let mut prev = 0usize; + for row in probs.chunks_exact(nc) { + let mut best = 0usize; + let mut bestv = row[0]; + for (c, &v) in row.iter().enumerate().skip(1) { + if v > bestv { + bestv = v; + best = c; + } + } + if best != prev && best != 0 { + if let Some(ch) = chars.get(best) { + out.push_str(ch); + } + } + prev = best; + } + out +} + +pub(crate) fn luma(p: &Rgb) -> f32 { + 0.299 * p[0] as f32 + 0.587 * p[1] as f32 + 0.114 * p[2] as f32 +} + +/// Split a region crop into text lines via a horizontal ink-projection profile. +/// Returns tight `(l, t, r, b)` boxes in crop pixels. +pub fn segment_lines(crop: &RgbImage) -> Vec<(u32, u32, u32, u32)> { + let (w, h) = crop.dimensions(); + if w == 0 || h == 0 { + return Vec::new(); + } + let mean: f32 = crop.pixels().map(luma).sum::() / (w * h) as f32; + let thresh = mean * 0.7; // ink = noticeably darker than the page average + let min_ink = ((w as f32) * 0.005).max(1.0) as u32; + + let mut profile = vec![0u32; h as usize]; + for y in 0..h { + let mut row = 0u32; + for x in 0..w { + if luma(crop.get_pixel(x, y)) < thresh { + row += 1; + } + } + profile[y as usize] = row; + } + + // Maximal runs of text rows, separated by (near-)blank rows. + let mut runs: Vec<(u32, u32)> = Vec::new(); + let mut start: Option = None; + for y in 0..h { + let text = profile[y as usize] >= min_ink; + if text && start.is_none() { + start = Some(y); + } else if !text { + if let Some(s) = start.take() { + if y - s >= 4 { + runs.push((s, y)); + } + } + } + } + if let Some(s) = start { + if h - s >= 4 { + runs.push((s, h)); + } + } + + // Tighten each line to its horizontal ink bounds. + runs.into_iter() + .map(|(t, b)| { + let (mut l, mut r) = (w, 0u32); + for y in t..b { + for x in 0..w { + if luma(crop.get_pixel(x, y)) < thresh { + l = l.min(x); + r = r.max(x + 1); + } + } + } + if l >= r { + (0, t, w, b) + } else { + (l, t, r, b) + } + }) + .collect() +} + +/// Layout labels whose content is recognised as running text. +pub fn is_text_label(label: &str) -> bool { + matches!( + label, + "text" + | "title" + | "section_header" + | "list_item" + | "caption" + | "footnote" + | "code" + | "formula" + ) +} + +/// A line's page-point bounding box, `(l, t, r, b)`. +pub type LineBox = (f32, f32, f32, f32); + +/// Gather every text-region line crop on a page, in page order: crop each +/// text region (page points × `scale` → image px), split it into lines, prep +/// each line, and keep the line's page-point bbox. The exact gathering the +/// native `ocr_page` does — shared so the browser path produces the same +/// cells given the same probabilities. +pub fn prep_region_lines( + img: &RgbImage, + regions: &[crate::layout::Region], + scale: f32, +) -> (Vec, Vec) { + let (iw, ih) = img.dimensions(); + let mut bboxes = Vec::new(); + let mut lines = Vec::new(); + for region in regions { + if !is_text_label(region.label) { + continue; + } + let l = (region.l * scale).max(0.0) as u32; + let t = (region.t * scale).max(0.0) as u32; + let r = ((region.r * scale).max(0.0) as u32).min(iw); + let b = ((region.b * scale).max(0.0) as u32).min(ih); + if r <= l || b <= t { + continue; + } + let crop = imageops::crop_imm(img, l, t, r - l, b - t).to_image(); + for (lx, ly, rx, ry) in segment_lines(&crop) { + let line = imageops::crop_imm(&crop, lx, ly, rx - lx, ry - ly).to_image(); + let Some(pl) = prep_line(&line) else { + continue; + }; + bboxes.push(( + (l + lx) as f32 / scale, + (t + ly) as f32 / scale, + (l + rx) as f32 / scale, + (t + ry) as f32 / scale, + )); + lines.push(pl); + } + } + (bboxes, lines) +} + +/// Split a text line into word tokens by a vertical ink-projection profile: +/// runs of ink columns separated by whitespace wider than ~0.6× the line +/// height (an inter-word/-column gap, not an inter-character one). Returns tight +/// `(l, 0, r, h)` boxes in line-crop pixels. The browser table path recognizes +/// these individually so each word carries its own box for column matching — +/// the native pipeline gets word boxes from the pdfium text layer instead. +pub fn segment_words(line: &RgbImage) -> Vec<(u32, u32, u32, u32)> { + let (w, h) = line.dimensions(); + if w == 0 || h == 0 { + return Vec::new(); + } + let mean: f32 = line.pixels().map(luma).sum::() / (w * h) as f32; + let thresh = mean * 0.7; + let mut col_ink = vec![0u32; w as usize]; + for y in 0..h { + for x in 0..w { + if luma(line.get_pixel(x, y)) < thresh { + col_ink[x as usize] += 1; + } + } + } + let min_gap = ((h as f32) * 0.6).max(4.0) as u32; + let mut words = Vec::new(); + let mut start: Option = None; + let mut last_ink = 0u32; + let mut gap = 0u32; + for x in 0..w { + if col_ink[x as usize] > 0 { + if start.is_none() { + start = Some(x); + } + last_ink = x; + gap = 0; + } else if let Some(s) = start { + gap += 1; + if gap >= min_gap { + words.push((s, 0, last_ink + 1, h)); + start = None; + } + } + } + if let Some(s) = start { + words.push((s, 0, last_ink + 1, h)); + } + words +} + +/// Gather word crops from a page's *table* regions (browser table path, #157 +/// stage 3): crop each table region, split it into lines, split each line into +/// words ([`segment_words`]), prep each word for recognition, and keep the +/// word's page-point bbox. Recognizing table interiors is what gives the cell +/// matcher the word boxes it needs — the native pipeline reads those from +/// pdfium's text layer, which the browser doesn't have. NOT used by native. +pub fn prep_table_words( + img: &RgbImage, + regions: &[crate::layout::Region], + scale: f32, +) -> (Vec, Vec) { + let (iw, ih) = img.dimensions(); + let mut bboxes = Vec::new(); + let mut lines = Vec::new(); + for region in regions { + if !crate::assemble::is_table_like(region.label) { + continue; + } + let l = (region.l * scale).max(0.0) as u32; + let t = (region.t * scale).max(0.0) as u32; + let r = ((region.r * scale).max(0.0) as u32).min(iw); + let b = ((region.b * scale).max(0.0) as u32).min(ih); + if r <= l || b <= t { + continue; + } + let crop = imageops::crop_imm(img, l, t, r - l, b - t).to_image(); + for (lx, ly, rx, ry) in segment_lines(&crop) { + let line = imageops::crop_imm(&crop, lx, ly, rx - lx, ry - ly).to_image(); + for (wx0, _, wx1, _) in segment_words(&line) { + let word = imageops::crop_imm(&line, wx0, 0, wx1 - wx0, ry - ly).to_image(); + let Some(pl) = prep_line(&word) else { + continue; + }; + bboxes.push(( + (l + lx + wx0) as f32 / scale, + (t + ly) as f32 / scale, + (l + lx + wx1) as f32 / scale, + (t + ry) as f32 / scale, + )); + lines.push(pl); + } + } + } + (bboxes, lines) +} + +/// Normalize an image to the scan polarity every stage assumes — dark ink +/// on light paper (the segmentation threshold and the recognition model's +/// training data both bake it in): a predominantly dark page (mean luma +/// below mid-gray — a dark-mode screenshot, an inverted scan) is inverted. +/// Browser-path helper; the native pipeline never calls it (its input is +/// scanned paper, and the conformance baseline stays untouched). +pub fn normalize_polarity(mut img: RgbImage) -> RgbImage { + let (w, h) = img.dimensions(); + if w == 0 || h == 0 { + return img; + } + let mean: f32 = img.pixels().map(luma).sum::() / (w * h) as f32; + if mean < 128.0 { + for px in img.pixels_mut() { + px.0 = [255 - px.0[0], 255 - px.0[1], 255 - px.0[2]]; + } + } + img +} + +/// Whole-image line preparation for the browser OCR path (no layout model: +/// the page itself is the single text region). Returns page-order prepared +/// lines; callers that need geometry use [`segment_lines`] directly. +pub fn prep_page_lines(img: &RgbImage) -> Vec { + segment_lines(img) + .into_iter() + .filter_map(|(l, t, r, b)| { + let line = imageops::crop_imm(img, l, t, r - l, b - t).to_image(); + prep_line(&line) + }) + .collect() +} + +/// Deterministic recognition batching: page-order line indices grouped by +/// exact width (equal widths share a run — bit-identical to one-at-a-time +/// recognition, see `ocr.rs`), each group split into [`REC_BATCH`] chunks. +pub fn width_batches(lines: &[PrepLine]) -> Vec<(usize, Vec)> { + let mut by_width: std::collections::BTreeMap> = + std::collections::BTreeMap::new(); + for (ix, pl) in lines.iter().enumerate() { + by_width.entry(pl.w).or_default().push(ix); + } + let mut out = Vec::new(); + for (w, ixs) in by_width { + for chunk in ixs.chunks(REC_BATCH) { + out.push((w, chunk.to_vec())); + } + } + out +} + +/// Pack one width-batch into the model's `(N, 3, H, W)` input buffer. +pub fn batch_input(w: usize, chunk: &[usize], lines: &[PrepLine]) -> Vec { + let hw = REC_HEIGHT as usize * w; + let mut data = vec![0f32; chunk.len() * 3 * hw]; + for (i, &ix) in chunk.iter().enumerate() { + data[i * 3 * hw..(i + 1) * 3 * hw].copy_from_slice(&lines[ix].data); + } + data +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A synthetic "page": white background with two black text bars. + fn page() -> RgbImage { + let mut img = RgbImage::from_pixel(200, 100, Rgb([255, 255, 255])); + for y in 20..30 { + for x in 10..190 { + img.put_pixel(x, y, Rgb([0, 0, 0])); + } + } + for y in 60..72 { + for x in 10..120 { + img.put_pixel(x, y, Rgb([0, 0, 0])); + } + } + img + } + + #[test] + fn segments_and_preps_page_lines() { + let lines = prep_page_lines(&page()); + assert_eq!(lines.len(), 2); + for pl in &lines { + assert_eq!(pl.data.len(), 3 * REC_HEIGHT as usize * pl.w); + } + // Different aspect ratios → different widths → separate batches. + let batches = width_batches(&lines); + assert_eq!(batches.len(), 2); + let (w0, chunk0) = &batches[0]; + assert_eq!( + batch_input(*w0, chunk0, &lines).len(), + 3 * REC_HEIGHT as usize * w0 + ); + } + + #[test] + fn dark_mode_pages_normalize_to_scan_polarity() { + // The same two-bar page, inverted (light text on dark) — the raw + // segmentation misfires (it thresholds the dark *background* as ink); + // polarity normalization recovers the true structure. + let mut dark = page(); + for px in dark.pixels_mut() { + px.0 = [255 - px.0[0], 255 - px.0[1], 255 - px.0[2]]; + } + assert_ne!(prep_page_lines(&dark).len(), 2); + let fixed = normalize_polarity(dark); + assert_eq!(prep_page_lines(&fixed).len(), 2); + // A light page passes through untouched. + let light = page(); + assert_eq!(normalize_polarity(light.clone()), light); + } + + #[test] + fn ctc_decode_collapses_repeats_and_blanks() { + // 3 classes: blank, "a", "b"; timesteps a a blank b b → "ab". + let chars = dict_chars("a\nb"); + assert_eq!(chars.len(), 4); // blank, a, b, space + let probs = [ + 0.1, 0.8, 0.1, 0.0, // a + 0.1, 0.8, 0.1, 0.0, // a (repeat collapses) + 0.9, 0.05, 0.05, 0.0, // blank + 0.1, 0.1, 0.8, 0.0, // b + 0.1, 0.1, 0.8, 0.0, // b (repeat collapses) + ]; + assert_eq!(decode_row(&chars, &probs, 4), "ab"); + } +} diff --git a/crates/docling-pdf/src/pdfium_backend.rs b/crates/docling-pdf/src/pdfium_backend.rs index 10688d32..f2d8a932 100644 --- a/crates/docling-pdf/src/pdfium_backend.rs +++ b/crates/docling-pdf/src/pdfium_backend.rs @@ -53,6 +53,28 @@ pub struct PdfPage { pub links: Vec, } +impl PdfPage { + /// A page built from recognized cells alone — the browser pipeline's + /// shape (#157), where the bitmap lives on the JS side. Exists so callers + /// compile identically with and without the `ml` feature: under a + /// feature-unified workspace build the struct carries the `image` field, + /// which a plain literal in a non-`ml` consumer can't spell. + #[cfg(feature = "ocr-prep")] + pub fn from_cells(width: f32, height: f32, scale: f32, cells: Vec) -> Self { + Self { + width, + height, + scale, + cells, + code_cells: Vec::new(), + word_cells: Vec::new(), + #[cfg(feature = "ml")] + image: RgbImage::new(0, 0), + links: Vec::new(), + } + } +} + /// A PDF link annotation: its rectangle (top-left page coordinates, matching /// [`TextCell`]) and target URI. #[derive(Debug, Clone)] diff --git a/crates/docling-pdf/src/scanned.rs b/crates/docling-pdf/src/scanned.rs new file mode 100644 index 00000000..fbfb1cc8 --- /dev/null +++ b/crates/docling-pdf/src/scanned.rs @@ -0,0 +1,67 @@ +//! Scanned-page assembly facade for the browser build (#157 stage 2). +//! +//! The native `Worker::process` chain, re-exposed without the `ml` feature so +//! `docling-wasm` can run it around JS-delegated inference: refine the raw +//! layout detections exactly like the native pipeline, then assemble the page +//! with the geometric table fallback (the lite profile — TableFormer is +//! stage 3) and no enrichments. One shared implementation, one behavior. + +use docling_core::{DoclingDocument, Node}; + +/// One assembled page: its nodes and `(anchor, href)` hyperlink pairs. +pub type AssembledPage = (Vec, Vec<(String, String)>); + +use crate::layout::{label_threshold, Region}; +use crate::pdfium_backend::{PdfPage, TextCell}; + +/// The native pipeline's region-refinement chain, in its exact order: +/// per-label score thresholds → overlap resolution → orphan text regions for +/// detector-missed cells → false-picture drops → contained-regular drops. +/// `cells` is the page's (possibly empty, pre-OCR) text-cell set. +pub fn refine_regions( + regions: Vec, + cells: &[TextCell], + page_w: f32, + page_h: f32, +) -> Vec { + let mut regions = regions; + regions.retain(|r| r.score >= label_threshold(r.label)); + let mut regions = crate::assemble::resolve(regions); + crate::assemble::add_orphan_regions(&mut regions, cells); + crate::assemble::drop_false_pictures(&mut regions, cells, page_w, page_h); + crate::assemble::drop_contained_regulars(&mut regions); + regions +} + +/// Assemble one refined page — geometric tables (no TableFormer), no +/// enrichments — into its nodes and hyperlink pairs. +pub fn assemble_page(page: &PdfPage, regions: Vec) -> AssembledPage { + assemble_page_with_tables(page, regions, vec![None; 0].into_iter().collect()) +} + +/// Like [`assemble_page`] but with pre-computed table rows per region: +/// `table_rows[i] = Some(rows)` for a table region whose structure the browser +/// TableFormer path (#157 stage 3) resolved, `None` for the geometric fallback. +/// An empty `table_rows` means "all geometric" (what [`assemble_page`] passes). +/// Same assembly as the native pipeline. +pub fn assemble_page_with_tables( + page: &PdfPage, + regions: Vec, + mut table_rows: Vec>>>, +) -> AssembledPage { + table_rows.resize(regions.len(), None); + let enrich = vec![None; regions.len()]; + crate::assemble::assemble_page(page, regions, &table_rows, &enrich) +} + +/// Stitch per-page results into a document: cross-page paragraph +/// continuations merge exactly like the native pipeline's final pass. +pub fn finish_document(name: &str, pages: Vec) -> DoclingDocument { + let mut doc = DoclingDocument::new(name); + for (nodes, links) in pages { + doc.nodes.extend(nodes); + doc.links.extend(links); + } + crate::assemble::merge_continuations(&mut doc.nodes); + doc +} diff --git a/crates/docling-pdf/src/tableformer.rs b/crates/docling-pdf/src/tableformer.rs index aad55ba6..f8693cbc 100644 --- a/crates/docling-pdf/src/tableformer.rs +++ b/crates/docling-pdf/src/tableformer.rs @@ -5,52 +5,23 @@ //! docling runs). See docs/PDF_CONFORMANCE.md. use crate::pdfium_backend::TextCell; +// The ONNX-free half (preprocessing, structure corrections, bbox bookkeeping, +// span merge, OTSL→grid) lives in tf_core so the browser build (#157 stage 3) +// runs the same logic; this file owns the three `ort` sessions and the +// owned-value KV-cache fast path. +use crate::tf_core::{ + argmax, build_table_cells, correct, merge_spans, preprocess_input, BboxBook, TableCell, END, + MAX_STEPS, START, UCEL, +}; use image::RgbImage; use ort::session::Session; use ort::value::{DynValue, Tensor}; -const SIDE: u32 = 448; -// Verbatim from docling's tm_config.json image_normalization (more digits than -// f32 holds; kept exact for provenance). -#[allow(clippy::excessive_precision)] -const MEAN: [f32; 3] = [0.94247851, 0.94254675, 0.94292611]; -#[allow(clippy::excessive_precision)] -const STD: [f32; 3] = [0.17910956, 0.17940403, 0.17931663]; -const MAX_STEPS: usize = 1024; +const SIDE: usize = crate::tf_core::SIDE as usize; +const EMBED_DIM: usize = crate::tf_core::EMBED_DIM; /// Decoder geometry, fixed by the exported TableModel04_rs graph: the cached /// decoder threads a `[N_LAYERS, past, 1, EMBED_DIM]` per-layer state cache. const N_LAYERS: usize = 6; -const EMBED_DIM: usize = 512; - -/// OTSL structure tokens (TableModel04_rs wordmap indices). -pub const START: i64 = 2; -pub const END: i64 = 3; -pub const ECEL: i64 = 4; // empty cell -pub const FCEL: i64 = 5; // full (content) cell -pub const LCEL: i64 = 6; // left-looking: extends the cell to its left (colspan) -pub const UCEL: i64 = 7; // up-looking: extends the cell above (rowspan) -pub const XCEL: i64 = 8; // cross: spans both ways -pub const NL: i64 = 9; // new row -pub const CHED: i64 = 10; // column header -pub const RHED: i64 = 11; // row header -pub const SROW: i64 = 12; // section row - -/// A predicted table cell: an OTSL grid position (with spans) + its box in the -/// 448 image normalized cxcywh, the OTSL tag, and the bbox decoder's cell -/// class (docling's `cell_class`; 2 = full, ≤1 = predicted empty). -#[derive(Debug, Clone)] -pub struct TableCell { - pub row: usize, - pub col: usize, - pub colspan: usize, - pub rowspan: usize, - pub tag: i64, - pub class: i64, - pub cx: f32, - pub cy: f32, - pub w: f32, - pub h: f32, -} pub struct TableFormer { encoder: Session, @@ -399,8 +370,9 @@ impl TableFormer { /// Predict the OTSL structure-token sequence for a table-region image. pub fn predict_otsl(&mut self, img: &RgbImage) -> Result, String> { let enc = self.encode(img)?; - // The two structure corrections mirror docling's `predict` exactly — note - // its `line_num` is never incremented, so `xcel→lcel` applies on every row. + // Structure corrections live in tf_core::correct (shared with the wasm + // path); docling's line_num is never incremented, so xcel→lcel fires on + // every row. let mut tags: Vec = vec![START]; let mut out: Vec = Vec::new(); let mut prev_ucel = false; @@ -408,13 +380,7 @@ impl TableFormer { let empty = self.empty_cache()?; while out.len() < MAX_STEPS { let (raw, _hidden) = self.decode_step(&tags, &enc, &mut cache, &empty)?; - let mut tag = raw; - if tag == XCEL { - tag = LCEL; - } - if prev_ucel && tag == LCEL { - tag = FCEL; - } + let tag = correct(raw, prev_ucel); if tag == END { break; } @@ -433,58 +399,21 @@ impl TableFormer { pub fn predict_table_structure(&mut self, img: &RgbImage) -> Result, String> { let enc = self.encode(img)?; - let mut tags: Vec = vec![START]; - let mut otsl: Vec = Vec::new(); - let mut hiddens: Vec = Vec::new(); // flattened [n, 512] - let mut n = 0usize; - let mut prev_ucel = false; - let mut skip = true; // first tag after is skipped - let mut first_lcel = true; - let mut bbox_ind = 0usize; - let mut cur_bbox_ind = 0usize; - let mut merge: std::collections::HashMap = std::collections::HashMap::new(); + // The autoregressive loop's bbox bookkeeping lives in tf_core::BboxBook + // (shared with the wasm path); this loop only steps the decoder. + let mut book = BboxBook::new(); let mut cache = DecodeCache::default(); let empty = self.empty_cache()?; - while otsl.len() < MAX_STEPS { - let (raw, hidden) = self.decode_step(&tags, &enc, &mut cache, &empty)?; - let mut tag = raw; - if tag == XCEL { - tag = LCEL; - } - if prev_ucel && tag == LCEL { - tag = FCEL; - } - if tag == END { + while book.otsl.len() < MAX_STEPS { + let (raw, hidden) = self.decode_step(&book.tags, &enc, &mut cache, &empty)?; + if !book.step(raw, &hidden) { break; } - // docling's tag_H_buf / bboxes_to_merge bookkeeping. - if !skip && matches!(tag, FCEL | ECEL | CHED | RHED | SROW | NL | UCEL) { - hiddens.extend_from_slice(&hidden); - n += 1; - if !first_lcel { - merge.insert(cur_bbox_ind, bbox_ind as i64); - } - bbox_ind += 1; - } - if tag != LCEL { - first_lcel = true; - } else if first_lcel { - hiddens.extend_from_slice(&hidden); - n += 1; - first_lcel = false; - cur_bbox_ind = bbox_ind; - merge.insert(cur_bbox_ind, -1); - bbox_ind += 1; - } - skip = matches!(tag, NL | UCEL | XCEL); - prev_ucel = tag == UCEL; - otsl.push(tag); - tags.push(tag); } - if n == 0 { + if book.n == 0 { return Ok(Vec::new()); } - let tag_h = Tensor::from_array(([n, 512usize], hiddens)) + let tag_h = Tensor::from_array(([book.n, EMBED_DIM], std::mem::take(&mut book.hiddens))) .map_err(|e| format!("tableformer: tag_h: {e}"))?; let bout = self .bbox @@ -502,8 +431,8 @@ impl TableFormer { .try_extract_tensor::() .map_err(|e| format!("tableformer: classes: {e}"))?; let classes: Vec = craw.chunks_exact(3).map(|c| argmax(c) as i64).collect(); - let (merged, merged_classes) = merge_spans(&boxes, &classes, &merge); - Ok(build_table_cells(&otsl, &merged, &merged_classes)) + let (merged, merged_classes) = merge_spans(&boxes, &classes, &book.merge); + Ok(build_table_cells(&book.otsl, &merged, &merged_classes)) } /// Predict a table region's Markdown grid: crop the region (docling's @@ -548,317 +477,10 @@ impl TableFormer { if cells.is_empty() { return None; } - // Words that belong to the table: non-empty text, ≥80 % of the word's - // area inside the table region (docling's `get_cells_in_bbox` ios test). - // Ids stay the page-level word indices so text joins in stream order. - let table_words: Vec = words - .iter() - .enumerate() - .filter(|(_, w)| !w.text.trim().is_empty()) - .filter_map(|(wi, w)| { - let (l, t, r, b) = (w.l as f64, w.t as f64, w.r as f64, w.b as f64); - let area = (r - l) * (b - t); - let iw = (r.min(region[2] as f64) - l.max(region[0] as f64)).max(0.0); - let ih = (b.min(region[3] as f64) - t.max(region[1] as f64)).max(0.0); - if area > 0.0 && iw * ih / area > 0.8 { - Some(crate::tf_match::PdfWord { - id: wi, - bbox: [l, t, r, b], - text: w.text.trim().to_string(), - }) - } else { - None - } - }) - .collect(); - - if !table_words.is_empty() && !simple_match() { - return docling_match_rows(&cells, region, &table_words, words); - } - - let (rw, rh) = (region[2] - region[0], region[3] - region[1]); - - // Cell boxes in page points (top-left), aligned with `cells`. - let boxes: Vec<[f32; 4]> = cells - .iter() - .map(|c| { - [ - region[0] + (c.cx - c.w / 2.0) * rw, - region[1] + (c.cy - c.h / 2.0) * rh, - region[0] + (c.cx + c.w / 2.0) * rw, - region[1] + (c.cy + c.h / 2.0) * rh, - ] - }) - .collect(); - - // Assign each word to the cell it overlaps most (intersection / word area). - let mut cell_words: Vec> = vec![Vec::new(); cells.len()]; - for (wi, w) in words.iter().enumerate() { - let wa = ((w.r - w.l) * (w.b - w.t)).max(1.0); - let mut best: Option<(f32, usize)> = None; - for (ci, b) in boxes.iter().enumerate() { - let ix = (w.r.min(b[2]) - w.l.max(b[0])).max(0.0); - let iy = (w.b.min(b[3]) - w.t.max(b[1])).max(0.0); - let io = ix * iy / wa; - if io > 0.0 && best.is_none_or(|(bo, _)| io > bo) { - best = Some((io, ci)); - } - } - if let Some((_, ci)) = best { - cell_words[ci].push(wi); - } - } - - let num_rows = cells.iter().map(|c| c.row + c.rowspan).max().unwrap_or(0); - let num_cols = cells.iter().map(|c| c.col + c.colspan).max().unwrap_or(0); - if num_rows == 0 || num_cols == 0 { - return None; - } - let mut grid = vec![vec![String::new(); num_cols]; num_rows]; - for (ci, c) in cells.iter().enumerate() { - // Keep words in text-stream order (the order they were collected = - // their word index), matching docling's cell text assembly — geometric - // re-sorting scrambles wrapped cells (`Inference time (secs)`). - let wis = std::mem::take(&mut cell_words[ci]); - let text = wis - .iter() - .map(|&i| words[i].text.trim()) - .collect::>() - .join(" "); - let text = normalize_cell_text(text); - // Spanned cells repeat their text across the covered grid positions. - for row in grid.iter_mut().skip(c.row).take(c.rowspan) { - for cell in row.iter_mut().skip(c.col).take(c.colspan) { - *cell = text.clone(); - } - } - } - Some(grid) - } -} - -/// Append one JSON line per table into `/tf_match_dump.jsonl` with the -/// exact matcher inputs (hand-rolled JSON to avoid a serde dependency). -fn dump_match_inputs( - dir: &str, - tf_cells: &[crate::tf_match::TfCell], - words: &[crate::tf_match::PdfWord], -) { - use std::io::Write; - let cells: Vec = tf_cells - .iter() - .map(|c| { - format!( - r#"{{"bbox":[{},{},{},{}],"cell_id":{},"row_id":{},"column_id":{},"cell_class":{},"colspan_val":{},"rowspan_val":{}}}"#, - c.bbox[0], c.bbox[1], c.bbox[2], c.bbox[3], - c.cell_id, c.row_id, c.column_id, c.cell_class, - c.colspan_val, c.rowspan_val - ) - }) - .collect(); - let ws: Vec = words - .iter() - .map(|w| { - format!( - r#"{{"id":{},"bbox":[{},{},{},{}],"text":{}}}"#, - w.id, - w.bbox[0], - w.bbox[1], - w.bbox[2], - w.bbox[3], - serde_json_escape(&w.text) - ) - }) - .collect(); - let line = format!( - r#"{{"table_cells":[{}],"pdf_cells":[{}]}}"#, - cells.join(","), - ws.join(",") - ); - if let Ok(mut f) = std::fs::OpenOptions::new() - .create(true) - .append(true) - .open(format!("{dir}/tf_match_dump.jsonl")) - { - let _ = writeln!(f, "{line}"); - } -} - -/// Minimal JSON string escaping for the parity dump. -fn serde_json_escape(s: &str) -> String { - let mut out = String::with_capacity(s.len() + 2); - out.push('"'); - for ch in s.chars() { - match ch { - '"' => out.push_str("\\\""), - '\\' => out.push_str("\\\\"), - '\n' => out.push_str("\\n"), - '\r' => out.push_str("\\r"), - '\t' => out.push_str("\\t"), - c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)), - c => out.push(c), - } - } - out.push('"'); - out -} - -/// `DOCLING_RS_TF_SIMPLE_MATCH=1` reverts to the pre-#60 best-overlap word -/// assignment (A/B escape hatch for the docling matching post-processor). -fn simple_match() -> bool { - std::env::var("DOCLING_RS_TF_SIMPLE_MATCH").is_ok_and(|v| !v.is_empty() && v != "0") -} - -/// docling glues `@` to whatever follows it (`mAP @0.5`, an email): -/// the PDF's word cells split `@` from the next token, and joining them -/// with a space would widen the cell and — via the column pad — shift -/// every row of the table. The groundtruth never contains "@ ", so this -/// is always the right normalization. -fn normalize_cell_text(text: String) -> String { - text.replace("@ ", "@") -} - -/// docling's matched-cell grid assembly (`tf_predictor.predict` with -/// `do_cell_matching=True`): run the ported matching post-processor, group the -/// word→cell assignments per grid position, compress the surviving row/column -/// ids to sequential indexes, and expand spans into a dense `rows × cols` text -/// grid. Matching runs in docling's coordinate space — the table bbox rounded -/// to integers, everything ×2 (its page scale) — so the post-processor's -/// absolute rounding agrees. -fn docling_match_rows( - cells: &[TableCell], - region: [f32; 4], - table_words: &[crate::tf_match::PdfWord], - words: &[TextCell], -) -> Option>> { - const SCALE: f64 = 2.0; // docling's table-structure page scale - let sl = (region[0] as f64).round_ties_even() * SCALE; - let st = (region[1] as f64).round_ties_even() * SCALE; - let sr = (region[2] as f64).round_ties_even() * SCALE; - let sb = (region[3] as f64).round_ties_even() * SCALE; - let (w2, h2) = (sr - sl, sb - st); - - let tf_cells: Vec = cells - .iter() - .enumerate() - .map(|(i, c)| { - let (cx, cy) = (c.cx as f64, c.cy as f64); - let (w, h) = (c.w as f64, c.h as f64); - crate::tf_match::TfCell { - bbox: [ - sl + (cx - w / 2.0) * w2, - st + (cy - h / 2.0) * h2, - sl + (cx + w / 2.0) * w2, - st + (cy + h / 2.0) * h2, - ], - cell_id: i, - row_id: c.row, - column_id: c.col, - cell_class: c.class, - colspan_val: if c.colspan > 1 { c.colspan } else { 0 }, - rowspan_val: if c.rowspan > 1 { c.rowspan } else { 0 }, - } - }) - .collect(); - - let scaled_words: Vec = table_words - .iter() - .map(|w| crate::tf_match::PdfWord { - id: w.id, - bbox: [ - w.bbox[0] * SCALE, - w.bbox[1] * SCALE, - w.bbox[2] * SCALE, - w.bbox[3] * SCALE, - ], - text: w.text.clone(), - }) - .collect(); - - // Debug: dump the matcher inputs as JSON lines for a side-by-side run - // against docling's Python post-processor (parity harness, not a feature). - if let Ok(dir) = std::env::var("DOCLING_RS_TF_MATCH_DUMP") { - if !dir.is_empty() { - dump_match_inputs(&dir, &tf_cells, &scaled_words); - } - } - - let (cells_wo, final_matches) = - crate::tf_match::match_and_post_process(tf_cells, &scaled_words); - - // `_merge_tf_output`: group per (column, row) in ascending-pdf-id order; - // the first word's table cell fixes the group's offsets and spans. - struct Merged { - start_row: usize, - start_col: usize, - row_span: usize, - col_span: usize, - word_ids: Vec, - } - let mut merged: Vec = Vec::new(); - let mut key_ix: std::collections::HashMap<(usize, usize), usize> = - std::collections::HashMap::new(); - for (&pdf_id, list) in &final_matches { - let tm = list[0].table_cell_id; - let Some(cell) = cells_wo.iter().find(|c| c.cell_id == tm) else { - continue; - }; - match key_ix.entry((cell.column_id, cell.row_id)) { - std::collections::hash_map::Entry::Occupied(e) => { - merged[*e.get()].word_ids.push(pdf_id); - } - std::collections::hash_map::Entry::Vacant(e) => { - e.insert(merged.len()); - merged.push(Merged { - start_row: cell.row_id, - start_col: cell.column_id, - row_span: cell.rowspan_val.max(1), - col_span: cell.colspan_val.max(1), - word_ids: vec![pdf_id], - }); - } - } - } - if merged.is_empty() { - return None; - } - - // `multi_table_predict`'s sort_row_col_indexes: compress the surviving - // row/column ids to gap-free indexes. - let mut start_cols: Vec = merged.iter().map(|m| m.start_col).collect(); - start_cols.sort_unstable(); - start_cols.dedup(); - let mut start_rows: Vec = merged.iter().map(|m| m.start_row).collect(); - start_rows.sort_unstable(); - start_rows.dedup(); - let mut num_rows = 0; - let mut num_cols = 0; - for m in &mut merged { - m.start_col = start_cols.binary_search(&m.start_col).expect("own value"); - m.start_row = start_rows.binary_search(&m.start_row).expect("own value"); - num_cols = num_cols.max(m.start_col + m.col_span); - num_rows = num_rows.max(m.start_row + m.row_span); - } - if num_rows == 0 || num_cols == 0 { - return None; - } - - let mut grid = vec![vec![String::new(); num_cols]; num_rows]; - for m in &merged { - let text = m - .word_ids - .iter() - .map(|&i| words[i].text.trim()) - .collect::>() - .join(" "); - let text = normalize_cell_text(text); - for row in grid.iter_mut().skip(m.start_row).take(m.row_span) { - for cell in row.iter_mut().skip(m.start_col).take(m.col_span) { - *cell = text.clone(); - } - } + // The ort-free tail (word matching + grid assembly) is shared with the + // browser path in tf_core. + crate::tf_core::table_rows(&cells, region, words) } - Some(grid) } /// Note once per process that TableFormer's ONNX graphs weren't found, so tables @@ -884,139 +506,6 @@ fn warn_missing_once(enc: &str, dec: &str, bbx: &str) { /// (2,1,0), so width is the major spatial axis. The page→1024px box-average /// (cv2.INTER_AREA) is the caller's job. fn preprocess(img: &RgbImage) -> Result, String> { - let nn = (SIDE * SIDE) as usize; - let side = SIDE as usize; - let (sw, sh) = (img.width() as i32, img.height() as i32); - let sxr = sw as f32 / SIDE as f32; - let syr = sh as f32 / SIDE as f32; - let mut data = vec![0f32; 3 * nn]; - for h in 0..side { - let fy = (h as f32 + 0.5) * syr - 0.5; - let wy = fy - fy.floor(); - let y0c = (fy.floor() as i32).clamp(0, sh - 1) as u32; - let y1c = (fy.floor() as i32 + 1).clamp(0, sh - 1) as u32; - for w in 0..side { - let fx = (w as f32 + 0.5) * sxr - 0.5; - let wx = fx - fx.floor(); - let x0c = (fx.floor() as i32).clamp(0, sw - 1) as u32; - let x1c = (fx.floor() as i32 + 1).clamp(0, sw - 1) as u32; - let p00 = img.get_pixel(x0c, y0c); - let p01 = img.get_pixel(x1c, y0c); - let p10 = img.get_pixel(x0c, y1c); - let p11 = img.get_pixel(x1c, y1c); - let idx = w * side + h; // (C, W, H): c*n + w*H + h - for c in 0..3 { - let top = p00[c] as f32 * (1.0 - wx) + p01[c] as f32 * wx; - let bot = p10[c] as f32 * (1.0 - wx) + p11[c] as f32 * wx; - let v = top * (1.0 - wy) + bot * wy; - data[c * nn + idx] = (v / 255.0 - MEAN[c]) / STD[c]; - } - } - } - Tensor::from_array(([1usize, 3, side, side], data)) + Tensor::from_array(([1usize, 3, SIDE, SIDE], preprocess_input(img))) .map_err(|e| format!("tableformer: input: {e}")) } - -/// docling's `mergebboxes` (cxcywh): the union box of a horizontal span's first -/// and last cell. -fn mergebboxes(b1: [f32; 4], b2: [f32; 4]) -> [f32; 4] { - let new_w = (b2[0] + b2[2] / 2.0) - (b1[0] - b1[2] / 2.0); - let new_h = (b2[1] + b2[3] / 2.0) - (b1[1] - b1[3] / 2.0); - let new_left = b1[0] - b1[2] / 2.0; - let new_top = (b2[1] - b2[3] / 2.0).min(b1[1] - b1[3] / 2.0); - [new_left + new_w / 2.0, new_top + new_h / 2.0, new_w, new_h] -} - -/// Apply docling's span merges: each merge key combines its box with the partner -/// (`-1` → the last box); partners are dropped. The merged cell keeps the -/// *first* box's class, matching docling's `outputs_class1.append(cls1)`. -fn merge_spans( - boxes: &[[f32; 4]], - classes: &[i64], - merge: &std::collections::HashMap, -) -> (Vec<[f32; 4]>, Vec) { - let skip: std::collections::HashSet = merge - .values() - .filter(|&&v| v >= 0) - .map(|&v| v as usize) - .collect(); - let mut out = Vec::new(); - let mut out_classes = Vec::new(); - for (i, &b) in boxes.iter().enumerate() { - let class = classes.get(i).copied().unwrap_or(2); - if let Some(&j) = merge.get(&i) { - let partner = if j < 0 { boxes.len() - 1 } else { j as usize }; - out.push(mergebboxes(b, boxes[partner.min(boxes.len() - 1)])); - out_classes.push(class); - } else if !skip.contains(&i) { - out.push(b); - out_classes.push(class); - } - } - (out, out_classes) -} - -const CELL_TAGS: [i64; 6] = [FCEL, ECEL, XCEL, CHED, RHED, SROW]; - -/// Lay the OTSL tag stream onto a grid (docling's `_build_table_cells`, OTSL -/// mode): cell tags create cells at (row, col); `lcel`/`ucel`/`xcel` are spans -/// (counted toward the column index but not cells). Colspan/rowspan are read off -/// the grid (consecutive `lcel`/`ucel` to the right/below). `boxes` are indexed -/// by cell order and aligned with the cells. -fn build_table_cells(otsl: &[i64], boxes: &[[f32; 4]], classes: &[i64]) -> Vec { - // 2D grid of tags (rows split on NL) for span lookups. - let mut grid: Vec> = vec![Vec::new()]; - for &t in otsl { - if t == NL { - grid.push(Vec::new()); - } else { - grid.last_mut().unwrap().push(t); - } - } - let mut cells = Vec::new(); - let mut cell_id = 0usize; - for (r, row) in grid.iter().enumerate() { - for (c, &tag) in row.iter().enumerate() { - if !CELL_TAGS.contains(&tag) { - continue; - } - let mut colspan = 1; - while c + colspan < row.len() && matches!(row[c + colspan], LCEL | XCEL) { - colspan += 1; - } - let mut rowspan = 1; - while r + rowspan < grid.len() - && grid[r + rowspan] - .get(c) - .is_some_and(|&t| matches!(t, UCEL | XCEL)) - { - rowspan += 1; - } - let b = boxes.get(cell_id).copied().unwrap_or([0.0; 4]); - // docling defaults a class-less cell to 2 (full). - let class = classes.get(cell_id).copied().unwrap_or(2); - cells.push(TableCell { - row: r, - col: c, - colspan, - rowspan, - tag, - class, - cx: b[0], - cy: b[1], - w: b[2], - h: b[3], - }); - cell_id += 1; - } - } - cells -} - -fn argmax(v: &[f32]) -> usize { - v.iter() - .enumerate() - .max_by(|a, b| a.1.total_cmp(b.1)) - .map(|(i, _)| i) - .unwrap_or(0) -} diff --git a/crates/docling-pdf/src/tf_core.rs b/crates/docling-pdf/src/tf_core.rs new file mode 100644 index 00000000..c5426bd9 --- /dev/null +++ b/crates/docling-pdf/src/tf_core.rs @@ -0,0 +1,672 @@ +//! ONNX-free half of the TableFormer pipeline: the 448 encoder-input +//! preprocessing, the autoregressive loop's structure corrections and bbox +//! bookkeeping, span merging, and the OTSL→grid layout. Everything here is pure +//! Rust (no `ort`), so the browser build (#157 stage 3) runs the *same* +//! conformance-critical logic the native pipeline does and delegates only the +//! three ONNX graphs (encoder / decoder / bbox) to ONNX Runtime Web. Keeping +//! one implementation is what keeps the wasm table structure identical to the +//! native CPU path; drift can only come from the runtime kernels. +//! +//! `tableformer.rs` (native, `ml` feature) owns the `ort` sessions and the +//! owned-value KV-cache fast path; it calls into here for the parts that don't +//! touch the runtime. + +use image::RgbImage; + +/// The encoder's fixed square input side. +pub const SIDE: u32 = 448; +// Verbatim from docling's tm_config.json image_normalization (more digits than +// f32 holds; kept exact for provenance). +#[allow(clippy::excessive_precision)] +pub const MEAN: [f32; 3] = [0.94247851, 0.94254675, 0.94292611]; +#[allow(clippy::excessive_precision)] +pub const STD: [f32; 3] = [0.17910956, 0.17940403, 0.17931663]; +/// Cap on decode steps (docling's generation limit). +pub const MAX_STEPS: usize = 1024; +/// The decoder hidden width, and the bbox decoder's per-cell `tag_h` stride. +pub const EMBED_DIM: usize = 512; + +/// OTSL structure tokens (TableModel04_rs wordmap indices). +pub const START: i64 = 2; +pub const END: i64 = 3; +pub const ECEL: i64 = 4; // empty cell +pub const FCEL: i64 = 5; // full (content) cell +pub const LCEL: i64 = 6; // left-looking: extends the cell to its left (colspan) +pub const UCEL: i64 = 7; // up-looking: extends the cell above (rowspan) +pub const XCEL: i64 = 8; // cross: spans both ways +pub const NL: i64 = 9; // new row +pub const CHED: i64 = 10; // column header +pub const RHED: i64 = 11; // row header +pub const SROW: i64 = 12; // section row + +const CELL_TAGS: [i64; 6] = [FCEL, ECEL, XCEL, CHED, RHED, SROW]; + +/// A predicted table cell: an OTSL grid position (with spans) + its box in the +/// 448 image normalized cxcywh, the OTSL tag, and the bbox decoder's cell +/// class (docling's `cell_class`; 2 = full, ≤1 = predicted empty). +#[derive(Debug, Clone)] +pub struct TableCell { + pub row: usize, + pub col: usize, + pub colspan: usize, + pub rowspan: usize, + pub tag: i64, + pub class: i64, + pub cx: f32, + pub cy: f32, + pub w: f32, + pub h: f32, +} + +/// Resize `img` to `SIDE×SIDE` (bilinear, aligned to docling's half-pixel +/// centers) and normalize, laid out `(C, W, H)` as the exported encoder expects +/// — the raw `[1,3,SIDE,SIDE]` float buffer. The native path wraps this in an +/// `ort` tensor; the browser path hands it to ONNX Runtime Web directly. +pub fn preprocess_input(img: &RgbImage) -> Vec { + let nn = (SIDE * SIDE) as usize; + let side = SIDE as usize; + let (sw, sh) = (img.width() as i32, img.height() as i32); + let sxr = sw as f32 / SIDE as f32; + let syr = sh as f32 / SIDE as f32; + let mut data = vec![0f32; 3 * nn]; + for h in 0..side { + let fy = (h as f32 + 0.5) * syr - 0.5; + let wy = fy - fy.floor(); + let y0c = (fy.floor() as i32).clamp(0, sh - 1) as u32; + let y1c = (fy.floor() as i32 + 1).clamp(0, sh - 1) as u32; + for w in 0..side { + let fx = (w as f32 + 0.5) * sxr - 0.5; + let wx = fx - fx.floor(); + let x0c = (fx.floor() as i32).clamp(0, sw - 1) as u32; + let x1c = (fx.floor() as i32 + 1).clamp(0, sw - 1) as u32; + let p00 = img.get_pixel(x0c, y0c); + let p01 = img.get_pixel(x1c, y0c); + let p10 = img.get_pixel(x0c, y1c); + let p11 = img.get_pixel(x1c, y1c); + let idx = w * side + h; // (C, W, H): c*n + w*H + h + for c in 0..3 { + let top = p00[c] as f32 * (1.0 - wx) + p01[c] as f32 * wx; + let bot = p10[c] as f32 * (1.0 - wx) + p11[c] as f32 * wx; + let v = top * (1.0 - wy) + bot * wy; + data[c * nn + idx] = (v / 255.0 - MEAN[c]) / STD[c]; + } + } + } + data +} + +/// docling's two structure corrections, applied to a raw argmax tag: `xcel` +/// collapses to `lcel` (its `line_num` is never incremented, so this fires on +/// every row), and an `lcel` right after a `ucel` becomes a full cell. +pub fn correct(raw: i64, prev_ucel: bool) -> i64 { + let mut tag = raw; + if tag == XCEL { + tag = LCEL; + } + if prev_ucel && tag == LCEL { + tag = FCEL; + } + tag +} + +/// The autoregressive loop's per-step state, mirroring docling's `predict` +/// bookkeeping (`tag_H_buf` / `bboxes_to_merge`): which decoder hidden states +/// feed the bbox decoder and how horizontal spans merge. Both the native and +/// browser loops step the decoder themselves (sync vs async `ort`) and feed +/// each result through [`step`](Self::step); everything else stays here so the +/// two paths can't drift. +#[derive(Default)] +pub struct BboxBook { + /// The decoder input prefix (`[START]`, then every emitted tag). + pub tags: Vec, + /// The emitted OTSL structure tokens (no `START`/`END`). + pub otsl: Vec, + /// Per-bbox-cell decoder hidden states, flattened `[n, EMBED_DIM]`. + pub hiddens: Vec, + /// Number of hidden states collected (`hiddens.len() / EMBED_DIM`). + pub n: usize, + /// Span merges: `cur_bbox_ind → partner` (`-1` → the last box). + pub merge: std::collections::HashMap, + prev_ucel: bool, + skip: bool, + first_lcel: bool, + bbox_ind: usize, + cur_bbox_ind: usize, +} + +impl BboxBook { + pub fn new() -> Self { + Self { + tags: vec![START], + skip: true, // first tag after is skipped + first_lcel: true, + ..Default::default() + } + } + + /// Feed one raw decoded tag and its hidden state. Returns `false` when the + /// corrected tag is `END` (stop decoding) — the tag is not recorded then. + pub fn step(&mut self, raw: i64, hidden: &[f32]) -> bool { + let tag = correct(raw, self.prev_ucel); + if tag == END { + return false; + } + // docling's tag_H_buf / bboxes_to_merge bookkeeping. + if !self.skip && matches!(tag, FCEL | ECEL | CHED | RHED | SROW | NL | UCEL) { + self.hiddens.extend_from_slice(hidden); + self.n += 1; + if !self.first_lcel { + self.merge.insert(self.cur_bbox_ind, self.bbox_ind as i64); + } + self.bbox_ind += 1; + } + if tag != LCEL { + self.first_lcel = true; + } else if self.first_lcel { + self.hiddens.extend_from_slice(hidden); + self.n += 1; + self.first_lcel = false; + self.cur_bbox_ind = self.bbox_ind; + self.merge.insert(self.cur_bbox_ind, -1); + self.bbox_ind += 1; + } + self.skip = matches!(tag, NL | UCEL | XCEL); + self.prev_ucel = tag == UCEL; + self.otsl.push(tag); + self.tags.push(tag); + true + } +} + +/// docling's `mergebboxes` (cxcywh): the union box of a horizontal span's first +/// and last cell. +fn mergebboxes(b1: [f32; 4], b2: [f32; 4]) -> [f32; 4] { + let new_w = (b2[0] + b2[2] / 2.0) - (b1[0] - b1[2] / 2.0); + let new_h = (b2[1] + b2[3] / 2.0) - (b1[1] - b1[3] / 2.0); + let new_left = b1[0] - b1[2] / 2.0; + let new_top = (b2[1] - b2[3] / 2.0).min(b1[1] - b1[3] / 2.0); + [new_left + new_w / 2.0, new_top + new_h / 2.0, new_w, new_h] +} + +/// Apply docling's span merges: each merge key combines its box with the partner +/// (`-1` → the last box); partners are dropped. The merged cell keeps the +/// *first* box's class, matching docling's `outputs_class1.append(cls1)`. +pub fn merge_spans( + boxes: &[[f32; 4]], + classes: &[i64], + merge: &std::collections::HashMap, +) -> (Vec<[f32; 4]>, Vec) { + let skip: std::collections::HashSet = merge + .values() + .filter(|&&v| v >= 0) + .map(|&v| v as usize) + .collect(); + let mut out = Vec::new(); + let mut out_classes = Vec::new(); + for (i, &b) in boxes.iter().enumerate() { + let class = classes.get(i).copied().unwrap_or(2); + if let Some(&j) = merge.get(&i) { + let partner = if j < 0 { boxes.len() - 1 } else { j as usize }; + out.push(mergebboxes(b, boxes[partner.min(boxes.len() - 1)])); + out_classes.push(class); + } else if !skip.contains(&i) { + out.push(b); + out_classes.push(class); + } + } + (out, out_classes) +} + +/// Lay the OTSL tag stream onto a grid (docling's `_build_table_cells`, OTSL +/// mode): cell tags create cells at (row, col); `lcel`/`ucel`/`xcel` are spans +/// (counted toward the column index but not cells). Colspan/rowspan are read off +/// the grid (consecutive `lcel`/`ucel` to the right/below). `boxes` are indexed +/// by cell order and aligned with the cells. +pub fn build_table_cells(otsl: &[i64], boxes: &[[f32; 4]], classes: &[i64]) -> Vec { + // 2D grid of tags (rows split on NL) for span lookups. + let mut grid: Vec> = vec![Vec::new()]; + for &t in otsl { + if t == NL { + grid.push(Vec::new()); + } else { + grid.last_mut().unwrap().push(t); + } + } + let mut cells = Vec::new(); + let mut cell_id = 0usize; + for (r, row) in grid.iter().enumerate() { + for (c, &tag) in row.iter().enumerate() { + if !CELL_TAGS.contains(&tag) { + continue; + } + let mut colspan = 1; + while c + colspan < row.len() && matches!(row[c + colspan], LCEL | XCEL) { + colspan += 1; + } + let mut rowspan = 1; + while r + rowspan < grid.len() + && grid[r + rowspan] + .get(c) + .is_some_and(|&t| matches!(t, UCEL | XCEL)) + { + rowspan += 1; + } + let b = boxes.get(cell_id).copied().unwrap_or([0.0; 4]); + // docling defaults a class-less cell to 2 (full). + let class = classes.get(cell_id).copied().unwrap_or(2); + cells.push(TableCell { + row: r, + col: c, + colspan, + rowspan, + tag, + class, + cx: b[0], + cy: b[1], + w: b[2], + h: b[3], + }); + cell_id += 1; + } + } + cells +} + +/// Index of the maximum. Uses Rust's `max_by` (ties resolve to the *last* +/// index; the decoder/bbox float logits don't produce exact ties in practice). +/// Kept verbatim from the native path so the two stay bit-identical. +pub fn argmax(v: &[f32]) -> usize { + v.iter() + .enumerate() + .max_by(|a, b| a.1.total_cmp(b.1)) + .map(|(i, _)| i) + .unwrap_or(0) +} + +use crate::pdfium_backend::TextCell; +use crate::tf_match::{PdfWord, TfCell}; + +/// The ONNX-free tail of TableFormer row prediction: match the page's word +/// cells into the predicted structure `cells` and expand spans into a dense +/// `rows × cols` text grid. `region` is `(l, t, r, b)` in page points; `cells` +/// are in the 448 image (normalized cxcywh). Shared by the native pipeline +/// (after `predict_table_structure`) and the browser path (after the ort-web +/// decode loop) — identical from here on. `None` when nothing matched. +pub fn table_rows( + cells: &[TableCell], + region: [f32; 4], + words: &[TextCell], +) -> Option>> { + // Words that belong to the table: non-empty text, ≥80 % of the word's area + // inside the table region (docling's `get_cells_in_bbox` ios test). Ids stay + // the page-level word indices so text joins in stream order. + let table_words: Vec = words + .iter() + .enumerate() + .filter(|(_, w)| !w.text.trim().is_empty()) + .filter_map(|(wi, w)| { + let (l, t, r, b) = (w.l as f64, w.t as f64, w.r as f64, w.b as f64); + let area = (r - l) * (b - t); + let iw = (r.min(region[2] as f64) - l.max(region[0] as f64)).max(0.0); + let ih = (b.min(region[3] as f64) - t.max(region[1] as f64)).max(0.0); + if area > 0.0 && iw * ih / area > 0.8 { + Some(PdfWord { + id: wi, + bbox: [l, t, r, b], + text: w.text.trim().to_string(), + }) + } else { + None + } + }) + .collect(); + + if !table_words.is_empty() && !simple_match() { + return docling_match_rows(cells, region, &table_words, words); + } + + let (rw, rh) = (region[2] - region[0], region[3] - region[1]); + + // Cell boxes in page points (top-left), aligned with `cells`. + let boxes: Vec<[f32; 4]> = cells + .iter() + .map(|c| { + [ + region[0] + (c.cx - c.w / 2.0) * rw, + region[1] + (c.cy - c.h / 2.0) * rh, + region[0] + (c.cx + c.w / 2.0) * rw, + region[1] + (c.cy + c.h / 2.0) * rh, + ] + }) + .collect(); + + // Assign each word to the cell it overlaps most (intersection / word area). + let mut cell_words: Vec> = vec![Vec::new(); cells.len()]; + for (wi, w) in words.iter().enumerate() { + let wa = ((w.r - w.l) * (w.b - w.t)).max(1.0); + let mut best: Option<(f32, usize)> = None; + for (ci, b) in boxes.iter().enumerate() { + let ix = (w.r.min(b[2]) - w.l.max(b[0])).max(0.0); + let iy = (w.b.min(b[3]) - w.t.max(b[1])).max(0.0); + let io = ix * iy / wa; + if io > 0.0 && best.is_none_or(|(bo, _)| io > bo) { + best = Some((io, ci)); + } + } + if let Some((_, ci)) = best { + cell_words[ci].push(wi); + } + } + + let num_rows = cells.iter().map(|c| c.row + c.rowspan).max().unwrap_or(0); + let num_cols = cells.iter().map(|c| c.col + c.colspan).max().unwrap_or(0); + if num_rows == 0 || num_cols == 0 { + return None; + } + let mut grid = vec![vec![String::new(); num_cols]; num_rows]; + for (ci, c) in cells.iter().enumerate() { + // Keep words in text-stream order (their word index), matching docling's + // cell text assembly — geometric re-sorting scrambles wrapped cells. + let wis = std::mem::take(&mut cell_words[ci]); + let text = wis + .iter() + .map(|&i| words[i].text.trim()) + .collect::>() + .join(" "); + let text = normalize_cell_text(text); + // Spanned cells repeat their text across the covered grid positions. + for row in grid.iter_mut().skip(c.row).take(c.rowspan) { + for cell in row.iter_mut().skip(c.col).take(c.colspan) { + *cell = text.clone(); + } + } + } + Some(grid) +} + +/// `DOCLING_RS_TF_SIMPLE_MATCH=1` reverts to the pre-#60 best-overlap word +/// assignment (A/B escape hatch for the docling matching post-processor). +fn simple_match() -> bool { + std::env::var("DOCLING_RS_TF_SIMPLE_MATCH").is_ok_and(|v| !v.is_empty() && v != "0") +} + +/// docling glues `@` to whatever follows it (`mAP @0.5`, an email): the PDF's +/// word cells split `@` from the next token, and joining with a space would +/// widen the cell and — via the column pad — shift every row. The groundtruth +/// never contains "@ ", so this is always the right normalization. +fn normalize_cell_text(text: String) -> String { + text.replace("@ ", "@") +} + +/// docling's matched-cell grid assembly (`tf_predictor.predict` with +/// `do_cell_matching=True`): run the ported matching post-processor, group the +/// word→cell assignments per grid position, compress the surviving row/column +/// ids to sequential indexes, and expand spans into a dense `rows × cols` text +/// grid. Matching runs in docling's coordinate space — the table bbox rounded +/// to integers, everything ×2 (its page scale) — so the post-processor's +/// absolute rounding agrees. +fn docling_match_rows( + cells: &[TableCell], + region: [f32; 4], + table_words: &[PdfWord], + words: &[TextCell], +) -> Option>> { + const SCALE: f64 = 2.0; // docling's table-structure page scale + let sl = (region[0] as f64).round_ties_even() * SCALE; + let st = (region[1] as f64).round_ties_even() * SCALE; + let sr = (region[2] as f64).round_ties_even() * SCALE; + let sb = (region[3] as f64).round_ties_even() * SCALE; + let (w2, h2) = (sr - sl, sb - st); + + let tf_cells: Vec = cells + .iter() + .enumerate() + .map(|(i, c)| { + let (cx, cy) = (c.cx as f64, c.cy as f64); + let (w, h) = (c.w as f64, c.h as f64); + TfCell { + bbox: [ + sl + (cx - w / 2.0) * w2, + st + (cy - h / 2.0) * h2, + sl + (cx + w / 2.0) * w2, + st + (cy + h / 2.0) * h2, + ], + cell_id: i, + row_id: c.row, + column_id: c.col, + cell_class: c.class, + colspan_val: if c.colspan > 1 { c.colspan } else { 0 }, + rowspan_val: if c.rowspan > 1 { c.rowspan } else { 0 }, + } + }) + .collect(); + + let scaled_words: Vec = table_words + .iter() + .map(|w| PdfWord { + id: w.id, + bbox: [ + w.bbox[0] * SCALE, + w.bbox[1] * SCALE, + w.bbox[2] * SCALE, + w.bbox[3] * SCALE, + ], + text: w.text.clone(), + }) + .collect(); + + // Debug (native only): dump the matcher inputs as JSON lines for a + // side-by-side run against docling's Python post-processor. + #[cfg(feature = "ml")] + if let Ok(dir) = std::env::var("DOCLING_RS_TF_MATCH_DUMP") { + if !dir.is_empty() { + dump_match_inputs(&dir, &tf_cells, &scaled_words); + } + } + + let (cells_wo, final_matches) = + crate::tf_match::match_and_post_process(tf_cells, &scaled_words); + + // `_merge_tf_output`: group per (column, row) in ascending-pdf-id order; the + // first word's table cell fixes the group's offsets and spans. + struct Merged { + start_row: usize, + start_col: usize, + row_span: usize, + col_span: usize, + word_ids: Vec, + } + let mut merged: Vec = Vec::new(); + let mut key_ix: std::collections::HashMap<(usize, usize), usize> = + std::collections::HashMap::new(); + for (&pdf_id, list) in &final_matches { + let tm = list[0].table_cell_id; + let Some(cell) = cells_wo.iter().find(|c| c.cell_id == tm) else { + continue; + }; + match key_ix.entry((cell.column_id, cell.row_id)) { + std::collections::hash_map::Entry::Occupied(e) => { + merged[*e.get()].word_ids.push(pdf_id); + } + std::collections::hash_map::Entry::Vacant(e) => { + e.insert(merged.len()); + merged.push(Merged { + start_row: cell.row_id, + start_col: cell.column_id, + row_span: cell.rowspan_val.max(1), + col_span: cell.colspan_val.max(1), + word_ids: vec![pdf_id], + }); + } + } + } + if merged.is_empty() { + return None; + } + + // `multi_table_predict`'s sort_row_col_indexes: compress the surviving + // row/column ids to gap-free indexes. + let mut start_cols: Vec = merged.iter().map(|m| m.start_col).collect(); + start_cols.sort_unstable(); + start_cols.dedup(); + let mut start_rows: Vec = merged.iter().map(|m| m.start_row).collect(); + start_rows.sort_unstable(); + start_rows.dedup(); + let mut num_rows = 0; + let mut num_cols = 0; + for m in &mut merged { + m.start_col = start_cols.binary_search(&m.start_col).expect("own value"); + m.start_row = start_rows.binary_search(&m.start_row).expect("own value"); + num_cols = num_cols.max(m.start_col + m.col_span); + num_rows = num_rows.max(m.start_row + m.row_span); + } + if num_rows == 0 || num_cols == 0 { + return None; + } + + let mut grid = vec![vec![String::new(); num_cols]; num_rows]; + for m in &merged { + let text = m + .word_ids + .iter() + .map(|&i| words[i].text.trim()) + .collect::>() + .join(" "); + let text = normalize_cell_text(text); + for row in grid.iter_mut().skip(m.start_row).take(m.row_span) { + for cell in row.iter_mut().skip(m.start_col).take(m.col_span) { + *cell = text.clone(); + } + } + } + Some(grid) +} + +/// Append one JSON line per table into `/tf_match_dump.jsonl` with the +/// exact matcher inputs (hand-rolled JSON to avoid a serde dependency). +#[cfg(feature = "ml")] +fn dump_match_inputs(dir: &str, tf_cells: &[TfCell], words: &[PdfWord]) { + use std::io::Write; + let cells: Vec = tf_cells + .iter() + .map(|c| { + format!( + r#"{{"bbox":[{},{},{},{}],"cell_id":{},"row_id":{},"column_id":{},"cell_class":{},"colspan_val":{},"rowspan_val":{}}}"#, + c.bbox[0], c.bbox[1], c.bbox[2], c.bbox[3], + c.cell_id, c.row_id, c.column_id, c.cell_class, + c.colspan_val, c.rowspan_val + ) + }) + .collect(); + let ws: Vec = words + .iter() + .map(|w| { + format!( + r#"{{"id":{},"bbox":[{},{},{},{}],"text":{}}}"#, + w.id, + w.bbox[0], + w.bbox[1], + w.bbox[2], + w.bbox[3], + serde_json_escape(&w.text) + ) + }) + .collect(); + let line = format!( + r#"{{"table_cells":[{}],"pdf_cells":[{}]}}"#, + cells.join(","), + ws.join(",") + ); + if let Ok(mut f) = std::fs::OpenOptions::new() + .create(true) + .append(true) + .open(format!("{dir}/tf_match_dump.jsonl")) + { + let _ = writeln!(f, "{line}"); + } +} + +/// Minimal JSON string escaping for the parity dump. +#[cfg(feature = "ml")] +fn serde_json_escape(s: &str) -> String { + let mut out = String::with_capacity(s.len() + 2); + out.push('"'); + for ch in s.chars() { + match ch { + '"' => out.push_str("\\\""), + '\\' => out.push_str("\\\\"), + '\n' => out.push_str("\\n"), + '\r' => out.push_str("\\r"), + '\t' => out.push_str("\\t"), + c if (c as u32) < 0x20 => out.push_str(&format!("\\u{:04x}", c as u32)), + c => out.push(c), + } + } + out.push('"'); + out +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn corrections() { + assert_eq!(correct(XCEL, false), LCEL); // xcel → lcel + assert_eq!(correct(LCEL, true), FCEL); // lcel after ucel → fcel + assert_eq!(correct(XCEL, true), FCEL); // xcel → lcel → fcel + assert_eq!(correct(FCEL, false), FCEL); + assert_eq!(correct(LCEL, false), LCEL); + } + + #[test] + fn argmax_behaviour() { + assert_eq!(argmax(&[0.1, 0.9, 0.3]), 1); + assert_eq!(argmax(&[0.5, 0.5]), 1); // Rust max_by ties to the last index + assert_eq!(argmax(&[]), 0); + } + + #[test] + fn book_skips_first_and_collects_hiddens() { + // then a 2x1 row: FCEL FCEL NL, END. The first FCEL after start + // is NOT skipped (skip only guards the tag right after start-consumed + // rows); verify hidden collection count and merge stays empty. + let mut b = BboxBook::new(); + let h = [1.0f32; EMBED_DIM]; + assert!(b.step(FCEL, &h)); // skip=true initially → not collected + assert!(b.step(FCEL, &h)); + assert!(b.step(NL, &h)); + assert!(!b.step(END, &h)); // stop + assert_eq!(b.otsl, vec![FCEL, FCEL, NL]); + // first FCEL skipped (skip=true), second FCEL + NL collected → n=2 + assert_eq!(b.n, 2); + assert_eq!(b.hiddens.len(), 2 * EMBED_DIM); + assert!(b.merge.is_empty()); + } + + #[test] + fn book_merges_horizontal_span() { + // FCEL LCEL: the LCEL is the first-lcel of a horizontal span → records a + // merge partner (-1 placeholder) for the span's leading cell. + let mut b = BboxBook::new(); + let h = [0.0f32; EMBED_DIM]; + b.step(FCEL, &h); // skipped (skip=true) + b.step(FCEL, &h); // collected, bbox_ind 0→1 + b.step(LCEL, &h); // first-lcel: cur=1, merge{1:-1}, bbox_ind 1→2 + assert_eq!(b.merge.get(&1), Some(&-1)); + } + + #[test] + fn build_cells_spans() { + // Row 0: FCEL LCEL (a 1x2 colspan) + // Row 1: FCEL ECEL + let otsl = vec![FCEL, LCEL, NL, FCEL, ECEL]; + let boxes = vec![[0.0; 4]; 3]; + let classes = vec![2, 2, 2]; + let cells = build_table_cells(&otsl, &boxes, &classes); + assert_eq!(cells.len(), 3); + assert_eq!((cells[0].colspan, cells[0].rowspan), (2, 1)); + assert_eq!((cells[0].row, cells[0].col), (0, 0)); + assert_eq!((cells[1].row, cells[1].col), (1, 0)); + } +} diff --git a/crates/docling-wasm/Cargo.toml b/crates/docling-wasm/Cargo.toml index 762b90cb..1183f0ae 100644 --- a/crates/docling-wasm/Cargo.toml +++ b/crates/docling-wasm/Cargo.toml @@ -16,10 +16,30 @@ publish = false # wasm-bindgen needs a cdylib; rlib keeps `cargo test` runnable on the host. crate-type = ["cdylib", "rlib"] +[features] +# Browser OCR (#157 stage 1): the ONNX-free half of the PP-OCRv3 recognition +# pipeline compiled in, with inference delegated to ONNX Runtime Web on the +# JS side. Pure Rust — the ort/pdfium native stack stays out. +default = ["ocr"] +ocr = [ + "dep:docling-core", + "dep:docling-pdf", + "dep:image", + "dep:js-sys", + "dep:wasm-bindgen-futures", +] + [dependencies] # The pure-Rust declarative slice only: no pdfium/ONNX (pdf, asr), no HTTP # client (fetch-images) — exactly the set that compiles for wasm32 (#79). docling = { path = "../docling", version = "0.49.0", default-features = false, features = ["pdf-text"] } +docling-core = { path = "../docling-core", version = "0.49.0", optional = true } +# ocr_prep only: line prep + CTC decode shared with the native pipeline. +docling-pdf = { path = "../docling-pdf", version = "0.49.0", default-features = false, features = ["ocr-prep"], optional = true } +# Input decoding for the OCR path; codecs mirror the browser's usual suspects. +image = { version = "0.25", default-features = false, features = ["png", "jpeg", "bmp", "tiff", "webp"], optional = true } +js-sys = { version = "0.3", optional = true } +wasm-bindgen-futures = { version = "0.4", optional = true } serde_json = "1" wasm-bindgen = "0.2" # Panic messages land in the browser console instead of "unreachable executed". diff --git a/crates/docling-wasm/LICENSE b/crates/docling-wasm/LICENSE new file mode 100644 index 00000000..8a2bc3ea --- /dev/null +++ b/crates/docling-wasm/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Artem Kustikov + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/crates/docling-wasm/README.md b/crates/docling-wasm/README.md index b23622fb..6a5011fe 100644 --- a/crates/docling-wasm/README.md +++ b/crates/docling-wasm/README.md @@ -62,6 +62,122 @@ Verified end-to-end in headless Chromium: Markdown/DOCX→md, DOCX→JSON, a corpus PDF→md through the text-layer path, and the scanned-PDF error path all exercised through the real wasm module. +## In-browser ML pipeline — OCR, layout, TableFormer (#157) + +With the default `ocr` feature the module runs the **full scanned-document ML +pipeline** client-side. Every ONNX-free step — image preprocessing, line/word +segmentation, CTC decoding, RT-DETR box decoding, the TableFormer +autoregressive loop and its bbox bookkeeping, span merging, OTSL→grid layout, +cell matching, region refinement and reading-order assembly — runs in Rust, +**the exact same functions the native pipeline uses** +(`docling_pdf::{ocr_prep, layout, scanned, tf_core, tf_match, resample}`). Only +the four ONNX graphs are delegated to +[ONNX Runtime Web](https://onnxruntime.ai/docs/tutorials/web/) via small JS +wrappers around `ort.InferenceSession`. One implementation is what keeps the +wasm output matching native — drift can only come from the runtime kernels. + +| Stage | What | Models (from the [`models-v1` release](https://github.com/docling-project/docling.rs/releases/tag/models-v1)) | +|---|---|---| +| **1** — `ocr_image` | OCR a single scanned **image** | PP-OCRv3 recognition (~10 MB) | +| **2** — `ScannedConverter` / `convert_scanned_image` | Scanned **PDF/image** → Markdown: layout + OCR + reading order, tables via geometric reconstruction | + RT-DETR layout (`layout_heron_int8.onnx`, ~68 MB) | +| **3** — `ScannedConverter.addPageTf` | Real **TableFormer** table structure instead of geometric | + `tableformer/{encoder,decoder_kv,bbox}.onnx` (~380 MB) | + +Stage 1 recognition wrapper: + +```js +const rec = { + run: async (n, h, w, data) => { + const out = await session.run({ [session.inputNames[0]]: + new ort.Tensor("float32", data, [n, 3, h, w]) }); + const t = out[session.outputNames[0]]; + return { data: t.data, dims: Array.from(t.dims) }; + }, +}; +const markdown = await ocr_image(imageBytes, dictText, rec, "md"); +``` + +Stage 2 rasterizes PDF pages with pdf.js at `{scale: 2}` (2 px/point — the +native `RENDER_SCALE`) and feeds bitmaps page by page: + +```js +const conv = new ScannedConverter(dictText); +for (let p = 1; p <= pdf.numPages; p++) { + const viewport = page.getViewport({ scale: 2 }); + // … render to canvas, then: + await conv.add_page(rgba, canvas.width, canvas.height, 2.0, layout, rec); +} +const markdown = conv.finish(file.name, "md"); +``` + +Stage 3 is the same but `conv.addPageTf(rgba, w, h, 2.0, layout, rec, tf)`, +where `tf` is a stateful session over the three TableFormer graphs (the heavy +cross-attention K/V and the growing decoder KV-cache stay on the JS side, so +each decode step marshals only the last tag and gets back logits+hidden — see +[`www/pipeline.js`](./www/pipeline.js)'s `JsTfSession`). The demos wire all of +this up: [`www/ocr.html`](./www/ocr.html) (stage 1), +[`www/scan.html`](./www/scan.html) (stages 2–3, pdf.js + a TableFormer toggle + +a device model-file picker). + +## Run it on your phone + +The demo is a static page — the models are the only heavy part, and they load +same-origin, from a CORS host, **or straight from files on your device**. + +1. **Fetch the models** (once), on a desktop or the phone: + ```bash + curl -fsSL https://raw.githubusercontent.com/docling-project/docling.rs/master/scripts/install/download_dependencies.sh | sh + ``` + This writes `models/` (layout, OCR, TableFormer, …). + +2. **Serve the page** from any static host that sends the right MIME types — + GitHub Pages, `raw.githack.com`, or a local `python3 -m http.server -d + crates/docling-wasm/www 8901`. (Plain `raw.githubusercontent.com` and most + GitHub CDNs serve HTML as `text/plain`, so the page won't render — use + Pages or githack.) Build the wasm `pkg/` first (see **Build** above). + +3. **Provide the models to the page**, either: + - **from your device** — tap the *"Models from device"* picker in + `scan.html` and select `layout_heron_int8.onnx` and, for TableFormer, + `encoder.onnx`, `decoder_kv.onnx` (+`.data`), `bbox.onnx` (+`.data`). Each + file is read to an `ArrayBuffer` on the client and used directly — no + download, and a single allocation per file (a 225 MB encoder over + `fetch` can OOM a mobile tab; a `File` read does not); or + - **from a CORS mirror** — a Hugging Face model repo serves + `Access-Control-Allow-Origin` (GitHub Release assets do not); point + `MODEL_BASE` in `pipeline.js` at yours. The recognition model already + streams from Hugging Face. + +Resolution order for every model is **device file → local `./models/` → +`MODEL_BASE` (Hugging Face)**. Tables need no upload beyond the model files; +the image never leaves the page. + +## Performance & threading + +- **Threads.** ONNX Runtime Web runs single-threaded unless the page is + [cross-origin isolated](https://web.dev/coop-coep/) (SharedArrayBuffer). A + static host can't set the COOP/COEP headers, so [`www/coi.js`](./www/coi.js) + is a service worker that adds them (`COEP: credentialless`, which keeps the + CDN/HF fetches working); `pipeline.js` then sets + `ort.env.wasm.numThreads` to the core count. The demo's ready line reports + the thread count — if it says `1 thread`, isolation didn't take (some mobile + browsers), and everything runs serially. +- **int8.** The layout model is the conv-only static-INT8 export (~68 MB, + ~2.4× faster than fp32 at unchanged conformance). The TableFormer **encoder + is fp32 (~225 MB) and runs once per table region** — it dominates wall time + on mobile (a multi-table page can take minutes). An int8 encoder is *not* the + easy win it looks like: it's a ResNet backbone (~20 Conv) feeding a 6-layer + transformer, and the transformer Gemms — not the convs — are the cost. + Measured on a native x86 build over 15 real table crops: conv-only static + QDQ keeps fidelity (enc_out/cross cosine ≥ 0.995) but is *not* faster than + ORT's fp32 conv kernels, and quantizing the attention MatMuls collapses + cross-attention fidelity to ~0.85 (garbled structure) for no speed gain. So + the encoder stays fp32; the real mobile levers are running heavy tables on a + desktop and batching multi-table pages (below), not weight quantization. +- **Where to run heavy tables.** TableFormer's 380 MB of models and per-region + encode make a desktop (more cores, more memory) far faster than a phone; the + geometric table path (stage 2, no TableFormer) already captures all cell text + and is the light option for mobile. + ## Host-side tests `cargo test -p docling-wasm` runs the conversion body natively (the diff --git a/crates/docling-wasm/src/lib.rs b/crates/docling-wasm/src/lib.rs index 99aec907..c80f37fc 100644 --- a/crates/docling-wasm/src/lib.rs +++ b/crates/docling-wasm/src/lib.rs @@ -19,6 +19,17 @@ use docling::{DocumentConverter, InputFormat, SourceDocument}; use wasm_bindgen::prelude::*; +#[cfg(feature = "ocr")] +mod ocr; +#[cfg(feature = "ocr")] +mod scanned; +#[cfg(feature = "ocr")] +mod tableformer; +#[cfg(feature = "ocr")] +pub use ocr::ocr_image; +#[cfg(feature = "ocr")] +pub use scanned::{convert_scanned_image, ScannedConverter}; + #[wasm_bindgen(start)] fn start() { // Panics surface as readable messages in the browser console instead of diff --git a/crates/docling-wasm/src/ocr.rs b/crates/docling-wasm/src/ocr.rs new file mode 100644 index 00000000..3b0d030e --- /dev/null +++ b/crates/docling-wasm/src/ocr.rs @@ -0,0 +1,142 @@ +//! Browser OCR — stage 1 of #157 (`#79 phase 2: ML pipeline in the browser`). +//! +//! Scanned **images** convert fully client-side: all pre/post-processing +//! (line segmentation, crop preparation, width-batching, CTC decoding) runs +//! in Rust via `docling_pdf::ocr_prep` — the *same* code the native pipeline +//! uses — and only the PP-OCRv3 recognition inference is delegated to +//! [ONNX Runtime Web](https://onnxruntime.ai/docs/tutorials/web/) on the JS +//! side through the [`RecSession`] interop interface. One shared +//! implementation is what keeps the wasm output comparable to the native CPU +//! path; drift can then only come from the runtime's kernels. +//! +//! Without a layout model (that's stage 2), the whole image is treated as a +//! single text region and split into lines by the ink-projection profile — +//! fine for typical single-column scans, degraded on complex layouts. +//! +//! ```js +//! import * as ort from "onnxruntime-web"; +//! import init, { ocr_image } from "./pkg/docling_wasm.js"; +//! await init(); +//! const session = await ort.InferenceSession.create("ocr_rec_en.onnx"); +//! const dict = await (await fetch("en_dict.txt")).text(); +//! const md = await ocr_image(imageBytes, dict, { +//! run: async (n, h, w, data) => { +//! const out = (await session.run({ x: new ort.Tensor("float32", data, [n, 3, h, w]) })).softmax_2.data +//! ?? Object.values(await session.run({ x: new ort.Tensor("float32", data, [n, 3, h, w]) }))[0]; +//! // return the output tensor: { data: Float32Array, dims: [n, t, c] } +//! }, +//! }); +//! ``` +//! (see `www/ocr.html` for the complete wiring, including output-name +//! discovery and model/dict caching.) + +use docling_core::{DoclingDocument, Node}; +use docling_pdf::ocr_prep::{ + batch_input, decode_row, dict_chars, normalize_polarity, prep_page_lines, width_batches, + REC_HEIGHT, +}; +use wasm_bindgen::prelude::*; + +#[wasm_bindgen] +extern "C" { + /// The JS-side recognition session: a wrapper object around an + /// `ort.InferenceSession` exposing one method, + /// `run(count, height, width, data)`, which feeds the `(N,3,H,W)` CHW + /// float buffer to the model and resolves to + /// `{ data: Float32Array, dims: [n, t, c] }` — the recognition + /// probabilities tensor. + pub type RecSession; + + #[wasm_bindgen(method, catch)] + pub async fn run( + this: &RecSession, + count: u32, + height: u32, + width: u32, + data: js_sys::Float32Array, + ) -> Result; +} + +/// OCR a scanned image entirely in the browser: `bytes` is the image file +/// (PNG/JPEG/…), `dict` the recognition dictionary text (`en_dict.txt` for +/// the default English model), `session` the JS inference wrapper. Returns +/// Markdown (default) or docling JSON per `to`, one paragraph per recognized +/// line. +#[wasm_bindgen] +pub async fn ocr_image( + bytes: &[u8], + dict: &str, + session: &RecSession, + to: Option, +) -> Result { + // Screenshots and dark-mode captures arrive light-on-dark; the whole + // pipeline assumes scan polarity, so normalize first. + let img = normalize_polarity( + image::load_from_memory(bytes) + .map_err(|e| JsError::new(&format!("decode image: {e}")))? + .to_rgb8(), + ); + let lines = prep_page_lines(&img); + let chars = dict_chars(dict); + let mut texts = vec![String::new(); lines.len()]; + for (w, chunk) in width_batches(&lines) { + let input = batch_input(w, &chunk, &lines); + let out = session + .run( + chunk.len() as u32, + REC_HEIGHT, + w as u32, + js_sys::Float32Array::from(input.as_slice()), + ) + .await + .map_err(|e| JsError::new(&format!("session.run: {e:?}")))?; + let (probs, t_len, nc) = tensor_parts(&out)?; + if probs.len() < chunk.len() * t_len * nc { + return Err(JsError::new("session.run returned a short tensor")); + } + for (i, &ix) in chunk.iter().enumerate() { + texts[ix] = decode_row(&chars, &probs[i * t_len * nc..(i + 1) * t_len * nc], nc); + } + } + let mut doc = DoclingDocument::new("image"); + doc.nodes.extend( + texts + .into_iter() + .map(|t| t.trim().to_string()) + .filter(|t| !t.is_empty()) + .map(|text| Node::Paragraph { text }), + ); + match to.as_deref().unwrap_or("md") { + "md" | "markdown" => Ok(doc.export_to_markdown()), + "json" => Ok(doc.export_to_json()), + other => Err(JsError::new(&format!( + "unknown output format {other:?} (expected \"md\" or \"json\")" + ))), + } +} + +/// Pull `{ data: Float32Array, dims: [..] }` out of a JS tensor object; +/// returns `(data, dims[len-2], dims[len-1])` — the trailing two dims are +/// what every consumer indexes by. +pub(crate) fn tensor_parts(out: &JsValue) -> Result<(Vec, usize, usize), JsError> { + let get = |k: &str| { + js_sys::Reflect::get(out, &JsValue::from_str(k)) + .map_err(|_| JsError::new(&format!("session.run result has no `{k}`"))) + }; + let data: js_sys::Float32Array = get("data")? + .dyn_into() + .map_err(|_| JsError::new("`data` is not a Float32Array"))?; + let dims: js_sys::Array = get("dims")? + .dyn_into() + .map_err(|_| JsError::new("`dims` is not an array"))?; + let len = dims.length(); + if len < 2 { + return Err(JsError::new("`dims` must have at least 2 entries")); + } + let t_len = dims.get(len - 2).as_f64().unwrap_or(0.0) as usize; + let nc = dims.get(len - 1).as_f64().unwrap_or(0.0) as usize; + if t_len == 0 || nc == 0 { + return Err(JsError::new("trailing `dims` entries must be > 0")); + } + Ok((data.to_vec(), t_len, nc)) +} diff --git a/crates/docling-wasm/src/scanned.rs b/crates/docling-wasm/src/scanned.rs new file mode 100644 index 00000000..6e322737 --- /dev/null +++ b/crates/docling-wasm/src/scanned.rs @@ -0,0 +1,307 @@ +//! Browser scanned-document pipeline — stage 2 of #157 (the *lite profile*): +//! RT-DETR layout detection + PP-OCRv3 recognition, both via ONNX Runtime Web +//! on the JS side; region refinement, region-cropped OCR, geometric table +//! reconstruction and reading-order assembly in Rust — the same code as the +//! native pipeline (`docling_pdf::{layout, ocr_prep, scanned}`). TableFormer +//! and the enrichment models are stage 3; table regions fall back to the +//! geometric reconstruction the native `--no-table-former` flag uses. +//! +//! Pages arrive as raw RGBA bitmaps: for a scanned PDF the host page renders +//! them with pdf.js (`page.getViewport({scale: 2})` — 2 px per PDF point, +//! matching the native pipeline's `RENDER_SCALE`); a standalone image is its +//! own page at scale 1, exactly like the native image path. +//! +//! ```js +//! const conv = new ScannedConverter(dictText); +//! for (const bitmap of pages) { +//! await conv.add_page(bitmap.data, bitmap.width, bitmap.height, 2.0, layout, rec); +//! } +//! const markdown = conv.finish("scan.pdf", "md"); +//! ``` +//! (`www/scan.html` is the complete wiring.) + +use docling_pdf::layout::{decode_layout, layout_input, SIDE}; +use docling_pdf::ocr_prep::{ + batch_input, decode_row, dict_chars, normalize_polarity, prep_region_lines, prep_table_words, + width_batches, PrepLine, REC_HEIGHT, +}; +use docling_pdf::pdfium_backend::{PdfPage, TextCell}; +use docling_pdf::scanned::{assemble_page_with_tables, finish_document, refine_regions}; +use image::RgbImage; +use wasm_bindgen::prelude::*; + +use crate::ocr::{tensor_parts, RecSession}; +use crate::tableformer::TfSession; + +#[wasm_bindgen] +extern "C" { + // On-page diagnostic sink (defined by the host — worker.js forwards it to + // the main thread, which shows it; console isn't reachable on a phone). + #[wasm_bindgen(js_name = __docling_diag)] + fn diag(s: &str); +} + +#[wasm_bindgen] +extern "C" { + /// The JS-side layout session: a wrapper around an `ort.InferenceSession` + /// over the RT-DETR layout model exposing `run(data)` — feed the + /// `(1, 3, 640, 640)` CHW float buffer, resolve to + /// `{ logits: {data, dims: [1, q, c]}, boxes: {data, dims: [1, q, 4]} }`. + pub type LayoutSession; + + #[wasm_bindgen(method, catch)] + pub async fn run(this: &LayoutSession, data: js_sys::Float32Array) -> Result; +} + +/// Multi-page scanned-document converter (lite profile). Feed pages in +/// order, then [`finish`](Self::finish) — cross-page paragraph continuations +/// merge exactly like the native pipeline. +#[wasm_bindgen] +pub struct ScannedConverter { + chars: Vec, + pages: Vec, +} + +#[wasm_bindgen] +impl ScannedConverter { + /// `dict` is the recognition dictionary text (`en_dict.txt` for the + /// default English model). + #[wasm_bindgen(constructor)] + pub fn new(dict: &str) -> Self { + Self { + chars: dict_chars(dict), + pages: Vec::new(), + } + } + + /// Convert one page (lite profile — geometric tables): `rgba` is the + /// rendered bitmap (canvas ImageData), `scale` its pixels-per-PDF-point + /// (2.0 for pdf.js `{scale: 2}`; 1.0 for a standalone image). + pub async fn add_page( + &mut self, + rgba: &[u8], + px_w: u32, + px_h: u32, + scale: f32, + layout: &LayoutSession, + rec: &RecSession, + ) -> Result<(), JsError> { + self.add_page_impl(rgba, px_w, px_h, scale, layout, rec, None) + .await + } + + /// Convert one page with TableFormer (#157 stage 3): table regions get the + /// ONNX table-structure model + docling's cell matcher instead of the + /// geometric reconstruction. `tf` is the JS-side session over the encoder / + /// decoder / bbox graphs. + #[wasm_bindgen(js_name = addPageTf)] + #[allow(clippy::too_many_arguments)] // wasm-bindgen entry: page bitmap + 3 sessions + pub async fn add_page_tf( + &mut self, + rgba: &[u8], + px_w: u32, + px_h: u32, + scale: f32, + layout: &LayoutSession, + rec: &RecSession, + tf: &TfSession, + ) -> Result<(), JsError> { + self.add_page_impl(rgba, px_w, px_h, scale, layout, rec, Some(tf)) + .await + } +} + +impl ScannedConverter { + /// Recognize a set of prepared line/word crops: width-batch them and run + /// each batch through the JS recognition session, greedy-CTC-decoding the + /// probabilities. Returns one string per input crop, in order. + async fn ocr_lines( + &self, + rec: &RecSession, + lines: &[PrepLine], + ) -> Result, JsError> { + let mut texts = vec![String::new(); lines.len()]; + for (w, chunk) in width_batches(lines) { + let data = batch_input(w, &chunk, lines); + let out = rec + .run( + chunk.len() as u32, + REC_HEIGHT, + w as u32, + js_sys::Float32Array::from(data.as_slice()), + ) + .await + .map_err(|e| JsError::new(&format!("rec session.run: {e:?}")))?; + let (probs, t_len, nc) = tensor_parts(&out)?; + if probs.len() < chunk.len() * t_len * nc { + return Err(JsError::new("rec session.run returned a short tensor")); + } + for (i, &ix) in chunk.iter().enumerate() { + texts[ix] = decode_row( + &self.chars, + &probs[i * t_len * nc..(i + 1) * t_len * nc], + nc, + ); + } + } + Ok(texts) + } + + #[allow(clippy::too_many_arguments)] + async fn add_page_impl( + &mut self, + rgba: &[u8], + px_w: u32, + px_h: u32, + scale: f32, + layout: &LayoutSession, + rec: &RecSession, + tf: Option<&TfSession>, + ) -> Result<(), JsError> { + if rgba.len() != (px_w as usize) * (px_h as usize) * 4 { + return Err(JsError::new("rgba buffer size does not match dimensions")); + } + let mut img = RgbImage::new(px_w, px_h); + for (i, px) in img.pixels_mut().enumerate() { + px.0 = [rgba[i * 4], rgba[i * 4 + 1], rgba[i * 4 + 2]]; + } + // Dark-mode screenshots invert scan polarity; normalize before both + // layout and OCR (each assumes dark ink on light paper). + let img = normalize_polarity(img); + let (page_w, page_h) = (px_w as f32 / scale, px_h as f32 / scale); + + // Layout: Rust preprocessing → JS inference → Rust decoding. + let input = layout_input(&img); + let out = layout + .run(js_sys::Float32Array::from(input.as_slice())) + .await + .map_err(|e| JsError::new(&format!("layout session.run: {e:?}")))?; + let get = |k: &str| { + js_sys::Reflect::get(&out, &JsValue::from_str(k)) + .map_err(|_| JsError::new(&format!("layout result has no `{k}`"))) + }; + let (logits, q, c) = tensor_parts(&get("logits")?)?; + let (boxes, bq, four) = tensor_parts(&get("boxes")?)?; + if bq != q || four != 4 { + return Err(JsError::new("layout boxes dims must be [1, q, 4]")); + } + let regions = decode_layout(&logits, &boxes, q, c, page_w, page_h); + let regions = refine_regions(regions, &[], page_w, page_h); + + // Diagnostic: the region-label histogram, so it's visible (browser + // console) whether the layout model flagged any `table` region on this + // page — tables only render when it does (geometric or TableFormer). + { + let mut hist: std::collections::BTreeMap<&str, usize> = + std::collections::BTreeMap::new(); + for r in ®ions { + *hist.entry(r.label).or_default() += 1; + } + let summary: Vec = hist.iter().map(|(l, n)| format!("{l}×{n}")).collect(); + diag(&format!("layout regions: {}", summary.join(", "))); + } + + // OCR the text regions (same gather/batch/decode as native ocr_page). + let (bboxes, lines) = prep_region_lines(&img, ®ions, scale); + let texts = self.ocr_lines(rec, &lines).await?; + let mut cells = Vec::new(); + for ((l, t, r, b), text) in bboxes.into_iter().zip(texts) { + let text = text.trim().to_string(); + if text.is_empty() { + continue; + } + cells.push(TextCell { text, l, t, r, b }); + } + + // Table interiors carry no words yet (prep_region_lines skips non-text + // labels, and the browser has no pdfium text layer). Recognize their + // word crops so the cell matcher — geometric or TableFormer — can fill + // the grid; assemble routes these cells into the table region, not into + // stray paragraphs. + let (tbboxes, tlines) = prep_table_words(&img, ®ions, scale); + let ttexts = self.ocr_lines(rec, &tlines).await?; + for ((l, t, r, b), text) in tbboxes.into_iter().zip(ttexts) { + let text = text.trim().to_string(); + if text.is_empty() { + continue; + } + cells.push(TextCell { text, l, t, r, b }); + } + + // TableFormer (opt-in): resolve each table region's structure through + // the ONNX graphs + shared matcher; other regions stay `None` (geometric + // fallback). The lite path passes no session → all geometric. + let table_rows = if let Some(tf) = tf { + let mut rows = Vec::with_capacity(regions.len()); + for r in ®ions { + if r.label == "table" { + rows.push( + crate::tableformer::predict_table_rows( + tf, + &img, + [r.l, r.t, r.r, r.b], + &cells, + ) + .await, + ); + } else { + rows.push(None); + } + } + rows + } else { + Vec::new() // assemble_page_with_tables resizes to all-None + }; + + let page = PdfPage::from_cells(page_w, page_h, scale, cells); + self.pages + .push(assemble_page_with_tables(&page, regions, table_rows)); + Ok(()) + } +} + +#[wasm_bindgen] +impl ScannedConverter { + /// Number of pages converted so far (progress display). + pub fn page_count(&self) -> usize { + self.pages.len() + } + + /// Assemble the accumulated pages into the final document and render it + /// as `"md"` (default) or `"json"`. Resets the converter. + pub fn finish(&mut self, name: &str, to: Option) -> Result { + let doc = finish_document(name, std::mem::take(&mut self.pages)); + match to.as_deref().unwrap_or("md") { + "md" | "markdown" => Ok(doc.export_to_markdown()), + "json" => Ok(doc.export_to_json()), + other => Err(JsError::new(&format!( + "unknown output format {other:?} (expected \"md\" or \"json\")" + ))), + } + } +} + +/// One-shot scanned-image conversion through the full lite profile (layout + +/// OCR + assembly) — the browser counterpart of the native image path +/// (a standalone image is its own page at scale 1). +#[wasm_bindgen] +pub async fn convert_scanned_image( + bytes: &[u8], + name: &str, + dict: &str, + layout: &LayoutSession, + rec: &RecSession, + to: Option, +) -> Result { + let img = image::load_from_memory(bytes) + .map_err(|e| JsError::new(&format!("decode image: {e}")))? + .to_rgba8(); + let (w, h) = img.dimensions(); + let mut conv = ScannedConverter::new(dict); + conv.add_page(img.as_raw(), w, h, 1.0, layout, rec).await?; + conv.finish(name, to) +} + +// Silence the unused warning for SIDE re-export path (the JS side sizes its +// tensor from the buffer length, but the constant documents the contract). +const _: u32 = SIDE; diff --git a/crates/docling-wasm/src/tableformer.rs b/crates/docling-wasm/src/tableformer.rs new file mode 100644 index 00000000..4b23862d --- /dev/null +++ b/crates/docling-wasm/src/tableformer.rs @@ -0,0 +1,141 @@ +//! Browser TableFormer — stage 3 of #157. The autoregressive structure loop +//! and every ONNX-free step (448 preprocessing, tag corrections, bbox +//! bookkeeping, span merge, OTSL→grid, cell matching) run in Rust via +//! `docling_pdf::tf_core` — the *same* code the native pipeline uses — and only +//! the three ONNX graphs (encoder / decoder / bbox) are delegated to ONNX +//! Runtime Web through the [`TfSession`] interop object. +//! +//! The heavy tensors never cross the wasm boundary: [`TfSession`] holds the +//! encoder's constant cross-attention K/V and `enc_out`, and the growing +//! decoder KV-cache, entirely on the JS side. Each decode step sends only the +//! last tag (one int) and gets back `logits` + `hidden` (525 floats). See +//! `www/scan.html` for the JS wiring (the `decoder_kv` graph: N_LAYERS=6, +//! KV_HEADS=8, head_dim=64, cross length 784). + +use docling_pdf::pdfium_backend::TextCell; +use docling_pdf::tf_core::{ + argmax, build_table_cells, merge_spans, preprocess_input, table_rows, BboxBook, TableCell, + MAX_STEPS, +}; +use image::RgbImage; +use wasm_bindgen::prelude::*; + +#[wasm_bindgen] +extern "C" { + /// The JS-side TableFormer session: a stateful wrapper around the three + /// `ort.InferenceSession`s. `encode` runs the image encoder and stashes the + /// constant cross tensors + `enc_out` and resets the KV-cache; `step` runs + /// one decoder step (feeding the stored cross + growing cache) and returns + /// `{ logits: Float32Array, hidden: Float32Array }`; `bbox` runs the bbox + /// decoder over the collected per-cell hidden states and returns + /// `{ boxes: Float32Array, classes: Float32Array }`. + pub type TfSession; + + #[wasm_bindgen(method, catch)] + async fn encode(this: &TfSession, image: js_sys::Float32Array) -> Result; + + #[wasm_bindgen(method, catch)] + async fn step(this: &TfSession, tag: i32) -> Result; + + #[wasm_bindgen(method, catch)] + async fn bbox( + this: &TfSession, + tag_h: js_sys::Float32Array, + n: u32, + ) -> Result; +} + +/// Pull a named `Float32Array` out of a JS result object. +fn get_f32(obj: &JsValue, key: &str) -> Result, JsError> { + let v = js_sys::Reflect::get(obj, &JsValue::from_str(key)) + .map_err(|_| JsError::new(&format!("tf result has no `{key}`")))?; + let arr: js_sys::Float32Array = v + .dyn_into() + .map_err(|_| JsError::new(&format!("tf `{key}` is not a Float32Array")))?; + Ok(arr.to_vec()) +} + +/// Run the full structure model on a `SIDE×SIDE`-croppable region image: encode +/// once, step the decoder to the OTSL sequence (`tf_core::BboxBook` drives the +/// exact bbox bookkeeping), run the bbox decoder, then merge spans and lay the +/// cells onto the grid — all shared with native. +async fn predict_structure( + session: &TfSession, + crop: &RgbImage, +) -> Result, JsError> { + let input = preprocess_input(crop); + session + .encode(js_sys::Float32Array::from(input.as_slice())) + .await + .map_err(|e| JsError::new(&format!("tf encode: {e:?}")))?; + + let mut book = BboxBook::new(); + while book.otsl.len() < MAX_STEPS { + let last = *book.tags.last().expect("decode starts from "); + let out = session + .step(last as i32) + .await + .map_err(|e| JsError::new(&format!("tf decode step: {e:?}")))?; + let logits = get_f32(&out, "logits")?; + let hidden = get_f32(&out, "hidden")?; + let raw = argmax(&logits) as i64; + if !book.step(raw, &hidden) { + break; + } + } + if book.n == 0 { + return Ok(Vec::new()); + } + + let out = session + .bbox( + js_sys::Float32Array::from(book.hiddens.as_slice()), + book.n as u32, + ) + .await + .map_err(|e| JsError::new(&format!("tf bbox: {e:?}")))?; + let boxes: Vec<[f32; 4]> = get_f32(&out, "boxes")? + .chunks_exact(4) + .map(|c| [c[0], c[1], c[2], c[3]]) + .collect(); + let classes: Vec = get_f32(&out, "classes")? + .chunks_exact(3) + .map(|c| argmax(c) as i64) + .collect(); + let (merged, merged_classes) = merge_spans(&boxes, &classes, &book.merge); + Ok(build_table_cells(&book.otsl, &merged, &merged_classes)) +} + +/// Predict a table region's grid, the browser counterpart of the native +/// `TableFormer::predict_table_rows`: docling's page→1024px box-average + crop +/// (`docling_pdf::resample`), the ONNX structure model, then the shared word +/// matcher (`tf_core::table_rows`). `page_image` is the rendered page (2 px per +/// point, like the native pipeline); `region` is `(l, t, r, b)` in page points. +/// `None` when no structure is predicted. +pub(crate) async fn predict_table_rows( + session: &TfSession, + page_image: &RgbImage, + region: [f32; 4], + words: &[TextCell], +) -> Option>> { + // page → 1024px height (cv2.INTER_AREA), then crop the table bbox — docling's + // coordinate chain with the same rounding the native path reproduces. + let sf = 1024.0 / page_image.height() as f32; + let pw = (page_image.width() as f32 * sf) as u32; + let page1024 = docling_pdf::resample::inter_area(page_image, pw, 1024); + let k = 2.0 * 1024.0 / page_image.height() as f64; + let px = |v: f32| (v as f64).round_ties_even() * k; + let x = (px(region[0]).round_ties_even()).max(0.0) as u32; + let y = (px(region[1]).round_ties_even()).max(0.0) as u32; + let x2 = (px(region[2]).round_ties_even() as u32).min(page1024.width()); + let y2 = (px(region[3]).round_ties_even() as u32).min(page1024.height()); + if x2 <= x || y2 <= y { + return None; + } + let crop = image::imageops::crop_imm(&page1024, x, y, x2 - x, y2 - y).to_image(); + let cells = predict_structure(session, &crop).await.ok()?; + if cells.is_empty() { + return None; + } + table_rows(&cells, region, words) +} diff --git a/crates/docling-wasm/www/coi.js b/crates/docling-wasm/www/coi.js new file mode 100644 index 00000000..c24bf4ae --- /dev/null +++ b/crates/docling-wasm/www/coi.js @@ -0,0 +1,55 @@ +// Cross-origin isolation shim, so ONNX Runtime Web can spin up multi-threaded +// wasm (SharedArrayBuffer needs `crossOriginIsolated`, which needs the page to +// be served with COOP+COEP). A plain static file server (python -m http.server, +// GitHub Pages) can't set those headers, so a service worker re-serves every +// response with them added — the standard "coi-serviceworker" trick. +// +// COEP is `credentialless` (not `require-corp`): cross-origin subresources — +// the jsDelivr ORT/pdf.js bundles and the Hugging Face recognition model — +// then load without needing their own CORP header, while the page still counts +// as cross-origin isolated. Browsers that don't support credentialless simply +// stay non-isolated and ORT falls back to a single thread (the page still +// works, just slower). +// +// One-liner to enable: before the module script. + +if (typeof window === "undefined") { + // --- service-worker context ----------------------------------------------- + self.addEventListener("install", () => self.skipWaiting()); + self.addEventListener("activate", (e) => e.waitUntil(self.clients.claim())); + + self.addEventListener("fetch", (event) => { + const req = event.request; + // `only-if-cached` is only valid with same-origin `no-cors`; leave it alone. + if (req.cache === "only-if-cached" && req.mode !== "same-origin") return; + + event.respondWith( + fetch(req) + .then((res) => { + if (res.status === 0) return res; // opaque response — pass through + const headers = new Headers(res.headers); + headers.set("Cross-Origin-Embedder-Policy", "credentialless"); + headers.set("Cross-Origin-Opener-Policy", "same-origin"); + return new Response(res.body, { + status: res.status, + statusText: res.statusText, + headers, + }); + }) + .catch((e) => console.error(e)), + ); + }); +} else { + // --- page context: register self, reload once to gain isolation ----------- + (async () => { + if (window.crossOriginIsolated) return; // already isolated — nothing to do + if (!window.isSecureContext || !navigator.serviceWorker) return; // can't + const reg = await navigator.serviceWorker + .register(document.currentScript.src) + .catch(() => null); + if (!reg) return; + // A fresh registration doesn't control this load yet; reload so the SW can + // stamp COOP/COEP on the document and crossOriginIsolated flips true. + if (!navigator.serviceWorker.controller) window.location.reload(); + })(); +} diff --git a/crates/docling-wasm/www/cyrillic_v5_dict.txt b/crates/docling-wasm/www/cyrillic_v5_dict.txt new file mode 100644 index 00000000..7625a74e --- /dev/null +++ b/crates/docling-wasm/www/cyrillic_v5_dict.txt @@ -0,0 +1,850 @@ +! +" +# +$ +% +& +' +( +) +* ++ +, +- +. +/ +0 +1 +2 +3 +4 +5 +6 +7 +8 +9 +: +; +< += +> +? +A +B +C +D +E +F +G +H +I +J +K +L +M +N +O +P +Q +R +S +T +U +V +W +X +Y +Z +[ +] +_ +` +a +b +c +d +e +f +g +h +i +j +k +l +m +n +o +p +q +r +s +t +u +v +w +x +y +z +© +‥ +{ +} +\ +| +@ +^ +~ +÷ +∕ +∙ +⋅ +· +± +∓ +∩ +∪ +□ +← +↔ +⇒ +⇐ +⇔ +∀ +∃ +∄ +∴ +∵ +∝ +∞ +⊥ +∟ +∠ +∡ +∢ +′ +″ +∥ +⊾ +⊿ +∂ +∫ +∬ +∭ +∮ +∯ +∰ +∑ +∏ +√ +∛ +∜ +∱ +∲ +∳ +∶ +∷ +∼ +® +℉ +Ω +℧ +Å +⌀ +ℏ +⅀ +⍺ +⍵ +¢ +€ +£ +¥ +₿ +Ⅰ +Ⅱ +Ⅲ +Ⅳ +Ⅴ +Ⅵ +Ⅶ +Ⅷ +Ⅸ +Ⅹ +Ⅺ +Ⅻ +ⅰ +ⅱ +ⅲ +ⅳ +ⅴ +ⅵ +ⅶ +ⅷ +ⅸ +ⅹ +ⅺ +ⅻ +➀ +➁ +➂ +➃ +➄ +➅ +➆ +➇ +➈ +➉ +➊ +➋ +➌ +➍ +➎ +➏ +➐ +➑ +➒ +➓ +❶ +❷ +❸ +❹ +❺ +❻ +❼ +❽ +❾ +❿ +① +② +③ +④ +⑤ +⑥ +⑦ +⑧ +⑨ +⑩ +● +▶ +𝑢 +︽ +– +﹥ +𝜓 +• +∋ +ƒ +० +⬆ +Ạ +◀ + +▫ +︾ +À +Á +Â +Ã +Ä +Å +Æ +Ç +È +É +Ê +Ë +Ì +Í +Î +Ï +Ð +Ñ +Ò +Ó +Ô +Õ +Ö +Ø +Ù +Ú +Û +Ü +Ý +Þ +à +á +â +ã +ä +å +æ +ç +è +é +ê +ë +ì +í +î +ï +ð +ñ +ò +ó +ô +õ +ö +ø +ù +ú +û +ü +ý +þ +ÿ +¡ +¤ +¦ +§ +¨ +ª +« +¬ +¯ +° +² +³ +´ +µ +¶ +¸ +¹ +º +» +¼ +½ +¾ +¿ +× +‐ +‑ +‒ +— +― +‖ +‗ +‘ +’ +‚ +‛ +“ +” +„ +‟ +† +‡ +‣ +․ +… +‧ +‰ +‴ +‵ +‶ +‷ +‸ +‹ +› +※ +‼ +‽ +‾ +₤ +₡ +₹ +− +∖ +∗ +≈ +≠ +≡ +≤ +≥ +⊂ +⊃ +↑ +→ +↓ +↕ +™ +Ω +℮ +∆ +✓ +✗ +✘ +▪ +◼ +✔ +✕ +☑ +☒ +№ +₽ +₴ +Α +α +Β +β +Γ +γ +Δ +δ +Ε +ε +Ζ +ζ +Η +η +Θ +θ +Ι +ι +Κ +κ +Λ +λ +Μ +μ +Ν +ν +Ξ +ξ +Ο +ο +Π +π +Ρ +ρ +Σ +σ +ς +Τ +τ +Υ +υ +Φ +φ +Χ +χ +Ψ +ψ +ω +Ѐ +Ё +Ђ +Ѓ +Є +Ѕ +І +Ї +Ј +Љ +Њ +Ћ +Ќ +Ѝ +Ў +Џ +А +Б +В +Г +Д +Е +Ж +З +И +Й +К +Л +М +Н +О +П +Р +С +Т +У +Ф +Х +Ц +Ч +Ш +Щ +Ъ +Ы +Ь +Э +Ю +Я +а +б +в +г +д +е +ж +з +и +й +к +л +м +н +о +п +р +с +т +у +ф +х +ц +ч +ш +щ +ъ +ы +ь +э +ю +я +ѐ +ё +ђ +ѓ +є +ѕ +і +ї +ј +љ +њ +ћ +ќ +ѝ +ў +џ +Ѡ +ѡ +Ѣ +ѣ +Ѥ +ѥ +Ѧ +ѧ +Ѩ +ѩ +Ѫ +ѫ +Ѭ +ѭ +Ѯ +ѯ +Ѱ +ѱ +Ѳ +ѳ +Ѵ +ѵ +Ѷ +ѷ +Ѹ +ѹ +Ѻ +ѻ +Ѽ +ѽ +Ѿ +ѿ +Ҁ +ҁ +Ҋ +ҋ +Ҍ +ҍ +Ҏ +ҏ +Ґ +ґ +Ғ +ғ +Ҕ +ҕ +Җ +җ +Ҙ +ҙ +Қ +қ +Ҝ +ҝ +Ҟ +ҟ +Ҡ +ҡ +Ң +ң +Ҧ +ҧ +Ҩ +ҩ +Ҫ +ҫ +Ҭ +ҭ +Ү +ү +Ұ +ұ +Ҳ +ҳ +Ҷ +ҷ +Ҹ +ҹ +Һ +һ +Ҽ +ҽ +Ҿ +ҿ +Ӏ +Ӂ +ӂ +Ӄ +ӄ +Ӆ +ӆ +Ӈ +ӈ +Ӊ +ӊ +Ӌ +ӌ +Ӎ +ӎ +ӏ +Ӑ +ӑ +Ӓ +ӓ +Ӗ +ӗ +Ә +ә +Ӛ +ӛ +Ӝ +ӝ +Ӟ +ӟ +Ӡ +ӡ +Ӣ +ӣ +Ӥ +ӥ +Ӧ +ӧ +Ө +ө +Ӫ +ӫ +Ӭ +ӭ +Ӯ +ӯ +Ӱ +ӱ +Ӳ +ӳ +Ӵ +ӵ +Ӷ +ӷ +Ӹ +ӹ +Ӻ +ӻ +Ӽ +ӽ +Ӿ +ӿ +Ԁ +ԁ +Ԃ +ԃ +Ԅ +ԅ +Ԇ +ԇ +Ԉ +ԉ +Ԋ +ԋ +Ԍ +ԍ +Ԏ +ԏ +Ԑ +ԑ +Ԓ +ԓ +Ԕ +ԕ +Ԗ +ԗ +Ԙ +ԙ +Ԛ +ԛ +Ԝ +ԝ +Ԟ +ԟ +Ԡ +ԡ +Ԣ +ԣ +Ԥ +ԥ +Ԧ +ԧ +Ԩ +ԩ +Ԫ +ԫ +Ԭ +ԭ +Ԯ +ԯ +ⷠ +ⷡ +ⷢ +ⷣ +ⷤ +ⷥ +ⷦ +ⷧ +ⷨ +ⷩ +ⷪ +ⷫ +ⷬ +ⷭ +ⷮ +ⷯ +ⷰ +ⷱ +ⷲ +ⷳ +ⷴ +ⷵ +ⷶ +ⷷ +ⷸ +ⷹ +ⷺ +ⷻ +ⷼ +ⷽ +ⷾ +ⷿ +Ꙁ +ꙁ +Ꙃ +ꙃ +Ꙅ +ꙅ +Ꙇ +ꙇ +Ꙉ +ꙉ +Ꙋ +ꙋ +Ꙍ +ꙍ +Ꙏ +ꙏ +Ꙑ +ꙑ +Ꙓ +ꙓ +Ꙕ +ꙕ +Ꙗ +ꙗ +Ꙙ +ꙙ +Ꙛ +ꙛ +Ꙝ +ꙝ +Ꙟ +ꙟ +Ꙡ +ꙡ +Ꙣ +ꙣ +Ꙥ +ꙥ +Ꙧ +ꙧ +Ꙩ +ꙩ +Ꙫ +ꙫ +Ꙭ +ꙭ +ꙮ +ꙴ +ꙵ +ꙶ +ꙷ +ꙸ +ꙹ +ꙺ +ꙻ +Ꚁ +ꚁ +Ꚃ +ꚃ +Ꚅ +ꚅ +Ꚇ +ꚇ +Ꚉ +ꚉ +Ꚋ +ꚋ +Ꚍ +ꚍ +Ꚏ +ꚏ +Ꚑ +ꚑ +Ꚓ +ꚓ +Ꚕ +ꚕ +Ꚗ +ꚗ +Ꚙ +ꚙ +Ꚛ +ꚛ +ꚜ +ꚝ +ꚞ +ꚟ diff --git a/crates/docling-wasm/www/ocr.html b/crates/docling-wasm/www/ocr.html new file mode 100644 index 00000000..7716c8cd --- /dev/null +++ b/crates/docling-wasm/www/ocr.html @@ -0,0 +1,219 @@ + + + + + + + + docling.rs — in-browser OCR (wasm + ONNX Runtime Web) + + + +

docling.rs → wasm: OCR a scanned image in your browser

+

+ Stage 1 of #157: + the PP-OCRv3 recognition model (~10 MB) runs via + ONNX Runtime Web; + line segmentation and CTC decoding run in Rust/wasm — the same code as the + native pipeline. The image never leaves this page. The model + dictionary + are fetched once and cached by the browser. Single-column scans work best + (layout detection is stage 2). +

+
drop a scanned image here (PNG/JPEG/TIFF/WebP) or click to pick
+ +

+ Recognition language: + + +

+ +
+ + + + diff --git a/crates/docling-wasm/www/pipeline.js b/crates/docling-wasm/www/pipeline.js new file mode 100644 index 00000000..c6d54015 --- /dev/null +++ b/crates/docling-wasm/www/pipeline.js @@ -0,0 +1,288 @@ +// wasm-side OCR pipeline (stage 2 of #157), used inside worker.js (default) or +// on the main thread (fallback). It owns ONLY the wasm work — model loading, +// per-page `add_page`, image decode. Rasterization stays on the MAIN thread +// (pdf.js + a real HTMLCanvas, which produces correct pixels and uses pdf.js's +// own worker); pages arrive here as ready RGBA buffers. Keeping pdf.js off the +// Web Worker avoids its "fake worker" fallback, which garbled the line crops. +// +// DOM-free: the only host concern is onStatus(msg, spinning) for progress. + +import * as ort from "https://cdn.jsdelivr.net/npm/onnxruntime-web/dist/ort.min.mjs"; +import init, { ScannedConverter, convert_scanned_image } from "./pkg/docling_wasm.js"; + +// Multi-threaded wasm when cross-origin isolated (coi.js); else one thread. +ort.env.wasm.numThreads = self.crossOriginIsolated + ? Math.min(navigator.hardwareConcurrency || 4, 8) + : 1; +export const THREADS = ort.env.wasm.numThreads; + +// Model bases, tried in order after any user-provided file: local ./models/ +// first (download_dependencies.sh for local dev), then a CORS-enabled Hugging +// Face mirror so the page works cross-origin (the hosted phone demo on +// raw.githack) — GitHub Release assets carry no CORS header. +const MODEL_BASE = "https://huggingface.co/pivozavrus/docling-rs-models/resolve/main/"; + +const REC_MODELS = { + en: { + model: "https://huggingface.co/SWHL/RapidOCR/resolve/main/PP-OCRv3/en_PP-OCRv3_rec_infer.onnx", + dict: "https://raw.githubusercontent.com/PaddlePaddle/PaddleOCR/main/ppocr/utils/en_dict.txt", + }, + cyrillic: { + // PP-OCRv5: markedly better Cyrillic accuracy than the v3 export (spaces + // and case survive); its dictionary only exists inside the repo's + // inference.yml, so a flattened copy ships next to this page. + model: "https://huggingface.co/PaddlePaddle/cyrillic_PP-OCRv5_mobile_rec_onnx/resolve/main/inference.onnx", + dict: "cyrillic_v5_dict.txt", + }, + ch: { + model: "https://huggingface.co/SWHL/RapidOCR/resolve/main/PP-OCRv3/ch_PP-OCRv3_rec_infer.onnx", + dict: "https://raw.githubusercontent.com/PaddlePaddle/PaddleOCR/main/ppocr/utils/ppocr_keys_v1.txt", + }, +}; + +// TableFormer graphs: local ./models/tableformer/ first (download_dependencies.sh), +// then the CORS Hugging Face mirror for the hosted demo. The encoder is +// self-contained; the decoder/bbox use ONNX external data, so their .onnx.data +// sidecar rides along via ort-web's externalData option (path = the location +// stored in the .onnx). Fetched lazily — ~380 MB total — only when a table +// profile is chosen. +const TF_DIRS = ["./models/tableformer/", MODEL_BASE]; + +/// The stateful session the wasm TfSession interop expects (see +/// src/tableformer.rs): `encode` runs the image encoder once and stashes the +/// constant cross-attention K/V + enc_out and resets the KV-cache; `step` runs +/// one decoder step feeding the stored cross + growing cache; `bbox` runs the +/// bbox decoder. The heavy tensors stay here — only tags and logits/hidden +/// cross the wasm boundary. Geometry: 6 layers, 8 KV heads, head_dim 64. +class JsTfSession { + constructor(enc, dec, bbox) { + this.enc = enc; + this.dec = dec; + this.bboxSess = bbox; + this.cross = null; + this.encOut = null; + this.cacheK = null; + this.cacheV = null; + } + async encode(image) { + const out = await this.enc.run({ image: new ort.Tensor("float32", image, [1, 3, 448, 448]) }); + this.cross = {}; + for (let i = 0; i < 6; i++) { + this.cross["cross_kt_" + i] = out["cross_kt_" + i]; + this.cross["cross_v_" + i] = out["cross_v_" + i]; + } + this.encOut = out["enc_out"]; + // Empty first-step KV-cache: [N_LAYERS, 1, KV_HEADS, past=0, head_dim]. + this.cacheK = new ort.Tensor("float32", new Float32Array(0), [6, 1, 8, 0, 64]); + this.cacheV = new ort.Tensor("float32", new Float32Array(0), [6, 1, 8, 0, 64]); + } + async step(tag) { + const out = await this.dec.run({ + tag: new ort.Tensor("int64", BigInt64Array.from([BigInt(tag)]), [1, 1]), + cache_k: this.cacheK, + cache_v: this.cacheV, + ...this.cross, + }); + this.cacheK = out["out_cache_k"]; + this.cacheV = out["out_cache_v"]; + return { logits: out["logits"].data, hidden: out["hidden"].data }; + } + async bbox(tagH, n) { + const out = await this.bboxSess.run({ + enc_out: this.encOut, + tag_h: new ort.Tensor("float32", tagH, [n, 512]), + }); + return { boxes: out["boxes"].data, classes: out["classes"].data }; + } +} + +export function createOcr({ onStatus }) { + const status = (msg, spinning = true) => onStatus && onStatus(msg, spinning); + + // Model files the user picked from the device (basename → ArrayBuffer), + // used instead of any network fetch — see scan.html's file picker. Reading a + // File to an ArrayBuffer is a single allocation, so it also sidesteps the + // double-buffering peak fetchProgress hits on the 225 MB encoder. + let provided = {}; + function setProvidedModels(map) { + provided = map || {}; + } + // Resolve a model by name: a user-provided file wins, else the first base + // that serves it. Returns { buf, base } (base is null for a provided file). + async function fetchModel(name, bases, label) { + if (provided[name]) return { buf: provided[name], base: null }; + let lastErr = null; + for (const b of bases) { + try { + return { buf: await fetchProgress(b + name, label), base: b }; + } catch (e) { + lastErr = e; + } + } + throw new Error(`${name} failed: ${(lastErr && lastErr.message) || lastErr}`); + } + + // fetch() with a live "x / y MB" progress line. + async function fetchProgress(url, label) { + const resp = await fetch(url, { cache: "force-cache" }); + if (!resp.ok) throw new Error(`${label}: HTTP ${resp.status}`); + const total = Number(resp.headers.get("Content-Length")) || 0; + if (!resp.body) return resp.arrayBuffer(); + const reader = resp.body.getReader(); + const chunks = []; + let got = 0; + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + chunks.push(value); + got += value.length; + const mb = (got / 1048576).toFixed(1); + status(total ? `${label} — ${mb} / ${(total / 1048576).toFixed(1)} MB` : `${label} — ${mb} MB`, true); + } + const buf = new Uint8Array(got); + let off = 0; + for (const c of chunks) { buf.set(c, off); off += c.length; } + return buf.buffer; + } + + const recCache = {}; + async function recFor(lang) { + if (!recCache[lang]) { + const [model, dict] = await Promise.all([ + fetchProgress(REC_MODELS[lang].model, `${lang} recognition model`), + fetch(REC_MODELS[lang].dict, { cache: "force-cache" }).then((r) => r.text()), + ]); + status(`starting ${lang} recognition session …`, true); + const session = await ort.InferenceSession.create(model, { + executionProviders: ["wasm"], + logSeverityLevel: 3, + }); + recCache[lang] = { + dict, + rec: { + run: async (n, h, w, data) => { + const results = await session.run({ + [session.inputNames[0]]: new ort.Tensor("float32", data, [n, 3, h, w]), + }); + const t = results[session.outputNames[0]]; + return { data: t.data, dims: Array.from(t.dims) }; + }, + }, + }; + } + return recCache[lang]; + } + + // Interop wrapper docling_wasm expects (see src/scanned.rs). + let layout = null; + let layoutKind = null; + async function loadLayout() { + for (const [name, kind] of [ + ["layout_heron_int8.onnx", "int8"], + ["layout_heron.onnx", "fp32"], + ]) { + let buf; + try { + ({ buf } = await fetchModel(name, ["./models/", MODEL_BASE], "layout model (first load only)")); + } catch (e) { + continue; // not provided and not fetchable → try the next candidate + } + status("starting layout session …", true); + const session = await ort.InferenceSession.create(buf, { + executionProviders: ["wasm"], + logSeverityLevel: 3, + }); + layoutKind = kind; + layout = { + run: async (data) => { + const results = await session.run({ + pixel_values: new ort.Tensor("float32", data, [1, 3, 640, 640]), + }); + const t = (n) => ({ data: results[n].data, dims: Array.from(results[n].dims) }); + return { logits: t("logits"), boxes: t("pred_boxes") }; + }, + }; + return layoutKind; + } + return null; + } + + // Bring wasm + the layout model up. Returns "int8" | "fp32" | null. + async function boot() { + status("loading wasm module …", true); + await init(); + return loadLayout(); + } + + // One blank inference through layout + rec so ORT's lazy kernel/thread init + // happens now, not inside the first real page. + async function warmup(lang) { + try { + await layout.run(new Float32Array(3 * 640 * 640)); + const { rec } = await recFor(lang); + await rec.run(1, 48, 320, new Float32Array(3 * 48 * 320)); + } catch (e) { + // best-effort — a failure just means the first page pays it. + } + } + + // TableFormer sessions, loaded lazily on first use (the encoder alone is + // ~225 MB, so this only downloads when the table profile is chosen). + let tf = null; + async function ensureTf() { + if (tf) return tf; + const load = async (name, external) => { + // A provided file wins; else the first base (local, then HF). External + // data comes from a provided file or the same base the .onnx came from. + const { buf: model, base } = await fetchModel(name + ".onnx", TF_DIRS, `tableformer ${name}`); + const opts = { executionProviders: ["wasm"], logSeverityLevel: 3 }; + if (external) { + const { buf: data } = await fetchModel( + name + ".onnx.data", + base ? [base] : TF_DIRS, + `tableformer ${name} data`, + ); + opts.externalData = [{ path: name + ".onnx.data", data: new Uint8Array(data) }]; + } + return ort.InferenceSession.create(model, opts); + }; + const enc = await load("encoder", false); + const dec = await load("decoder_kv", true); + const bbox = await load("bbox", true); + status("starting tableformer sessions …", true); + tf = new JsTfSession(enc, dec, bbox); + return tf; + } + + // Multi-page document lifecycle: startDoc → addPage* → finishDoc. Pages come + // in as RGBA (already rasterized on the main thread), one document at a time. + let cur = null; + async function startDoc(lang, useTf) { + const { dict, rec } = await recFor(lang); + const tfSess = useTf ? await ensureTf() : null; + cur = { conv: new ScannedConverter(dict), rec, tf: tfSess }; + } + async function addPage(rgba, w, h, scale) { + if (cur.tf) { + await cur.conv.addPageTf(rgba, w, h, scale, layout, cur.rec, cur.tf); + } else { + await cur.conv.add_page(rgba, w, h, scale, layout, cur.rec); + } + } + function finishDoc(name) { + const md = cur.conv.finish(name, "md"); + cur = null; + return md; + } + + // Standalone image: the wasm side decodes it (no canvas needed). + async function convertImage(bytes, name, lang) { + const { dict, rec } = await recFor(lang); + return convert_scanned_image(new Uint8Array(bytes), name, dict, layout, rec, "md"); + } + + return { + boot, warmup, recFor, startDoc, addPage, finishDoc, convertImage, setProvidedModels, + get layoutKind() { return layoutKind; }, + }; +} diff --git a/crates/docling-wasm/www/scan.html b/crates/docling-wasm/www/scan.html new file mode 100644 index 00000000..48623064 --- /dev/null +++ b/crates/docling-wasm/www/scan.html @@ -0,0 +1,347 @@ + + + + + + + + docling.rs — scanned PDFs in the browser (lite profile) + + + +

docling.rs → wasm: scanned PDF / image, full lite profile

+

+ Stage 2 of #157: + RT-DETR layout detection + PP-OCRv3 recognition via ONNX Runtime Web; + pdf.js rasterizes PDF pages at the native pipeline's 2 px/point; + region refinement, OCR and reading-order assembly run in Rust/wasm — the + same code as the native pipeline. Tables come out via the geometric + reconstruction (TableFormer is stage 3). Nothing leaves this page. +

+
+ The layout model is served from this page's own origin — fetch it once into ./models/ +

+ Unlike the recognition model (fetched from Hugging Face, which sends CORS + headers), the RT-DETR layout export lives on this repo's GitHub Release, + and GitHub Release assets can't be fetched cross-origin from a + browser (no Access-Control-Allow-Origin). So this + page loads the layout model same-origin, from a + models/ directory next to it. Run the installer from + this www/ directory — it writes ./models/ + relative to the current directory: +

+ 
cd crates/docling-wasm/www/
+curl -fsSL https://raw.githubusercontent.com/docling-project/docling.rs/master/scripts/install/download_dependencies.sh | sh
+

+ That populates ./models/ (this page reads + ./models/layout_heron_int8.onnx, ~68 MB, falling back to + ./models/layout_heron.onnx, ~172 MB — int8 needs a + release that hosts it, otherwise fp32 is fetched). These are the exact + conformance-validated models the native pipeline uses; both decode + identically here (same pixel_values → logits, pred_boxes + contract). The recognition model still streams from Hugging Face on demand. +

+
+
drop a scanned PDF or image here, or click to pick
+ +

+ Recognition language: + + + +

+

+ Models from device (optional — skips the download): + +
pick layout_heron_int8.onnx and, for TableFormer, encoder.onnx, + decoder_kv.onnx(+.data), bbox.onnx(+.data). +

+ +
+
+ + + + diff --git a/crates/docling-wasm/www/worker.js b/crates/docling-wasm/www/worker.js new file mode 100644 index 00000000..f4dffe9b --- /dev/null +++ b/crates/docling-wasm/www/worker.js @@ -0,0 +1,66 @@ +// Web Worker host for the wasm OCR pipeline (stage 2 of #157). The heavy work +// — layout + OCR inference and the synchronous Rust pre/post-processing inside +// add_page — runs here, so the main thread stays responsive. Rasterization is +// NOT here: the main thread rasterizes with pdf.js and transfers ready RGBA +// buffers in (pdf.js in a worker falls back to a "fake worker" that garbles the +// raster). +// +// RPC: every request carries an id; the reply is {type:"ok", id, ...data} or +// {type:"error", id, msg}. Progress is a broadcast {type:"status", msg, +// spinning}. Requests: boot{lang} | rec{lang} | doc-start{lang} | +// doc-page{rgba,w,h,scale} | doc-finish{name} | convert-image{bytes,name,lang}. + +import { createOcr, THREADS } from "./pipeline.js"; + +const post = (type, extra) => self.postMessage({ type, ...extra }); + +// On-page diagnostic sink the wasm calls (see src/scanned.rs) — forward to the +// main thread, which renders it (a phone can't open the console). +globalThis.__docling_diag = (msg) => post("diag", { msg }); + +const ocr = createOcr({ + onStatus: (msg, spinning) => post("status", { msg, spinning }), +}); + +async function handle(m) { + switch (m.type) { + case "set-models": + ocr.setProvidedModels(m.models); + return {}; + case "boot": { + const kind = await ocr.boot(); + if (!kind) return { noLayout: true }; + await ocr.recFor(m.lang); + await ocr.warmup(m.lang); + return { kind, threads: THREADS }; + } + case "rec": + await ocr.recFor(m.lang); + await ocr.warmup(m.lang); + return {}; + case "doc-start": + await ocr.startDoc(m.lang, m.useTf); + return {}; + case "doc-page": + await ocr.addPage(new Uint8Array(m.rgba), m.w, m.h, m.scale); + return {}; + case "doc-finish": + return { md: ocr.finishDoc(m.name) }; + case "convert-image": + return { md: await ocr.convertImage(m.bytes, m.name, m.lang) }; + default: + throw new Error(`unknown request ${m.type}`); + } +} + +self.onmessage = async (e) => { + const m = e.data; + try { + const r = (await handle(m)) || {}; + post("ok", { id: m.id, ...r }); + } catch (err) { + post("error", { id: m.id, msg: String((err && err.message) || err) }); + } +}; + +post("up"); diff --git a/crates/docling/src/backend/images.rs b/crates/docling/src/backend/images.rs index c5f7ad37..d0375e06 100644 --- a/crates/docling/src/backend/images.rs +++ b/crates/docling/src/backend/images.rs @@ -54,6 +54,10 @@ impl ImageResolver for NoFetch { /// itself fetched from the web. pub(crate) struct FsImageResolver { base_dir: Option, + /// Only read by the gated remote-URL resolution; without `fetch-images` + /// (e.g. the wasm32 build) it is stored-but-unused so `new`'s signature + /// stays the same across feature shapes. + #[cfg_attr(not(feature = "fetch-images"), allow(dead_code))] base_url: Option, /// Memoized remote fetches, keyed by the resolved absolute URL: fills as /// [`prefetch`](Self::prefetch) warms it (concurrently) and as `resolve` @@ -549,10 +553,3 @@ mod tests { std::env::remove_var("DOCLING_RS_ALLOW_PRIVATE_IP_FETCH"); } } - -/// Compiled without the `fetch-images` feature (no HTTP client — e.g. the -/// wasm32 build): remote images stay placeholders, same as a fetch failure. -#[cfg(not(feature = "fetch-images"))] -fn fetch_remote(_url: &str) -> Option { - None -}