|
| 1 | +//! Browser TableFormer — stage 3 of #157. The autoregressive structure loop |
| 2 | +//! and every ONNX-free step (448 preprocessing, tag corrections, bbox |
| 3 | +//! bookkeeping, span merge, OTSL→grid, cell matching) run in Rust via |
| 4 | +//! `docling_pdf::tf_core` — the *same* code the native pipeline uses — and only |
| 5 | +//! the three ONNX graphs (encoder / decoder / bbox) are delegated to ONNX |
| 6 | +//! Runtime Web through the [`TfSession`] interop object. |
| 7 | +//! |
| 8 | +//! The heavy tensors never cross the wasm boundary: [`TfSession`] holds the |
| 9 | +//! encoder's constant cross-attention K/V and `enc_out`, and the growing |
| 10 | +//! decoder KV-cache, entirely on the JS side. Each decode step sends only the |
| 11 | +//! last tag (one int) and gets back `logits` + `hidden` (525 floats). See |
| 12 | +//! `www/scan.html` for the JS wiring (the `decoder_kv` graph: N_LAYERS=6, |
| 13 | +//! KV_HEADS=8, head_dim=64, cross length 784). |
| 14 | +
|
| 15 | +use docling_pdf::pdfium_backend::TextCell; |
| 16 | +use docling_pdf::tf_core::{ |
| 17 | + argmax, build_table_cells, merge_spans, preprocess_input, table_rows, BboxBook, TableCell, |
| 18 | + MAX_STEPS, |
| 19 | +}; |
| 20 | +use image::RgbImage; |
| 21 | +use wasm_bindgen::prelude::*; |
| 22 | + |
| 23 | +#[wasm_bindgen] |
| 24 | +extern "C" { |
| 25 | + /// The JS-side TableFormer session: a stateful wrapper around the three |
| 26 | + /// `ort.InferenceSession`s. `encode` runs the image encoder and stashes the |
| 27 | + /// constant cross tensors + `enc_out` and resets the KV-cache; `step` runs |
| 28 | + /// one decoder step (feeding the stored cross + growing cache) and returns |
| 29 | + /// `{ logits: Float32Array, hidden: Float32Array }`; `bbox` runs the bbox |
| 30 | + /// decoder over the collected per-cell hidden states and returns |
| 31 | + /// `{ boxes: Float32Array, classes: Float32Array }`. |
| 32 | + pub type TfSession; |
| 33 | + |
| 34 | + #[wasm_bindgen(method, catch)] |
| 35 | + async fn encode(this: &TfSession, image: js_sys::Float32Array) -> Result<JsValue, JsValue>; |
| 36 | + |
| 37 | + #[wasm_bindgen(method, catch)] |
| 38 | + async fn step(this: &TfSession, tag: i32) -> Result<JsValue, JsValue>; |
| 39 | + |
| 40 | + #[wasm_bindgen(method, catch)] |
| 41 | + async fn bbox( |
| 42 | + this: &TfSession, |
| 43 | + tag_h: js_sys::Float32Array, |
| 44 | + n: u32, |
| 45 | + ) -> Result<JsValue, JsValue>; |
| 46 | +} |
| 47 | + |
| 48 | +/// Pull a named `Float32Array` out of a JS result object. |
| 49 | +fn get_f32(obj: &JsValue, key: &str) -> Result<Vec<f32>, JsError> { |
| 50 | + let v = js_sys::Reflect::get(obj, &JsValue::from_str(key)) |
| 51 | + .map_err(|_| JsError::new(&format!("tf result has no `{key}`")))?; |
| 52 | + let arr: js_sys::Float32Array = v |
| 53 | + .dyn_into() |
| 54 | + .map_err(|_| JsError::new(&format!("tf `{key}` is not a Float32Array")))?; |
| 55 | + Ok(arr.to_vec()) |
| 56 | +} |
| 57 | + |
| 58 | +/// Run the full structure model on a `SIDE×SIDE`-croppable region image: encode |
| 59 | +/// once, step the decoder to the OTSL sequence (`tf_core::BboxBook` drives the |
| 60 | +/// exact bbox bookkeeping), run the bbox decoder, then merge spans and lay the |
| 61 | +/// cells onto the grid — all shared with native. |
| 62 | +async fn predict_structure( |
| 63 | + session: &TfSession, |
| 64 | + crop: &RgbImage, |
| 65 | +) -> Result<Vec<TableCell>, JsError> { |
| 66 | + let input = preprocess_input(crop); |
| 67 | + session |
| 68 | + .encode(js_sys::Float32Array::from(input.as_slice())) |
| 69 | + .await |
| 70 | + .map_err(|e| JsError::new(&format!("tf encode: {e:?}")))?; |
| 71 | + |
| 72 | + let mut book = BboxBook::new(); |
| 73 | + while book.otsl.len() < MAX_STEPS { |
| 74 | + let last = *book.tags.last().expect("decode starts from <start>"); |
| 75 | + let out = session |
| 76 | + .step(last as i32) |
| 77 | + .await |
| 78 | + .map_err(|e| JsError::new(&format!("tf decode step: {e:?}")))?; |
| 79 | + let logits = get_f32(&out, "logits")?; |
| 80 | + let hidden = get_f32(&out, "hidden")?; |
| 81 | + let raw = argmax(&logits) as i64; |
| 82 | + if !book.step(raw, &hidden) { |
| 83 | + break; |
| 84 | + } |
| 85 | + } |
| 86 | + if book.n == 0 { |
| 87 | + return Ok(Vec::new()); |
| 88 | + } |
| 89 | + |
| 90 | + let out = session |
| 91 | + .bbox( |
| 92 | + js_sys::Float32Array::from(book.hiddens.as_slice()), |
| 93 | + book.n as u32, |
| 94 | + ) |
| 95 | + .await |
| 96 | + .map_err(|e| JsError::new(&format!("tf bbox: {e:?}")))?; |
| 97 | + let boxes: Vec<[f32; 4]> = get_f32(&out, "boxes")? |
| 98 | + .chunks_exact(4) |
| 99 | + .map(|c| [c[0], c[1], c[2], c[3]]) |
| 100 | + .collect(); |
| 101 | + let classes: Vec<i64> = get_f32(&out, "classes")? |
| 102 | + .chunks_exact(3) |
| 103 | + .map(|c| argmax(c) as i64) |
| 104 | + .collect(); |
| 105 | + let (merged, merged_classes) = merge_spans(&boxes, &classes, &book.merge); |
| 106 | + Ok(build_table_cells(&book.otsl, &merged, &merged_classes)) |
| 107 | +} |
| 108 | + |
| 109 | +/// Predict a table region's grid, the browser counterpart of the native |
| 110 | +/// `TableFormer::predict_table_rows`: docling's page→1024px box-average + crop |
| 111 | +/// (`docling_pdf::resample`), the ONNX structure model, then the shared word |
| 112 | +/// matcher (`tf_core::table_rows`). `page_image` is the rendered page (2 px per |
| 113 | +/// point, like the native pipeline); `region` is `(l, t, r, b)` in page points. |
| 114 | +/// `None` when no structure is predicted. |
| 115 | +pub(crate) async fn predict_table_rows( |
| 116 | + session: &TfSession, |
| 117 | + page_image: &RgbImage, |
| 118 | + region: [f32; 4], |
| 119 | + words: &[TextCell], |
| 120 | +) -> Option<Vec<Vec<String>>> { |
| 121 | + // page → 1024px height (cv2.INTER_AREA), then crop the table bbox — docling's |
| 122 | + // coordinate chain with the same rounding the native path reproduces. |
| 123 | + let sf = 1024.0 / page_image.height() as f32; |
| 124 | + let pw = (page_image.width() as f32 * sf) as u32; |
| 125 | + let page1024 = docling_pdf::resample::inter_area(page_image, pw, 1024); |
| 126 | + let k = 2.0 * 1024.0 / page_image.height() as f64; |
| 127 | + let px = |v: f32| (v as f64).round_ties_even() * k; |
| 128 | + let x = (px(region[0]).round_ties_even()).max(0.0) as u32; |
| 129 | + let y = (px(region[1]).round_ties_even()).max(0.0) as u32; |
| 130 | + let x2 = (px(region[2]).round_ties_even() as u32).min(page1024.width()); |
| 131 | + let y2 = (px(region[3]).round_ties_even() as u32).min(page1024.height()); |
| 132 | + if x2 <= x || y2 <= y { |
| 133 | + return None; |
| 134 | + } |
| 135 | + let crop = image::imageops::crop_imm(&page1024, x, y, x2 - x, y2 - y).to_image(); |
| 136 | + let cells = predict_structure(session, &crop).await.ok()?; |
| 137 | + if cells.is_empty() { |
| 138 | + return None; |
| 139 | + } |
| 140 | + table_rows(&cells, region, words) |
| 141 | +} |
0 commit comments