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