docling_pdf/
layout.rs

1//! Layout detection via the RT-DETR (`docling-layout-heron`) model exported to
2//! ONNX, run with `ort`. A port of docling-ibm-models' `LayoutPredictor`:
3//! resize the page image to 640×640 and rescale to `[0,1]` (the heron processor
4//! has `do_normalize=false`), run the model, then RT-DETR
5//! `post_process_object_detection` (sigmoid → top-k over query×class →
6//! center-to-corners boxes scaled to the page).
7
8#[cfg(feature = "ml")]
9use image::imageops::FilterType;
10#[cfg(feature = "ml")]
11use image::RgbImage;
12#[cfg(feature = "ml")]
13use ort::session::Session;
14#[cfg(feature = "ml")]
15use ort::value::Tensor;
16
17/// The 17 canonical layout classes, indexed by the model's class id
18/// (`config.json` `id2label`).
19pub const LABELS: [&str; 17] = [
20    "caption",
21    "footnote",
22    "formula",
23    "list_item",
24    "page_footer",
25    "page_header",
26    "picture",
27    "section_header",
28    "table",
29    "text",
30    "title",
31    "document_index",
32    "code",
33    "checkbox_selected",
34    "checkbox_unselected",
35    "form",
36    "key_value_region",
37];
38
39/// One detected region, in page points (top-left origin).
40#[derive(Debug, Clone)]
41pub struct Region {
42    pub label: &'static str,
43    pub score: f32,
44    pub l: f32,
45    pub t: f32,
46    pub r: f32,
47    pub b: f32,
48}
49
50/// Base confidence threshold (docling-ibm-models `base_threshold`): the raw
51/// RT-DETR floor before docling's `LayoutPostprocessor` applies its stricter
52/// per-label thresholds ([`label_threshold`]).
53const THRESHOLD: f32 = 0.3;
54/// RT-DETR's fixed square input side.
55pub const SIDE: u32 = 640;
56
57/// Per-label confidence threshold, ported from docling's
58/// `LayoutPostprocessor.CONFIDENCE_THRESHOLDS`. The raw predictor keeps every
59/// detection above the 0.3 base; the postprocessor then drops a cluster whose
60/// score is below its label's threshold. Applying it here (equivalent, since
61/// every per-label threshold is ≥ the 0.3 base) keeps low-confidence pictures /
62/// tables / list-items out of the assembly, matching docling.
63pub fn label_threshold(label: &str) -> f32 {
64    match label {
65        "section_header"
66        | "title"
67        | "code"
68        | "checkbox_selected"
69        | "checkbox_unselected"
70        | "form"
71        | "key_value_region"
72        | "document_index" => 0.45,
73        // caption, footnote, formula, list_item, page_footer, page_header,
74        // picture, table, text — all 0.5 in docling.
75        _ => 0.5,
76    }
77}
78
79#[cfg(feature = "ml")]
80pub struct LayoutModel {
81    session: Session,
82    /// Set when a multi-page inference fails — e.g. a locally built pre-#73
83    /// static graph (fixed batch=1) via `DOCLING_LAYOUT_ONNX` or a stale
84    /// `layout_heron_int8.onnx`. Batched calls then fall back to per-page runs
85    /// instead of failing the conversion.
86    batch_unsupported: bool,
87    /// The fp32 graph to escalate a suspicious page to, set only when the
88    /// *auto-selected* int8 graph loaded (an explicit `DOCLING_LAYOUT_ONNX` /
89    /// `DOCLING_RS_FP32` choice is respected). Int8 confidences sit close
90    /// enough to the 0.5 label thresholds that a different CPU's quantized
91    /// kernels (AVX-VNNI vs AVX2, CUDA's fallback mix) can flip a whole page's
92    /// detections — observed as a bill page whose tables all dissolved into
93    /// orphan lines on one machine while converting perfectly on another.
94    fp32_path: Option<String>,
95    /// Lazily-loaded session over `fp32_path` — most documents never pay for it.
96    fp32: Option<Session>,
97    /// Intra-op threads, kept for the lazy fp32 load.
98    intra: usize,
99}
100
101#[cfg(feature = "ml")]
102impl LayoutModel {
103    /// Load the ONNX model from `DOCLING_LAYOUT_ONNX`. Without the override,
104    /// prefers `models/layout_heron_int8.onnx` when present (the quantized
105    /// default; `DOCLING_RS_FP32=1` opts out), else `models/layout_heron.onnx`.
106    pub fn load() -> Result<Self, String> {
107        Self::load_with(crate::intra_threads())
108    }
109
110    /// Like [`load`](Self::load) but with an explicit intra-op thread count. A
111    /// parallel page-worker pool loads its helper models on a single thread each
112    /// and gets its speed-up from running pages concurrently instead.
113    pub fn load_with(intra: usize) -> Result<Self, String> {
114        let path = crate::model_path(
115            "DOCLING_LAYOUT_ONNX",
116            "models/layout_heron.onnx",
117            "models/layout_heron_int8.onnx",
118        );
119        if crate::timing::enabled() {
120            eprintln!("docling-pdf: layout model: {path}");
121        }
122        // Escalation target for the quant-robustness guard: only when the
123        // int8 graph was picked automatically and the fp32 one is also there.
124        let fp32_path = if std::env::var("DOCLING_LAYOUT_ONNX").is_err() {
125            let fp32 = crate::resolve_asset("models/layout_heron.onnx");
126            (path != fp32 && std::path::Path::new(&fp32).exists()).then_some(fp32)
127        } else {
128            None
129        };
130        let session = Self::open_session(&path, intra)?;
131        Ok(Self {
132            session,
133            batch_unsupported: false,
134            fp32_path,
135            fp32: None,
136            intra,
137        })
138    }
139
140    fn open_session(path: &str, intra: usize) -> Result<Session, String> {
141        let builder = Session::builder()
142            .map_err(|e| format!("layout: builder: {e}"))?
143            // Let inference use the available cores (ort otherwise defaults low);
144            // a large PDF runs this model once per page.
145            .with_intra_threads(intra)
146            .map_err(|e| format!("layout: intra_threads: {e}"))?;
147        crate::ep::apply(builder)
148            .map_err(|e| format!("layout: {e}"))?
149            .commit_from_file(path)
150            .map_err(|e| format!("layout: load {path}: {e}"))
151    }
152
153    /// Re-run one page through the fp32 graph — the escape hatch for a page
154    /// whose int8 detections look implausible (see `fp32_path`). `Ok(None)`
155    /// when there is nothing to escalate to: fp32 already loaded, an explicit
156    /// model override, or no fp32 file on disk.
157    pub fn predict_fp32_fallback(
158        &mut self,
159        img: &RgbImage,
160        page_w: f32,
161        page_h: f32,
162    ) -> Result<Option<Vec<Region>>, String> {
163        let Some(path) = self.fp32_path.clone() else {
164            return Ok(None);
165        };
166        if self.fp32.is_none() {
167            if crate::timing::enabled() {
168                eprintln!("docling-pdf: loading fp32 layout fallback: {path}");
169            }
170            self.fp32 = Some(Self::open_session(&path, self.intra)?);
171        }
172        let session = self.fp32.as_mut().expect("just loaded");
173        Ok(Some(
174            Self::run_on(session, &[(img, page_w, page_h)])?
175                .pop()
176                .expect("one result per input page"),
177        ))
178    }
179
180    /// Detect layout regions on a page image. `page_w`/`page_h` are the page size
181    /// in points; returned boxes are in those coordinates.
182    pub fn predict(
183        &mut self,
184        img: &RgbImage,
185        page_w: f32,
186        page_h: f32,
187    ) -> Result<Vec<Region>, String> {
188        Ok(self
189            .predict_batch(&[(img, page_w, page_h)])?
190            .pop()
191            .expect("one result per input page"))
192    }
193
194    /// Detect layout regions on several page images with **one** inference call
195    /// (issue #73). The ONNX export has a dynamic batch dimension, so a worker
196    /// can amortize the per-run framework overhead and keep its cores busier on
197    /// multi-page documents. Results are per-image, index-aligned with `pages`,
198    /// and identical to calling [`predict`](Self::predict) per page.
199    pub fn predict_batch(
200        &mut self,
201        pages: &[(&RgbImage, f32, f32)],
202    ) -> Result<Vec<Vec<Region>>, String> {
203        if pages.len() > 1 && self.batch_unsupported {
204            return self.predict_singly(pages);
205        }
206        match self.run_batch(pages) {
207            Err(e) if pages.len() > 1 => {
208                // A graph without the dynamic batch dim (pre-#73 export) fails
209                // only for batch > 1 — remember and recover per page. Warn once
210                // per process, not per worker: every worker owns a LayoutModel
211                // over the same graph file, so repeats carry no information.
212                static WARNED: std::sync::atomic::AtomicBool =
213                    std::sync::atomic::AtomicBool::new(false);
214                if !WARNED.swap(true, std::sync::atomic::Ordering::Relaxed) {
215                    eprintln!(
216                        "docling-pdf: layout model rejected a {}-page batch ({e}); \
217                         falling back to per-page inference — re-export with \
218                         scripts/install/export_layout.py for batched layout",
219                        pages.len()
220                    );
221                }
222                self.batch_unsupported = true;
223                self.predict_singly(pages)
224            }
225            other => other,
226        }
227    }
228
229    fn predict_singly(
230        &mut self,
231        pages: &[(&RgbImage, f32, f32)],
232    ) -> Result<Vec<Vec<Region>>, String> {
233        pages
234            .iter()
235            .map(|p| Ok(self.run_batch(&[*p])?.pop().expect("one result")))
236            .collect()
237    }
238
239    fn run_batch(&mut self, pages: &[(&RgbImage, f32, f32)]) -> Result<Vec<Vec<Region>>, String> {
240        Self::run_on(&mut self.session, pages)
241    }
242
243    fn run_on(
244        session: &mut Session,
245        pages: &[(&RgbImage, f32, f32)],
246    ) -> Result<Vec<Vec<Region>>, String> {
247        if pages.is_empty() {
248            return Ok(Vec::new());
249        }
250        // Resize each page to 640×640 (RT-DETR ignores aspect ratio), rescale to
251        // [0,1], lay out as NCHW.
252        let n = (SIDE * SIDE) as usize;
253        let batch = pages.len();
254        let mut data = vec![0f32; batch * 3 * n];
255        for (p, (img, _, _)) in pages.iter().enumerate() {
256            let resized = image::imageops::resize(*img, SIDE, SIDE, FilterType::Triangle);
257            let page_off = p * 3 * n;
258            for (i, px) in resized.pixels().enumerate() {
259                data[page_off + i] = px[0] as f32 / 255.0;
260                data[page_off + n + i] = px[1] as f32 / 255.0;
261                data[page_off + 2 * n + i] = px[2] as f32 / 255.0;
262            }
263        }
264        let input = Tensor::from_array(([batch, 3, SIDE as usize, SIDE as usize], data))
265            .map_err(|e| format!("layout: input tensor: {e}"))?;
266        let outputs = session
267            .run(ort::inputs!["pixel_values" => input])
268            .map_err(|e| format!("layout: inference: {e}"))?;
269        let (lshape, logits) = outputs["logits"]
270            .try_extract_tensor::<f32>()
271            .map_err(|e| format!("layout: extract logits: {e}"))?;
272        let (_, boxes) = outputs["pred_boxes"]
273            .try_extract_tensor::<f32>()
274            .map_err(|e| format!("layout: extract boxes: {e}"))?;
275
276        let num_queries = lshape[1] as usize;
277        let num_classes = lshape[2] as usize;
278
279        let mut all = Vec::with_capacity(batch);
280        for (p, (_, page_w, page_h)) in pages.iter().enumerate() {
281            let logits =
282                &logits[p * num_queries * num_classes..(p + 1) * num_queries * num_classes];
283            let boxes = &boxes[p * num_queries * 4..(p + 1) * num_queries * 4];
284            all.push(decode_layout(
285                logits,
286                boxes,
287                num_queries,
288                num_classes,
289                *page_w,
290                *page_h,
291            ));
292        }
293        Ok(all)
294    }
295}
296
297fn sigmoid(x: f32) -> f32 {
298    1.0 / (1.0 + (-x).exp())
299}
300
301/// Pack one page image into the model's `(1, 3, SIDE, SIDE)` input: resize
302/// (aspect ignored, RT-DETR convention), rescale to `[0,1]`, CHW. Shared
303/// with the browser build (#157), which delegates only the session call.
304#[cfg(feature = "ocr-prep")]
305pub fn layout_input(img: &image::RgbImage) -> Vec<f32> {
306    let n = (SIDE * SIDE) as usize;
307    let mut data = vec![0f32; 3 * n];
308    let resized = image::imageops::resize(img, SIDE, SIDE, image::imageops::FilterType::Triangle);
309    for (i, px) in resized.pixels().enumerate() {
310        data[i] = px[0] as f32 / 255.0;
311        data[n + i] = px[1] as f32 / 255.0;
312        data[2 * n + i] = px[2] as f32 / 255.0;
313    }
314    data
315}
316
317/// Decode one page's raw RT-DETR outputs into scored [`Region`]s in page
318/// points — sigmoid over every (query, class), top-`num_queries` kept, boxes
319/// converted center→corners and scaled. Shared with the browser build; the
320/// native batch path calls it per page, so both decode identically.
321pub fn decode_layout(
322    logits: &[f32],
323    boxes: &[f32],
324    num_queries: usize,
325    num_classes: usize,
326    page_w: f32,
327    page_h: f32,
328) -> Vec<Region> {
329    let mut scored: Vec<(f32, usize)> = (0..num_queries * num_classes)
330        .map(|idx| (sigmoid(logits[idx]), idx))
331        .collect();
332    scored.sort_unstable_by(|a, b| b.0.total_cmp(&a.0));
333    scored.truncate(num_queries);
334
335    let mut regions = Vec::new();
336    for (score, idx) in scored {
337        if score <= THRESHOLD {
338            continue;
339        }
340        let label_id = idx % num_classes;
341        let q = idx / num_classes;
342        let cx = boxes[q * 4];
343        let cy = boxes[q * 4 + 1];
344        let w = boxes[q * 4 + 2];
345        let h = boxes[q * 4 + 3];
346        // center_to_corners, then scale normalized coords to page points.
347        let l = (cx - w / 2.0) * page_w;
348        let t = (cy - h / 2.0) * page_h;
349        let r = (cx + w / 2.0) * page_w;
350        let b = (cy + h / 2.0) * page_h;
351        regions.push(Region {
352            label: LABELS.get(label_id).copied().unwrap_or("text"),
353            score,
354            l,
355            t,
356            r,
357            b,
358        });
359    }
360    regions
361}
docling_pdf/layout.rs

docling_pdf/
layout.rs
