|
| 1 | +//! Browser pipeline for **digital** PDFs — the ones that carry a text layer |
| 2 | +//! (#157). It is the scanned pipeline with the expensive half removed: the text |
| 3 | +//! comes from the PDF itself rather than from OCR, so only the layout model |
| 4 | +//! runs, and the result is both faster and exact (no recognition errors, no |
| 5 | +//! mangled umlauts). |
| 6 | +//! |
| 7 | +//! This is what the native pipeline does for a digital PDF — it OCRs a page |
| 8 | +//! only when `page.cells` is empty — so the same three steps happen here in the |
| 9 | +//! same order: layout detection, region refinement against the real text cells, |
| 10 | +//! then TableFormer (or the geometric fallback) per table region. |
| 11 | +//! |
| 12 | +//! Without it the browser had to choose between structure and fidelity: the |
| 13 | +//! pure text-layer path (`convert`) is milliseconds but emits flat paragraphs, |
| 14 | +//! because headings, tables and pictures are all things the layout model finds. |
| 15 | +//! |
| 16 | +//! ```js |
| 17 | +//! const conv = new DigitalConverter(pdfBytes); // throws when there is no text layer |
| 18 | +//! for (let i = 0; i < conv.page_count(); i++) { |
| 19 | +//! // rasterize page i with pdf.js at 2 px/point, then: |
| 20 | +//! await conv.add_page(i, rgba, w, h, 2.0, layout); |
| 21 | +//! } |
| 22 | +//! const markdown = conv.finish("bill.pdf", "md", "embedded"); |
| 23 | +//! ``` |
| 24 | +
|
| 25 | +use docling_pdf::assemble::{geometric_table_is_reliable, reconstruct_table}; |
| 26 | +use docling_pdf::layout::{decode_layout, layout_input}; |
| 27 | +use docling_pdf::pdfium_backend::PdfPage; |
| 28 | +use docling_pdf::scanned::{assemble_page_with_tables, finish_document, refine_regions}; |
| 29 | +use image::RgbImage; |
| 30 | +use wasm_bindgen::prelude::*; |
| 31 | + |
| 32 | +use crate::ocr::tensor_parts; |
| 33 | +use crate::scanned::{render, LayoutSession}; |
| 34 | +use crate::tableformer::TfSession; |
| 35 | + |
| 36 | +/// A digital PDF being converted page by page. Construct it from the file's |
| 37 | +/// bytes (the text layer is parsed once, in Rust), then feed the rasterized |
| 38 | +/// pages in order. |
| 39 | +#[wasm_bindgen] |
| 40 | +pub struct DigitalConverter { |
| 41 | + /// One entry per page, carrying that page's text cells in PDF points. |
| 42 | + pages: Vec<PdfPage>, |
| 43 | + out: Vec<docling_pdf::scanned::AssembledPage>, |
| 44 | +} |
| 45 | + |
| 46 | +#[wasm_bindgen] |
| 47 | +impl DigitalConverter { |
| 48 | + /// Parse the PDF's text layer. Fails when there is none — the caller should |
| 49 | + /// fall back to the scanned pipeline, exactly as the demo page does. |
| 50 | + #[wasm_bindgen(constructor)] |
| 51 | + pub fn new(bytes: &[u8]) -> Result<DigitalConverter, JsError> { |
| 52 | + let pages = docling_pdf::textparse::pdf_text_pages(bytes); |
| 53 | + if pages.iter().all(|p| p.cells.is_empty()) { |
| 54 | + return Err(JsError::new( |
| 55 | + "PDF has no embedded text layer (scanned/image-only?) — use the OCR pipeline", |
| 56 | + )); |
| 57 | + } |
| 58 | + Ok(Self { |
| 59 | + pages, |
| 60 | + out: Vec::new(), |
| 61 | + }) |
| 62 | + } |
| 63 | + |
| 64 | + /// Pages the text parser found. |
| 65 | + pub fn page_count(&self) -> usize { |
| 66 | + self.pages.len() |
| 67 | + } |
| 68 | + |
| 69 | + /// Convert page `index` (0-based) given its rendered bitmap: layout |
| 70 | + /// detection plus geometric tables. `scale` is the raster's pixels per PDF |
| 71 | + /// point (2.0 for pdf.js `{scale: 2}`). |
| 72 | + pub async fn add_page( |
| 73 | + &mut self, |
| 74 | + index: usize, |
| 75 | + rgba: &[u8], |
| 76 | + px_w: u32, |
| 77 | + px_h: u32, |
| 78 | + scale: f32, |
| 79 | + layout: &LayoutSession, |
| 80 | + ) -> Result<(), JsError> { |
| 81 | + self.add_page_impl(index, rgba, px_w, px_h, scale, layout, None) |
| 82 | + .await |
| 83 | + } |
| 84 | + |
| 85 | + /// [`add_page`](Self::add_page) with TableFormer for the table regions whose |
| 86 | + /// geometric reconstruction looks unreliable. |
| 87 | + #[wasm_bindgen(js_name = addPageTf)] |
| 88 | + #[allow(clippy::too_many_arguments)] |
| 89 | + pub async fn add_page_tf( |
| 90 | + &mut self, |
| 91 | + index: usize, |
| 92 | + rgba: &[u8], |
| 93 | + px_w: u32, |
| 94 | + px_h: u32, |
| 95 | + scale: f32, |
| 96 | + layout: &LayoutSession, |
| 97 | + tf: &TfSession, |
| 98 | + ) -> Result<(), JsError> { |
| 99 | + self.add_page_impl(index, rgba, px_w, px_h, scale, layout, Some(tf)) |
| 100 | + .await |
| 101 | + } |
| 102 | + |
| 103 | + /// Assemble the converted pages into a document and render it as `"md"` |
| 104 | + /// (default), `"json"` or `"doclang"`, with `images` picking how pictures |
| 105 | + /// render in Markdown. Resets the converter. |
| 106 | + pub fn finish( |
| 107 | + &mut self, |
| 108 | + name: &str, |
| 109 | + to: Option<String>, |
| 110 | + images: Option<String>, |
| 111 | + ) -> Result<String, JsError> { |
| 112 | + let doc = finish_document(name, std::mem::take(&mut self.out)); |
| 113 | + render(&doc, to.as_deref(), images.as_deref()) |
| 114 | + } |
| 115 | +} |
| 116 | + |
| 117 | +impl DigitalConverter { |
| 118 | + #[allow(clippy::too_many_arguments)] |
| 119 | + async fn add_page_impl( |
| 120 | + &mut self, |
| 121 | + index: usize, |
| 122 | + rgba: &[u8], |
| 123 | + px_w: u32, |
| 124 | + px_h: u32, |
| 125 | + scale: f32, |
| 126 | + layout: &LayoutSession, |
| 127 | + tf: Option<&TfSession>, |
| 128 | + ) -> Result<(), JsError> { |
| 129 | + if rgba.len() != (px_w as usize) * (px_h as usize) * 4 { |
| 130 | + return Err(JsError::new("rgba buffer size does not match dimensions")); |
| 131 | + } |
| 132 | + let mut page = self |
| 133 | + .pages |
| 134 | + .get(index) |
| 135 | + .cloned() |
| 136 | + .ok_or_else(|| JsError::new("page index is out of range"))?; |
| 137 | + |
| 138 | + let mut img = RgbImage::new(px_w, px_h); |
| 139 | + for (i, px) in img.pixels_mut().enumerate() { |
| 140 | + px.0 = [rgba[i * 4], rgba[i * 4 + 1], rgba[i * 4 + 2]]; |
| 141 | + } |
| 142 | + |
| 143 | + // Layout: Rust preprocessing → JS inference → Rust decoding. The page |
| 144 | + // geometry comes from the text parser, not the bitmap, so a raster at |
| 145 | + // any scale lines up with the cells (which are in PDF points). |
| 146 | + let (page_w, page_h) = (page.width, page.height); |
| 147 | + let input = layout_input(&img); |
| 148 | + let out = layout |
| 149 | + .run(js_sys::Float32Array::from(input.as_slice())) |
| 150 | + .await |
| 151 | + .map_err(|e| JsError::new(&format!("layout session.run: {e:?}")))?; |
| 152 | + let get = |k: &str| { |
| 153 | + js_sys::Reflect::get(&out, &JsValue::from_str(k)) |
| 154 | + .map_err(|_| JsError::new(&format!("layout result has no `{k}`"))) |
| 155 | + }; |
| 156 | + let (logits, q, c) = tensor_parts(&get("logits")?)?; |
| 157 | + let (boxes, bq, four) = tensor_parts(&get("boxes")?)?; |
| 158 | + if bq != q || four != 4 { |
| 159 | + return Err(JsError::new("layout boxes dims must be [1, q, 4]")); |
| 160 | + } |
| 161 | + let regions = decode_layout(&logits, &boxes, q, c, page_w, page_h); |
| 162 | + // Refine against the *real* text cells: unlike the OCR path, orphan-text |
| 163 | + // recovery and the false-picture drop have something to work with here. |
| 164 | + let regions = refine_regions(regions, &page.cells, page_w, page_h); |
| 165 | + |
| 166 | + // Attach the raster so picture regions crop out of it, and record what |
| 167 | + // it is scaled by (cells stay in points). |
| 168 | + page.image = img; |
| 169 | + page.scale = scale; |
| 170 | + |
| 171 | + let mut table_rows = Vec::with_capacity(regions.len()); |
| 172 | + for r in ®ions { |
| 173 | + if r.label != "table" { |
| 174 | + table_rows.push(None); |
| 175 | + continue; |
| 176 | + } |
| 177 | + // Same bargain as the scanned path: TableFormer costs seconds per |
| 178 | + // region, so skip it wherever the free geometric reconstruction is |
| 179 | + // already dense and well-formed. |
| 180 | + let geometric = reconstruct_table(r, &page.cells); |
| 181 | + match tf { |
| 182 | + Some(tf) if !geometric_table_is_reliable(&geometric) => { |
| 183 | + table_rows.push( |
| 184 | + crate::tableformer::predict_table_rows( |
| 185 | + tf, |
| 186 | + &page.image, |
| 187 | + [r.l, r.t, r.r, r.b], |
| 188 | + &page.word_cells, |
| 189 | + ) |
| 190 | + .await, |
| 191 | + ); |
| 192 | + } |
| 193 | + _ => table_rows.push(None), |
| 194 | + } |
| 195 | + } |
| 196 | + |
| 197 | + self.out |
| 198 | + .push(assemble_page_with_tables(&page, regions, table_rows)); |
| 199 | + Ok(()) |
| 200 | + } |
| 201 | +} |
0 commit comments