Skip to main content

docling_pdf/
ocr_det.rs

1//! PP-OCR text *detection* for bitmap pages (#429): the DB (Differentiable
2//! Binarization) detector RapidOCR runs in front of its recognizer, ported so
3//! text the layout model gives no region — a diagram's labels, a chart's
4//! ticks, a stamp, a page number in the margin — is still read.
5//!
6//! The pipeline's OCR has always been *recognition-only*: PP-OCRv3 rec runs
7//! on the lines inside layout regions, and nothing else on a page is looked
8//! at. docling's engines (RapidOCR, EasyOCR, Tesseract) instead detect text
9//! lines over the whole bitmap and every detected line becomes a cell — the
10//! ones no layout cluster claims turn into orphan text clusters (confidence
11//! 1.0) and read out in the flow. `ModalNet-19.png` shows the gap: the layout
12//! model scores only `SoftMax` above its 0.5 threshold (docling's own layout
13//! run yields *zero* clusters there), yet docling reads `MatMul`, `SoftMax`,
14//! `Mask (opt.)`, `Scale`, `Q`, `K` — all from its detector.
15//!
16//! Model: RapidOCR's `PP-OCRv6_det_small.onnx` (the one docling 2.127's
17//! RapidOCR default resolves to; `.models/ocr_det.onnx`, `DOCLING_OCR_DET_ONNX`
18//! overrides). Pre/post-processing follow `rapidocr/ch_ppocr_det`
19//! (`DetPreProcess` / `DBPostProcess`) with RapidOCR's config: shorter side
20//! scaled up to 736 (`limit_type: min`), sides rounded to multiples of 32,
21//! BGR channel order (cv2 input), `(x/255 − 0.5)/0.5`; probability threshold
22//! 0.3, 2×2 dilation, per-blob minimum-area rectangle, `box_score_fast` ≥ 0.5,
23//! unclip ratio 1.6, RapidOCR's row-then-column box order. Two deliberate
24//! simplifications: only outer blob boundaries are considered (OpenCV's
25//! `RETR_LIST` also walks hole contours, whose boxes fail the score gate
26//! anyway), and a detected quad is handed on as its axis-aligned bounding
27//! box — the recognizer's line prep crops rectangles, and rotated lines are
28//! not what documents lose today.
29//!
30//! The detector is a *supplement*: region-scoped recognition stays the source
31//! for text inside layout regions (so every existing snapshot of a scanned
32//! page keeps its lines), and only detected boxes not already covered by a
33//! recognized cell are cropped and recognized. Missing model → no detection,
34//! quietly (`DOCLING_RS_DEBUG` reports it), so an install without it behaves
35//! exactly as before.
36
37use image::RgbImage;
38
39/// RapidOCR `Det.limit_side_len` with `limit_type: min`.
40pub const LIMIT_SIDE_LEN: u32 = 736;
41/// `DBPostProcess(thresh=…)`: probability → binary mask.
42pub const THRESH: f32 = 0.3;
43/// `DBPostProcess(box_thresh=…)`: minimum mean probability inside a box.
44pub const BOX_THRESH: f32 = 0.5;
45/// `DBPostProcess(unclip_ratio=…)`.
46pub const UNCLIP_RATIO: f32 = 1.6;
47/// `DBPostProcess.min_size`: the shortest side a blob's rectangle may have.
48const MIN_SIZE: f32 = 3.0;
49/// `TextDetector._BOX_SORT_Y_THRESHOLD`: boxes whose top-left `y` differ by
50/// less than this share a row when ordering.
51const BOX_SORT_Y_THRESHOLD: f32 = 10.0;
52
53/// One detected text line: its axis-aligned box in *input image pixels* and
54/// the DB score (mean probability inside the pre-unclip rectangle).
55#[derive(Debug, Clone, Copy, PartialEq)]
56pub struct DetBox {
57    pub l: f32,
58    pub t: f32,
59    pub r: f32,
60    pub b: f32,
61    pub score: f32,
62}
63
64/// `DetPreProcess.resize`: the network input size for a `w × h` image —
65/// shorter side scaled up to [`LIMIT_SIDE_LEN`] (never down), both sides
66/// rounded to a multiple of 32. `None` when a side rounds to zero.
67pub fn det_input_size(w: u32, h: u32) -> Option<(u32, u32)> {
68    det_input_size_capped(w, h, max_side_cap())
69}
70
71/// Default cap on the detector input's longer side: PaddleOCR's own
72/// `det_limit_side_len` (`limit_type: max`). RapidOCR — and so docling —
73/// runs the uncapped shorter-side rule instead; measured on the snapshot
74/// corpus the cap cuts detection to about a third of its time (a Letter page
75/// at the 2.0 px/pt render goes 1216 × 1600 → 736 × 960) and moves bitmap
76/// outputs only by noise-level amounts in both directions, so speed wins by
77/// default and `DOCLING_RS_OCR_DET_MAX_SIDE=0` restores RapidOCR's input.
78pub const DEFAULT_MAX_SIDE: u32 = 960;
79
80/// `DOCLING_RS_OCR_DET_MAX_SIDE`: the cap on the detector input's longer
81/// side — [`DEFAULT_MAX_SIDE`] unless set, `0` = uncapped (RapidOCR's rule).
82/// The DB net is the costliest OCR stage on a scanned page and its cost is
83/// linear in input pixels; a tighter cap trades small-print recall for time.
84fn max_side_cap() -> u32 {
85    static CAP: std::sync::OnceLock<u32> = std::sync::OnceLock::new();
86    *CAP.get_or_init(|| {
87        docling_core::env::parse::<u32>("DOCLING_RS_OCR_DET_MAX_SIDE").unwrap_or(DEFAULT_MAX_SIDE)
88    })
89}
90
91/// [`det_input_size`] with an explicit longer-side cap (`0` = none): the cap
92/// scales the image down first, then RapidOCR's shorter-side rule applies to
93/// what is left (so a capped input never exceeds the cap).
94pub fn det_input_size_capped(w: u32, h: u32, max_side: u32) -> Option<(u32, u32)> {
95    if w == 0 || h == 0 {
96        return None;
97    }
98    let (w, h) = (w as f32, h as f32);
99    // RapidOCR's rule first (shorter side up to 736, never down), then the
100    // cap pulls the longer side back if that overshoots it.
101    let mut ratio = if w.min(h) < LIMIT_SIDE_LEN as f32 {
102        LIMIT_SIDE_LEN as f32 / w.min(h)
103    } else {
104        1.0
105    };
106    if max_side > 0 && w.max(h) * ratio > max_side as f32 {
107        ratio = max_side as f32 / w.max(h);
108    }
109    let round32 = |v: f32| ((v as i64 as f32 / 32.0).round() * 32.0) as i64;
110    let (rw, rh) = (round32(w * ratio), round32(h * ratio));
111    (rw > 0 && rh > 0).then_some((rw as u32, rh as u32))
112}
113
114/// The NCHW float input for the detector: the image resized to
115/// [`det_input_size`] (bilinear, cv2's default), channels in **BGR** order
116/// (RapidOCR feeds a cv2 image), normalized `(x/255 − 0.5)/0.5`. Returns the
117/// tensor and its `(width, height)`.
118pub fn prep_det_input(img: &RgbImage) -> Option<(Vec<f32>, u32, u32)> {
119    let (w, h) = det_input_size(img.width(), img.height())?;
120    let resized = if (w, h) == img.dimensions() {
121        img.clone()
122    } else {
123        resize_bilinear(img, w, h)
124    };
125    let n = (w * h) as usize;
126    let mut data = vec![0f32; 3 * n];
127    for (i, px) in resized.pixels().enumerate() {
128        // B, G, R planes.
129        data[i] = px[2] as f32 / 127.5 - 1.0;
130        data[n + i] = px[1] as f32 / 127.5 - 1.0;
131        data[2 * n + i] = px[0] as f32 / 127.5 - 1.0;
132    }
133    Some((data, w, h))
134}
135
136/// Bilinear resize (cv2's default `INTER_LINEAR`) — `fast_image_resize`'s SIMD
137/// convolution with the triangle kernel, the scalar `image` crate resize with
138/// `DOCLING_RS_SLOW_RESIZE=1` (same kernel, several times slower; the
139/// scalar path is also the fallback should the SIMD one refuse the buffer).
140fn resize_bilinear(img: &RgbImage, w: u32, h: u32) -> RgbImage {
141    #[cfg(feature = "ml")]
142    {
143        use fast_image_resize as fir;
144        static SLOW: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
145        let slow = *SLOW.get_or_init(|| docling_core::env::flag("DOCLING_RS_SLOW_RESIZE"));
146        if !slow {
147            let fast = || {
148                let src = fir::images::ImageRef::new(
149                    img.width(),
150                    img.height(),
151                    img.as_raw(),
152                    fir::PixelType::U8x3,
153                )
154                .ok()?;
155                let mut dst = fir::images::Image::new(w, h, fir::PixelType::U8x3);
156                fir::Resizer::new()
157                    .resize(
158                        &src,
159                        &mut dst,
160                        &fir::ResizeOptions::new()
161                            .resize_alg(fir::ResizeAlg::Convolution(fir::FilterType::Bilinear)),
162                    )
163                    .ok()?;
164                RgbImage::from_raw(w, h, dst.into_vec())
165            };
166            if let Some(out) = fast() {
167                return out;
168            }
169        }
170    }
171    image::imageops::resize(img, w, h, image::imageops::FilterType::Triangle)
172}
173
174/// `DBPostProcess.__call__`: text boxes from the detector's `w × h`
175/// probability map, mapped onto a `dest_w × dest_h` image (the one the map was
176/// computed from). Boxes come back in RapidOCR's reading order.
177pub fn db_boxes(prob: &[f32], w: usize, h: usize, dest_w: u32, dest_h: u32) -> Vec<DetBox> {
178    if prob.len() < w * h || w == 0 || h == 0 {
179        return Vec::new();
180    }
181    // Binarize, then dilate with a 2×2 kernel (cv2.dilate, anchor at the
182    // kernel center = its bottom-right cell, so a pixel lights up when it or
183    // its left / upper / upper-left neighbour is set).
184    let seg = |x: usize, y: usize| prob[y * w + x] > THRESH;
185    let mut mask = vec![false; w * h];
186    for y in 0..h {
187        for x in 0..w {
188            mask[y * w + x] = seg(x, y)
189                || (x > 0 && seg(x - 1, y))
190                || (y > 0 && seg(x, y - 1))
191                || (x > 0 && y > 0 && seg(x - 1, y - 1));
192        }
193    }
194    let mut quads: Vec<([(f32, f32); 4], f32)> = Vec::new();
195    let mut seen = vec![false; w * h];
196    let mut stack = Vec::new();
197    let mut component = Vec::new();
198    for start in 0..w * h {
199        if !mask[start] || seen[start] {
200            continue;
201        }
202        // 8-connected blob (cv2.findContours' outer boundary connectivity).
203        component.clear();
204        seen[start] = true;
205        stack.push(start);
206        while let Some(i) = stack.pop() {
207            component.push(i);
208            let (x, y) = (i % w, i / w);
209            for dy in -1i64..=1 {
210                for dx in -1i64..=1 {
211                    let (nx, ny) = (x as i64 + dx, y as i64 + dy);
212                    if nx < 0 || ny < 0 || nx >= w as i64 || ny >= h as i64 {
213                        continue;
214                    }
215                    let j = ny as usize * w + nx as usize;
216                    if mask[j] && !seen[j] {
217                        seen[j] = true;
218                        stack.push(j);
219                    }
220                }
221            }
222        }
223        if quads.len() >= 1000 {
224            // `max_candidates`.
225            break;
226        }
227        // The contour points are pixel centers on the blob's boundary; the
228        // minimum-area rectangle of the boundary is that of the whole blob.
229        let boundary: Vec<(f32, f32)> = component
230            .iter()
231            .copied()
232            .filter(|&i| {
233                let (x, y) = (i % w, i / w);
234                x == 0
235                    || y == 0
236                    || x + 1 == w
237                    || y + 1 == h
238                    || !mask[i - 1]
239                    || !mask[i + 1]
240                    || !mask[i - w]
241                    || !mask[i + w]
242            })
243            .map(|i| ((i % w) as f32, (i / w) as f32))
244            .collect();
245        let Some((corners, sside)) = min_area_rect(&boundary) else {
246            continue;
247        };
248        if sside < MIN_SIZE {
249            continue;
250        }
251        let score = box_score_fast(prob, w, h, &corners);
252        if score < BOX_THRESH {
253            continue;
254        }
255        let Some((expanded, sside)) = unclip(&corners) else {
256            continue;
257        };
258        if sside < MIN_SIZE + 2.0 {
259            continue;
260        }
261        // Into the source image's pixel grid.
262        let mapped: [(f32, f32); 4] = std::array::from_fn(|k| {
263            let (x, y) = expanded[k];
264            (
265                (x / w as f32 * dest_w as f32)
266                    .round()
267                    .clamp(0.0, dest_w as f32),
268                (y / h as f32 * dest_h as f32)
269                    .round()
270                    .clamp(0.0, dest_h as f32),
271            )
272        });
273        quads.push((mapped, score));
274    }
275    // `filter_det_res`: drop boxes whose rectangle is ≤ 3 px on a side.
276    let mut boxes: Vec<DetBox> = quads
277        .into_iter()
278        .filter_map(|(q, score)| {
279            let side =
280                |a: (f32, f32), b: (f32, f32)| ((a.0 - b.0).powi(2) + (a.1 - b.1).powi(2)).sqrt();
281            let (rw, rh) = (side(q[0], q[1]).floor(), side(q[0], q[3]).floor());
282            if rw <= 3.0 || rh <= 3.0 {
283                return None;
284            }
285            let xs = q.iter().map(|p| p.0);
286            let ys = q.iter().map(|p| p.1);
287            Some(DetBox {
288                l: xs.clone().fold(f32::MAX, f32::min),
289                t: ys.clone().fold(f32::MAX, f32::min),
290                r: xs.fold(f32::MIN, f32::max),
291                b: ys.fold(f32::MIN, f32::max),
292                score,
293            })
294        })
295        .collect();
296    sort_boxes(&mut boxes);
297    boxes
298}
299
300/// `TextDetector.sorted_boxes`: by top edge, rows joined while consecutive
301/// tops are closer than [`BOX_SORT_Y_THRESHOLD`], then left to right in a row.
302pub fn sort_boxes(boxes: &mut [DetBox]) {
303    boxes.sort_by(|a, b| a.t.total_cmp(&b.t));
304    let mut row = 0usize;
305    let mut rows = Vec::with_capacity(boxes.len());
306    for i in 0..boxes.len() {
307        if i > 0 && boxes[i].t - boxes[i - 1].t >= BOX_SORT_Y_THRESHOLD {
308            row += 1;
309        }
310        rows.push(row);
311    }
312    let mut order: Vec<usize> = (0..boxes.len()).collect();
313    order.sort_by(|&a, &b| {
314        rows[a]
315            .cmp(&rows[b])
316            .then(boxes[a].l.total_cmp(&boxes[b].l))
317    });
318    let sorted: Vec<DetBox> = order.iter().map(|&i| boxes[i]).collect();
319    boxes.copy_from_slice(&sorted);
320}
321
322/// A candidate enclosing rectangle: its area, corners and shorter side.
323type RectCandidate = (f32, [(f32, f32); 4], f32);
324
325/// `cv2.minAreaRect` over a point set (rotating calipers on the convex hull):
326/// the four corners of the smallest enclosing rectangle and its shorter side,
327/// in the order `get_mini_boxes` returns (top-left, top-right, bottom-right,
328/// bottom-left, for an upright box). `None` for fewer than one point.
329pub fn min_area_rect(points: &[(f32, f32)]) -> Option<([(f32, f32); 4], f32)> {
330    let hull = convex_hull(points);
331    if hull.is_empty() {
332        return None;
333    }
334    if hull.len() <= 2 {
335        // A single pixel or a straight run: a degenerate rectangle along it.
336        let (a, b) = (hull[0], *hull.last().unwrap());
337        return Some((order_corners([a, b, b, a]), 0.0));
338    }
339    let mut best: Option<RectCandidate> = None;
340    for i in 0..hull.len() {
341        let (p, q) = (hull[i], hull[(i + 1) % hull.len()]);
342        let (ex, ey) = (q.0 - p.0, q.1 - p.1);
343        let len = (ex * ex + ey * ey).sqrt();
344        if len < 1e-6 {
345            continue;
346        }
347        let (ux, uy) = (ex / len, ey / len);
348        let (vx, vy) = (-uy, ux);
349        let (mut umin, mut umax, mut vmin, mut vmax) = (f32::MAX, f32::MIN, f32::MAX, f32::MIN);
350        for &(x, y) in &hull {
351            let u = x * ux + y * uy;
352            let v = x * vx + y * vy;
353            umin = umin.min(u);
354            umax = umax.max(u);
355            vmin = vmin.min(v);
356            vmax = vmax.max(v);
357        }
358        let area = (umax - umin) * (vmax - vmin);
359        if best.as_ref().is_none_or(|(a, _, _)| area < *a) {
360            let corner = |u: f32, v: f32| (u * ux + v * vx, u * uy + v * vy);
361            let corners = [
362                corner(umin, vmin),
363                corner(umax, vmin),
364                corner(umax, vmax),
365                corner(umin, vmax),
366            ];
367            best = Some((area, corners, (umax - umin).min(vmax - vmin)));
368        }
369    }
370    best.map(|(_, corners, sside)| (order_corners(corners), sside))
371}
372
373/// `get_mini_boxes`' corner order: sort by x, then the left pair top-first and
374/// the right pair top-first → `[tl, tr, br, bl]`.
375fn order_corners(mut c: [(f32, f32); 4]) -> [(f32, f32); 4] {
376    c.sort_by(|a, b| a.0.total_cmp(&b.0));
377    let (i1, i4) = if c[1].1 > c[0].1 { (0, 1) } else { (1, 0) };
378    let (i2, i3) = if c[3].1 > c[2].1 { (2, 3) } else { (3, 2) };
379    [c[i1], c[i2], c[i3], c[i4]]
380}
381
382/// Andrew's monotone chain; counter-clockwise, no collinear duplicates.
383fn convex_hull(points: &[(f32, f32)]) -> Vec<(f32, f32)> {
384    let mut pts: Vec<(f32, f32)> = points.to_vec();
385    pts.sort_by(|a, b| a.0.total_cmp(&b.0).then(a.1.total_cmp(&b.1)));
386    pts.dedup();
387    if pts.len() < 3 {
388        return pts;
389    }
390    let cross = |o: (f32, f32), a: (f32, f32), b: (f32, f32)| {
391        (a.0 - o.0) * (b.1 - o.1) - (a.1 - o.1) * (b.0 - o.0)
392    };
393    let mut lower: Vec<(f32, f32)> = Vec::new();
394    for &p in &pts {
395        while lower.len() >= 2 && cross(lower[lower.len() - 2], lower[lower.len() - 1], p) <= 0.0 {
396            lower.pop();
397        }
398        lower.push(p);
399    }
400    let mut upper: Vec<(f32, f32)> = Vec::new();
401    for &p in pts.iter().rev() {
402        while upper.len() >= 2 && cross(upper[upper.len() - 2], upper[upper.len() - 1], p) <= 0.0 {
403            upper.pop();
404        }
405        upper.push(p);
406    }
407    lower.pop();
408    upper.pop();
409    lower.extend(upper);
410    lower
411}
412
413/// `box_score_fast`: mean probability over the pixels inside the rectangle
414/// (a convex quad — the pixel-center test replaces `cv2.fillPoly`).
415fn box_score_fast(prob: &[f32], w: usize, h: usize, quad: &[(f32, f32); 4]) -> f32 {
416    let xmin = quad
417        .iter()
418        .map(|p| p.0)
419        .fold(f32::MAX, f32::min)
420        .floor()
421        .clamp(0.0, (w - 1) as f32) as usize;
422    let xmax = quad
423        .iter()
424        .map(|p| p.0)
425        .fold(f32::MIN, f32::max)
426        .ceil()
427        .clamp(0.0, (w - 1) as f32) as usize;
428    let ymin = quad
429        .iter()
430        .map(|p| p.1)
431        .fold(f32::MAX, f32::min)
432        .floor()
433        .clamp(0.0, (h - 1) as f32) as usize;
434    let ymax = quad
435        .iter()
436        .map(|p| p.1)
437        .fold(f32::MIN, f32::max)
438        .ceil()
439        .clamp(0.0, (h - 1) as f32) as usize;
440    let (mut sum, mut n) = (0f64, 0usize);
441    for y in ymin..=ymax {
442        for x in xmin..=xmax {
443            if inside_convex(quad, (x as f32, y as f32)) {
444                sum += prob[y * w + x] as f64;
445                n += 1;
446            }
447        }
448    }
449    if n == 0 {
450        0.0
451    } else {
452        (sum / n as f64) as f32
453    }
454}
455
456/// Point-in-convex-polygon (boundary counts as inside), any winding.
457fn inside_convex(quad: &[(f32, f32); 4], p: (f32, f32)) -> bool {
458    let mut pos = false;
459    let mut neg = false;
460    for i in 0..4 {
461        let (a, b) = (quad[i], quad[(i + 1) % 4]);
462        let cross = (b.0 - a.0) * (p.1 - a.1) - (b.1 - a.1) * (p.0 - a.0);
463        pos |= cross > 1e-6;
464        neg |= cross < -1e-6;
465    }
466    !(pos && neg)
467}
468
469/// `unclip` + `get_mini_boxes`: grow the rectangle outward by
470/// `area · unclip_ratio / perimeter` (the polygon offset of a rectangle is a
471/// rounded rectangle whose minimum-area rectangle is the original grown by
472/// the offset on every side). Returns the corners and the shorter side.
473fn unclip(quad: &[(f32, f32); 4]) -> Option<([(f32, f32); 4], f32)> {
474    let side = |a: (f32, f32), b: (f32, f32)| ((a.0 - b.0).powi(2) + (a.1 - b.1).powi(2)).sqrt();
475    let (wlen, hlen) = (side(quad[0], quad[1]), side(quad[1], quad[2]));
476    let perimeter = 2.0 * (wlen + hlen);
477    if perimeter < 1e-6 {
478        return None;
479    }
480    let d = wlen * hlen * UNCLIP_RATIO / perimeter;
481    let (cx, cy) = (
482        quad.iter().map(|p| p.0).sum::<f32>() / 4.0,
483        quad.iter().map(|p| p.1).sum::<f32>() / 4.0,
484    );
485    // Unit axes of the rectangle.
486    let (ux, uy) = if wlen > 1e-6 {
487        (
488            (quad[1].0 - quad[0].0) / wlen,
489            (quad[1].1 - quad[0].1) / wlen,
490        )
491    } else {
492        (1.0, 0.0)
493    };
494    let (vx, vy) = if hlen > 1e-6 {
495        (
496            (quad[2].0 - quad[1].0) / hlen,
497            (quad[2].1 - quad[1].1) / hlen,
498        )
499    } else {
500        (-uy, ux)
501    };
502    let (hw, hh) = (wlen / 2.0 + d, hlen / 2.0 + d);
503    let corner = |su: f32, sv: f32| {
504        (
505            cx + su * hw * ux + sv * hh * vx,
506            cy + su * hw * uy + sv * hh * vy,
507        )
508    };
509    let corners = order_corners([
510        corner(-1.0, -1.0),
511        corner(1.0, -1.0),
512        corner(1.0, 1.0),
513        corner(-1.0, 1.0),
514    ]);
515    Some((corners, (wlen + 2.0 * d).min(hlen + 2.0 * d)))
516}
517
518/// The detected lines the recognizer still has to read (#429): those not
519/// already covered by the region-scoped pass. `detected` is in image pixels
520/// at `scale` px/pt; `regions` and `cells` are the page's layout regions and
521/// the cells recognized so far, in page points. A line counts as covered when
522/// it lies mostly (> 50 %) inside a text-like or table region — whose lines
523/// the region pass segmented and recognized, whatever it made of them — when
524/// the recognized cells inside it sum to > 30 % of its area, or when a line
525/// accepted earlier in reading order overlaps it by > 30 %. Cumulative on
526/// purpose: DB happily spans two adjacent columns in one box, and judged cell
527/// by cell such a box overlaps every cell a little while re-reading all of
528/// them (the `old_newspaper` scan doubled a paragraph that way); and DB can
529/// emit a word and its whole line. Returns `text` regions in page points for
530/// the recognizer's line prep. Shared by the native worker and the browser
531/// pipeline.
532pub fn uncovered_lines(
533    detected: &[DetBox],
534    scale: f32,
535    regions: &[crate::layout::Region],
536    cells: &[crate::pdfium_backend::TextCell],
537) -> Vec<crate::layout::Region> {
538    let mut accepted: Vec<crate::layout::Region> = Vec::new();
539    for d in detected.iter().map(|d| crate::layout::Region {
540        label: "text",
541        score: d.score,
542        l: d.l / scale,
543        t: d.t / scale,
544        r: d.r / scale,
545        b: d.b / scale,
546    }) {
547        let da = ((d.r - d.l) * (d.b - d.t)).max(1.0);
548        let inter = |l: f32, t: f32, r: f32, b: f32| {
549            (d.r.min(r) - d.l.max(l)).max(0.0) * (d.b.min(b) - d.t.max(t)).max(0.0)
550        };
551        let in_region = regions.iter().any(|r| {
552            (crate::ocr_prep::is_text_label(r.label) || crate::assemble::is_table_like(r.label))
553                && inter(r.l, r.t, r.r, r.b) / da > 0.5
554        });
555        let by_cells: f32 = cells.iter().map(|c| inter(c.l, c.t, c.r, c.b)).sum::<f32>() / da;
556        let by_accepted = accepted
557            .iter()
558            .any(|u| inter(u.l, u.t, u.r, u.b) / da > 0.3);
559        if !in_region && by_cells <= 0.3 && !by_accepted {
560            accepted.push(d);
561        }
562    }
563    accepted
564}
565
566#[cfg(feature = "ml")]
567pub use session::DetModel;
568
569#[cfg(feature = "ml")]
570mod session {
571    use super::{db_boxes, prep_det_input, DetBox};
572    use image::RgbImage;
573    use ort::session::Session;
574    use ort::value::Tensor;
575
576    /// The detector session. Loaded lazily by the page worker alongside the
577    /// recognizer; absent model → the pipeline runs recognition-only.
578    pub struct DetModel {
579        session: Session,
580    }
581
582    /// `DOCLING_OCR_DET_ONNX`, else `.models/ocr_det.onnx` through the asset
583    /// resolver (CWD, `DOCLING_RS_MODELS_DIR`, exe dir).
584    pub(crate) fn resolve_det_path() -> String {
585        docling_core::env::nonempty("DOCLING_OCR_DET_ONNX")
586            .unwrap_or_else(|| crate::resolve_asset(".models/ocr_det.onnx"))
587    }
588
589    impl DetModel {
590        /// Load the detector with `intra` intra-op threads (the worker's layout
591        /// thread budget; DB is a plain conv net, so the output is stable
592        /// across thread counts to the precision a 0.3 threshold sees).
593        pub fn load(intra: usize) -> Result<Self, String> {
594            let path = resolve_det_path();
595            if !std::path::Path::new(&path).exists() {
596                return Err(format!("text detection model not found at {path}"));
597            }
598            let builder = Session::builder()
599                .map_err(|e| format!("ocr-det: builder: {e}"))?
600                .with_intra_threads(intra.max(1))
601                .map_err(|e| format!("ocr-det: intra_threads: {e}"))?;
602            let builder = docling_onnx::apply(builder).map_err(|e| format!("ocr-det: {e}"))?;
603            let session = docling_onnx::commit(builder, &path, "det")
604                .map_err(|e| format!("ocr-det: load {path}: {e}"))?;
605            Ok(Self { session })
606        }
607
608        /// Detect text lines on `img`; boxes in `img` pixels, reading order.
609        pub fn detect(&mut self, img: &RgbImage) -> Result<Vec<DetBox>, String> {
610            let Some((data, w, h)) = crate::timing::timed("ocr.det.prep", || prep_det_input(img))
611            else {
612                return Ok(Vec::new());
613            };
614            let input = Tensor::from_array(([1usize, 3, h as usize, w as usize], data))
615                .map_err(|e| format!("ocr-det: input: {e}"))?;
616            let name = self.session.inputs()[0].name().to_string();
617            let outputs = crate::timing::timed("ocr.det.net", || {
618                self.session
619                    .run(ort::inputs![name.as_str() => input])
620                    .map_err(|e| format!("ocr-det: run: {e}"))
621            })?;
622            let (shape, prob) = outputs[0]
623                .try_extract_tensor::<f32>()
624                .map_err(|e| format!("ocr-det: output: {e}"))?;
625            let dims: Vec<usize> = shape.iter().map(|&d| d as usize).collect();
626            let (ph, pw) = match dims.as_slice() {
627                [_, _, ph, pw] => (*ph, *pw),
628                _ => return Err(format!("ocr-det: unexpected output shape {dims:?}")),
629            };
630            Ok(crate::timing::timed("ocr.det.post", || {
631                db_boxes(prob, pw, ph, img.width(), img.height())
632            }))
633        }
634    }
635}
636
637#[cfg(test)]
638mod tests {
639    use super::*;
640
641    #[test]
642    fn input_size_scales_the_short_side_to_736_in_multiples_of_32() {
643        // RapidOCR's uncapped rule. 445 × 884: short side 445 → ×1.654;
644        // 736 × 1462 → 736 × 1472.
645        assert_eq!(det_input_size_capped(445, 884, 0), Some((736, 1472)));
646        // Already ≥ 736 on the short side: unchanged bar the /32 rounding.
647        assert_eq!(det_input_size_capped(1335, 2652, 0), Some((1344, 2656)));
648        assert_eq!(det_input_size(0, 10), None);
649        // The default cap (960, PaddleOCR's) applies when the env knob is unset.
650        assert_eq!(det_input_size(1224, 1584), Some((736, 960)));
651        // A longer-side cap scales a big page down (1224 × 1584 → 736 × 960
652        // for 960) and leaves a small image's shorter-side upscale alone
653        // (445 × 884 still goes to 736 × 1472 under a 1500 cap, 480 × 960
654        // under 960).
655        assert_eq!(det_input_size_capped(1224, 1584, 960), Some((736, 960)));
656        assert_eq!(det_input_size_capped(445, 884, 1500), Some((736, 1472)));
657        assert_eq!(det_input_size_capped(445, 884, 960), Some((480, 960)));
658        assert_eq!(det_input_size_capped(1224, 1584, 0), Some((1216, 1600)));
659    }
660
661    #[test]
662    fn det_input_is_bgr_normalized() {
663        let mut img = RgbImage::new(736, 736);
664        img.put_pixel(0, 0, image::Rgb([255, 0, 128]));
665        let (data, w, h) = prep_det_input(&img).unwrap();
666        assert_eq!((w, h), (736, 736));
667        let n = (w * h) as usize;
668        // Blue plane first: pixel (0,0) has B=128 → ~0.0039; G=0 → -1; R=255 → 1.
669        assert!((data[0] - (128.0 / 127.5 - 1.0)).abs() < 1e-6);
670        assert_eq!(data[n], -1.0);
671        assert_eq!(data[2 * n], 1.0);
672    }
673
674    #[test]
675    fn min_area_rect_of_an_upright_and_a_tilted_blob() {
676        let pts: Vec<(f32, f32)> = (0..20)
677            .flat_map(|x| (0..5).map(move |y| (x as f32, y as f32)))
678            .collect();
679        let (c, sside) = min_area_rect(&pts).unwrap();
680        assert!((sside - 4.0).abs() < 1e-3);
681        assert!(
682            (c[0].0 - 0.0).abs() < 1e-3 && (c[0].1 - 0.0).abs() < 1e-3,
683            "{c:?}"
684        );
685        assert!(
686            (c[2].0 - 19.0).abs() < 1e-3 && (c[2].1 - 4.0).abs() < 1e-3,
687            "{c:?}"
688        );
689        // The same strip rotated 45°: the tight rectangle is ~19 × 4, not the
690        // axis-aligned 16 × 16 box.
691        let s = std::f32::consts::FRAC_1_SQRT_2;
692        let rot: Vec<(f32, f32)> = pts
693            .iter()
694            .map(|&(x, y)| (x * s - y * s + 50.0, x * s + y * s + 50.0))
695            .collect();
696        let (_, sside) = min_area_rect(&rot).unwrap();
697        assert!((sside - 4.0).abs() < 1e-2, "{sside}");
698    }
699
700    /// Two text-like blobs on a probability map → two boxes, grown by the
701    /// unclip distance (area·1.6/perimeter) on every side, in reading order,
702    /// scaled to the destination image; a faint blob below `box_thresh` and a
703    /// speck below `min_size` are dropped.
704    #[test]
705    fn db_boxes_from_a_synthetic_probability_map() {
706        let (w, h) = (128usize, 64usize);
707        let mut prob = vec![0f32; w * h];
708        let blob = |prob: &mut Vec<f32>, l: usize, t: usize, r: usize, b: usize, p: f32| {
709            for y in t..b {
710                for x in l..r {
711                    prob[y * w + x] = p;
712                }
713            }
714        };
715        blob(&mut prob, 70, 10, 110, 20, 0.9); // right, upper row
716        blob(&mut prob, 10, 12, 50, 22, 0.9); // left, same row (top within 10)
717        blob(&mut prob, 10, 40, 60, 48, 0.35); // above thresh but mean < box_thresh
718        blob(&mut prob, 100, 50, 102, 52, 0.9); // speck
719        let boxes = db_boxes(&prob, w, h, 256, 128);
720        assert_eq!(boxes.len(), 2, "{boxes:?}");
721        // Left blob first (same row, smaller x). Its rectangle spans pixel
722        // centers 10..49 × 12..21 (39 × 9 after dilation shifts by one:
723        // 10..50 × 12..22 → 40 × 10), unclip d = 400·1.6/100 = 6.4.
724        let a = &boxes[0];
725        assert!(a.l < boxes[1].l);
726        // The dilated rectangle takes in one zero-probability rim row and
727        // column, so the mean sits a little under the blob's 0.9 — RapidOCR
728        // scores the dilated contour on the raw map the same way.
729        assert!(a.score > 0.75 && a.score < 0.9, "{}", a.score);
730        // Doubled for the 2× destination scale: l ≈ (10 − 6.4)·2, r ≈ (50 + 6.4)·2.
731        assert!(
732            (a.l - 7.0).abs() <= 2.0 && (a.r - 113.0).abs() <= 2.0,
733            "{a:?}"
734        );
735        assert!(
736            (a.t - 11.0).abs() <= 2.0 && (a.b - 57.0).abs() <= 2.0,
737            "{a:?}"
738        );
739    }
740
741    /// The coverage rule: a line inside a text region, one whose area
742    /// recognized cells mostly fill (cumulatively — two half-covering cells
743    /// count), and a duplicate of an accepted line are all skipped; a line in
744    /// the open is kept, converted to page points.
745    #[test]
746    fn uncovered_lines_skip_what_the_region_pass_read() {
747        use crate::layout::Region;
748        use crate::pdfium_backend::TextCell;
749        let bx = |l: f32, t: f32, r: f32, b: f32| DetBox {
750            l,
751            t,
752            r,
753            b,
754            score: 0.9,
755        };
756        let regions = vec![Region {
757            label: "text",
758            score: 0.9,
759            l: 0.0,
760            t: 0.0,
761            r: 100.0,
762            b: 20.0,
763        }];
764        let cell = |l: f32, r: f32| TextCell {
765            text: "x".into(),
766            l,
767            t: 50.0,
768            r,
769            b: 60.0,
770        };
771        let cells = vec![cell(0.0, 50.0), cell(50.0, 100.0)];
772        let detected = vec![
773            bx(0.0, 0.0, 200.0, 40.0), // inside the text region (page pts 0..100 × 0..20)
774            bx(0.0, 100.0, 200.0, 120.0), // two cells cover it half each → covered
775            bx(0.0, 300.0, 200.0, 320.0), // in the open → kept
776            bx(20.0, 302.0, 100.0, 318.0), // a word of the accepted line → duplicate
777        ];
778        let out = uncovered_lines(&detected, 2.0, &regions, &cells);
779        assert_eq!(out.len(), 1, "{out:?}");
780        assert_eq!(
781            (out[0].l, out[0].t, out[0].r, out[0].b),
782            (0.0, 150.0, 100.0, 160.0)
783        );
784        assert_eq!(out[0].label, "text");
785    }
786
787    #[test]
788    fn boxes_sort_by_row_then_column() {
789        let bx = |l: f32, t: f32| DetBox {
790            l,
791            t,
792            r: l + 10.0,
793            b: t + 10.0,
794            score: 1.0,
795        };
796        let mut boxes = vec![
797            bx(50.0, 100.0),
798            bx(10.0, 105.0),
799            bx(30.0, 20.0),
800            bx(5.0, 200.0),
801        ];
802        sort_boxes(&mut boxes);
803        let order: Vec<(f32, f32)> = boxes.iter().map(|b| (b.l, b.t)).collect();
804        assert_eq!(
805            order,
806            vec![(30.0, 20.0), (10.0, 105.0), (50.0, 100.0), (5.0, 200.0)]
807        );
808    }
809}