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#[cfg(feature = "ml")]
51/// Base confidence threshold (docling-ibm-models `base_threshold`): the raw
52/// RT-DETR floor before docling's `LayoutPostprocessor` applies its stricter
53/// per-label thresholds ([`label_threshold`]).
54const THRESHOLD: f32 = 0.3;
55#[cfg(feature = "ml")]
56const SIDE: u32 = 640;
57
58/// Per-label confidence threshold, ported from docling's
59/// `LayoutPostprocessor.CONFIDENCE_THRESHOLDS`. The raw predictor keeps every
60/// detection above the 0.3 base; the postprocessor then drops a cluster whose
61/// score is below its label's threshold. Applying it here (equivalent, since
62/// every per-label threshold is ≥ the 0.3 base) keeps low-confidence pictures /
63/// tables / list-items out of the assembly, matching docling.
64pub fn label_threshold(label: &str) -> f32 {
65    match label {
66        "section_header"
67        | "title"
68        | "code"
69        | "checkbox_selected"
70        | "checkbox_unselected"
71        | "form"
72        | "key_value_region"
73        | "document_index" => 0.45,
74        // caption, footnote, formula, list_item, page_footer, page_header,
75        // picture, table, text — all 0.5 in docling.
76        _ => 0.5,
77    }
78}
79
80#[cfg(feature = "ml")]
81pub struct LayoutModel {
82    session: Session,
83    /// Set when a multi-page inference fails — e.g. a locally built pre-#73
84    /// static graph (fixed batch=1) via `DOCLING_LAYOUT_ONNX` or a stale
85    /// `layout_heron_int8.onnx`. Batched calls then fall back to per-page runs
86    /// instead of failing the conversion.
87    batch_unsupported: bool,
88}
89
90#[cfg(feature = "ml")]
91impl LayoutModel {
92    /// Load the ONNX model from `DOCLING_LAYOUT_ONNX`. Without the override,
93    /// prefers `models/layout_heron_int8.onnx` when present (the quantized
94    /// default; `DOCLING_RS_FP32=1` opts out), else `models/layout_heron.onnx`.
95    pub fn load() -> Result<Self, String> {
96        Self::load_with(crate::intra_threads())
97    }
98
99    /// Like [`load`](Self::load) but with an explicit intra-op thread count. A
100    /// parallel page-worker pool loads its helper models on a single thread each
101    /// and gets its speed-up from running pages concurrently instead.
102    pub fn load_with(intra: usize) -> Result<Self, String> {
103        let path = crate::model_path(
104            "DOCLING_LAYOUT_ONNX",
105            "models/layout_heron.onnx",
106            "models/layout_heron_int8.onnx",
107        );
108        if crate::timing::enabled() {
109            eprintln!("docling-pdf: layout model: {path}");
110        }
111        let builder = Session::builder()
112            .map_err(|e| format!("layout: builder: {e}"))?
113            // Let inference use the available cores (ort otherwise defaults low);
114            // a large PDF runs this model once per page.
115            .with_intra_threads(intra)
116            .map_err(|e| format!("layout: intra_threads: {e}"))?;
117        let session = crate::ep::apply(builder)
118            .map_err(|e| format!("layout: {e}"))?
119            .commit_from_file(&path)
120            .map_err(|e| format!("layout: load {path}: {e}"))?;
121        Ok(Self {
122            session,
123            batch_unsupported: false,
124        })
125    }
126
127    /// Detect layout regions on a page image. `page_w`/`page_h` are the page size
128    /// in points; returned boxes are in those coordinates.
129    pub fn predict(
130        &mut self,
131        img: &RgbImage,
132        page_w: f32,
133        page_h: f32,
134    ) -> Result<Vec<Region>, String> {
135        Ok(self
136            .predict_batch(&[(img, page_w, page_h)])?
137            .pop()
138            .expect("one result per input page"))
139    }
140
141    /// Detect layout regions on several page images with **one** inference call
142    /// (issue #73). The ONNX export has a dynamic batch dimension, so a worker
143    /// can amortize the per-run framework overhead and keep its cores busier on
144    /// multi-page documents. Results are per-image, index-aligned with `pages`,
145    /// and identical to calling [`predict`](Self::predict) per page.
146    pub fn predict_batch(
147        &mut self,
148        pages: &[(&RgbImage, f32, f32)],
149    ) -> Result<Vec<Vec<Region>>, String> {
150        if pages.len() > 1 && self.batch_unsupported {
151            return self.predict_singly(pages);
152        }
153        match self.run_batch(pages) {
154            Err(e) if pages.len() > 1 => {
155                // A graph without the dynamic batch dim (pre-#73 export) fails
156                // only for batch > 1 — remember and recover per page. Warn once
157                // per process, not per worker: every worker owns a LayoutModel
158                // over the same graph file, so repeats carry no information.
159                static WARNED: std::sync::atomic::AtomicBool =
160                    std::sync::atomic::AtomicBool::new(false);
161                if !WARNED.swap(true, std::sync::atomic::Ordering::Relaxed) {
162                    eprintln!(
163                        "docling-pdf: layout model rejected a {}-page batch ({e}); \
164                         falling back to per-page inference — re-export with \
165                         scripts/install/export_layout.py for batched layout",
166                        pages.len()
167                    );
168                }
169                self.batch_unsupported = true;
170                self.predict_singly(pages)
171            }
172            other => other,
173        }
174    }
175
176    fn predict_singly(
177        &mut self,
178        pages: &[(&RgbImage, f32, f32)],
179    ) -> Result<Vec<Vec<Region>>, String> {
180        pages
181            .iter()
182            .map(|p| Ok(self.run_batch(&[*p])?.pop().expect("one result")))
183            .collect()
184    }
185
186    fn run_batch(&mut self, pages: &[(&RgbImage, f32, f32)]) -> Result<Vec<Vec<Region>>, String> {
187        if pages.is_empty() {
188            return Ok(Vec::new());
189        }
190        // Resize each page to 640×640 (RT-DETR ignores aspect ratio), rescale to
191        // [0,1], lay out as NCHW.
192        let n = (SIDE * SIDE) as usize;
193        let batch = pages.len();
194        let mut data = vec![0f32; batch * 3 * n];
195        for (p, (img, _, _)) in pages.iter().enumerate() {
196            let resized = image::imageops::resize(*img, SIDE, SIDE, FilterType::Triangle);
197            let page_off = p * 3 * n;
198            for (i, px) in resized.pixels().enumerate() {
199                data[page_off + i] = px[0] as f32 / 255.0;
200                data[page_off + n + i] = px[1] as f32 / 255.0;
201                data[page_off + 2 * n + i] = px[2] as f32 / 255.0;
202            }
203        }
204        let input = Tensor::from_array(([batch, 3, SIDE as usize, SIDE as usize], data))
205            .map_err(|e| format!("layout: input tensor: {e}"))?;
206        let outputs = self
207            .session
208            .run(ort::inputs!["pixel_values" => input])
209            .map_err(|e| format!("layout: inference: {e}"))?;
210        let (lshape, logits) = outputs["logits"]
211            .try_extract_tensor::<f32>()
212            .map_err(|e| format!("layout: extract logits: {e}"))?;
213        let (_, boxes) = outputs["pred_boxes"]
214            .try_extract_tensor::<f32>()
215            .map_err(|e| format!("layout: extract boxes: {e}"))?;
216
217        let num_queries = lshape[1] as usize;
218        let num_classes = lshape[2] as usize;
219
220        let mut all = Vec::with_capacity(batch);
221        for (p, (_, page_w, page_h)) in pages.iter().enumerate() {
222            let logits =
223                &logits[p * num_queries * num_classes..(p + 1) * num_queries * num_classes];
224            let boxes = &boxes[p * num_queries * 4..(p + 1) * num_queries * 4];
225
226            // sigmoid over every (query, class); take the top `num_queries` scores.
227            let mut scored: Vec<(f32, usize)> = (0..num_queries * num_classes)
228                .map(|idx| (sigmoid(logits[idx]), idx))
229                .collect();
230            scored.sort_unstable_by(|a, b| b.0.total_cmp(&a.0));
231            scored.truncate(num_queries);
232
233            let mut regions = Vec::new();
234            for (score, idx) in scored {
235                if score <= THRESHOLD {
236                    continue;
237                }
238                let label_id = idx % num_classes;
239                let q = idx / num_classes;
240                let cx = boxes[q * 4];
241                let cy = boxes[q * 4 + 1];
242                let w = boxes[q * 4 + 2];
243                let h = boxes[q * 4 + 3];
244                // center_to_corners, then scale normalized coords to page points.
245                let l = (cx - w / 2.0) * page_w;
246                let t = (cy - h / 2.0) * page_h;
247                let r = (cx + w / 2.0) * page_w;
248                let b = (cy + h / 2.0) * page_h;
249                regions.push(Region {
250                    label: LABELS.get(label_id).copied().unwrap_or("text"),
251                    score,
252                    l,
253                    t,
254                    r,
255                    b,
256                });
257            }
258            all.push(regions);
259        }
260        Ok(all)
261    }
262}
263
264#[cfg(feature = "ml")]
265fn sigmoid(x: f32) -> f32 {
266    1.0 / (1.0 + (-x).exp())
267}
docling_pdf/layout.rs

docling_pdf/
layout.rs
