Skip to main content

docling_pdf/
ocr_prep.rs

1//! ONNX-free half of the PP-OCRv3 recognition pipeline: everything before
2//! and after the `session.run` call — line segmentation, crop preparation,
3//! width-batching, CTC decoding, dictionary handling.
4//!
5//! Split out of `ocr.rs` (which keeps the `ort` session) so the browser build
6//! can reuse it (issue #79 phase 2): `docling-wasm` runs these exact
7//! functions and delegates only the inference call to ONNX Runtime Web on
8//! the JS side. Keeping one implementation is what makes the wasm output
9//! byte-comparable to the native CPU path — any drift then comes from the
10//! runtime, not from pre/post-processing.
11
12use image::{imageops, imageops::FilterType, Rgb, RgbImage};
13
14/// PP-OCRv3's fixed input height.
15pub const REC_HEIGHT: u32 = 48;
16
17/// Cap on lines per recognition run: bounds peak input-tensor memory
18/// (16 × 3 × 48 × 2400 px ≈ 22 MB f32) without costing measurable batching
19/// benefit — same-width groups are rarely larger.
20pub const REC_BATCH: usize = 16;
21
22/// A text-line crop prepared for recognition: resized to the fixed model
23/// height, normalised to `[-1, 1]`, laid out CHW.
24pub struct PrepLine {
25    /// Width after the aspect-preserving resize to [`REC_HEIGHT`].
26    pub w: usize,
27    /// `3 * REC_HEIGHT * w` values.
28    pub data: Vec<f32>,
29}
30
31/// Prepare one line crop, or `None` for a degenerate (zero-sized) crop.
32pub fn prep_line(line: &RgbImage) -> Option<PrepLine> {
33    let (w, h) = line.dimensions();
34    if w == 0 || h == 0 {
35        return None;
36    }
37    let new_w = ((w as f32) * REC_HEIGHT as f32 / h as f32)
38        .round()
39        .clamp(8.0, 2400.0) as u32;
40    let resized = imageops::resize(line, new_w, REC_HEIGHT, FilterType::Triangle);
41    let n = (REC_HEIGHT * new_w) as usize;
42    // Normalise to [-1, 1]: (x/255 - 0.5) / 0.5.
43    let mut data = vec![0f32; 3 * n];
44    for (i, px) in resized.pixels().enumerate() {
45        data[i] = px[0] as f32 / 127.5 - 1.0;
46        data[n + i] = px[1] as f32 / 127.5 - 1.0;
47        data[2 * n + i] = px[2] as f32 / 127.5 - 1.0;
48    }
49    Some(PrepLine {
50        w: new_w as usize,
51        data,
52    })
53}
54
55/// The CTC class table for a recognition dictionary file: index 0 = blank,
56/// then one class per dictionary line, then the space class.
57pub fn dict_chars(dict: &str) -> Vec<String> {
58    let mut chars = vec![String::new()]; // blank at 0
59    chars.extend(dict.lines().map(|s| s.to_string()));
60    chars.push(" ".to_string());
61    chars
62}
63
64/// Greedy CTC decode of one row's `(T, C)` probabilities.
65pub fn decode_row(chars: &[String], probs: &[f32], nc: usize) -> String {
66    decode_row_scored(chars, probs, nc).0
67}
68
69/// [`decode_row`] plus the line's recognition confidence: the mean probability
70/// of the emitted characters (PaddleOCR's per-line score — what docling's
71/// `TextCell.confidence` carries for OCR cells, feeding the `ocr_score`
72/// confidence aggregate, #183). `0.0` when nothing decodes.
73pub fn decode_row_scored(chars: &[String], probs: &[f32], nc: usize) -> (String, f32) {
74    let mut out = String::new();
75    let mut prev = 0usize;
76    let mut conf_sum = 0.0f32;
77    let mut conf_n = 0usize;
78    for row in probs.chunks_exact(nc) {
79        let mut best = 0usize;
80        let mut bestv = row[0];
81        for (c, &v) in row.iter().enumerate().skip(1) {
82            if v > bestv {
83                bestv = v;
84                best = c;
85            }
86        }
87        if best != prev && best != 0 {
88            if let Some(ch) = chars.get(best) {
89                out.push_str(ch);
90                conf_sum += bestv;
91                conf_n += 1;
92            }
93        }
94        prev = best;
95    }
96    let conf = if conf_n == 0 {
97        0.0
98    } else {
99        conf_sum / conf_n as f32
100    };
101    (out, conf)
102}
103
104pub(crate) fn luma(p: &Rgb<u8>) -> f32 {
105    0.299 * p[0] as f32 + 0.587 * p[1] as f32 + 0.114 * p[2] as f32
106}
107
108/// Split a region crop into text lines via a horizontal ink-projection profile.
109/// Returns tight `(l, t, r, b)` boxes in crop pixels.
110pub fn segment_lines(crop: &RgbImage) -> Vec<(u32, u32, u32, u32)> {
111    let (w, h) = crop.dimensions();
112    if w == 0 || h == 0 {
113        return Vec::new();
114    }
115    let mean: f32 = crop.pixels().map(luma).sum::<f32>() / (w * h) as f32;
116    let thresh = mean * 0.7; // ink = noticeably darker than the page average
117    let min_ink = ((w as f32) * 0.005).max(1.0) as u32;
118
119    // A column inked in nearly every row is a vertical rule — a panel border, a
120    // table line — not text; left in the profile it bridges every inter-line
121    // gap and the whole crop collapses into one giant "line" (a framed
122    // terms-and-conditions box OCR'd as a single unreadable strip). Mask those
123    // columns out of the row profile; glyph columns never come close (a
124    // descender-to-ascender stack is still far from 90 % of the crop height).
125    let mut col_ink = vec![0u32; w as usize];
126    for y in 0..h {
127        for x in 0..w {
128            if luma(crop.get_pixel(x, y)) < thresh {
129                col_ink[x as usize] += 1;
130            }
131        }
132    }
133    // Borders are *thin*: when "always-inked" columns make up a noticeable
134    // share of the width it is not a frame but a polarity/threshold artifact
135    // (an inverted dark-mode page classifies its whole background as ink) —
136    // leave the profile alone and let polarity normalization handle it.
137    let rule_cols = col_ink
138        .iter()
139        .filter(|&&c| c as f32 > 0.9 * h as f32)
140        .count();
141    let mask_rules = (rule_cols as f32) < 0.15 * w as f32;
142    let rule = |x: u32| mask_rules && col_ink[x as usize] as f32 > 0.9 * h as f32;
143
144    let mut profile = vec![0u32; h as usize];
145    for y in 0..h {
146        let mut row = 0u32;
147        for x in 0..w {
148            if !rule(x) && luma(crop.get_pixel(x, y)) < thresh {
149                row += 1;
150            }
151        }
152        profile[y as usize] = row;
153    }
154
155    // Maximal runs of text rows, separated by (near-)blank rows.
156    let mut runs: Vec<(u32, u32)> = Vec::new();
157    let mut start: Option<u32> = None;
158    for y in 0..h {
159        let text = profile[y as usize] >= min_ink;
160        if text && start.is_none() {
161            start = Some(y);
162        } else if !text {
163            if let Some(s) = start.take() {
164                if y - s >= 4 {
165                    runs.push((s, y));
166                }
167            }
168        }
169    }
170    if let Some(s) = start {
171        if h - s >= 4 {
172            runs.push((s, h));
173        }
174    }
175
176    // Tighten each line to its horizontal ink bounds.
177    runs.into_iter()
178        .map(|(t, b)| {
179            let (mut l, mut r) = (w, 0u32);
180            for y in t..b {
181                for x in 0..w {
182                    if luma(crop.get_pixel(x, y)) < thresh {
183                        l = l.min(x);
184                        r = r.max(x + 1);
185                    }
186                }
187            }
188            if l >= r {
189                (0, t, w, b)
190            } else {
191                (l, t, r, b)
192            }
193        })
194        .collect()
195}
196
197/// Layout labels whose content is recognised as running text.
198pub fn is_text_label(label: &str) -> bool {
199    matches!(
200        label,
201        "text"
202            | "title"
203            | "section_header"
204            | "list_item"
205            | "caption"
206            | "footnote"
207            | "code"
208            | "formula"
209    )
210}
211
212/// A line's page-point bounding box, `(l, t, r, b)`.
213pub type LineBox = (f32, f32, f32, f32);
214
215/// Gather every text-region line crop on a page, in page order: crop each
216/// text region (page points × `scale` → image px), split it into lines, prep
217/// each line, and keep the line's page-point bbox. The exact gathering the
218/// native `ocr_page` does — shared so the browser path produces the same
219/// cells given the same probabilities.
220pub fn prep_region_lines(
221    img: &RgbImage,
222    regions: &[crate::layout::Region],
223    scale: f32,
224) -> (Vec<LineBox>, Vec<PrepLine>) {
225    let (iw, ih) = img.dimensions();
226    let mut bboxes = Vec::new();
227    let mut lines = Vec::new();
228    for region in regions {
229        if !is_text_label(region.label) {
230            continue;
231        }
232        let l = (region.l * scale).max(0.0) as u32;
233        let t = (region.t * scale).max(0.0) as u32;
234        let r = ((region.r * scale).max(0.0) as u32).min(iw);
235        let b = ((region.b * scale).max(0.0) as u32).min(ih);
236        if r <= l || b <= t {
237            continue;
238        }
239        let crop = imageops::crop_imm(img, l, t, r - l, b - t).to_image();
240        for (lx, ly, rx, ry) in segment_lines(&crop) {
241            let line = imageops::crop_imm(&crop, lx, ly, rx - lx, ry - ly).to_image();
242            let Some(pl) = prep_line(&line) else {
243                continue;
244            };
245            bboxes.push((
246                (l + lx) as f32 / scale,
247                (t + ly) as f32 / scale,
248                (l + rx) as f32 / scale,
249                (t + ry) as f32 / scale,
250            ));
251            lines.push(pl);
252        }
253    }
254    (bboxes, lines)
255}
256
257/// [`prep_region_lines`] with the text detector's boxes as the line source
258/// (#570): inside each text region, the detected lines whose center falls in
259/// the region are the crops (whole, in the detector's reading order — see
260/// `det_boxes_inside`); a region the detector found nothing in falls back to
261/// the projection segmentation. `detected` is in image pixels (the detector ran
262/// on `img`).
263///
264/// Why: the projection profile cuts a region into full-width strips, so on a
265/// form — several fields on one baseline, separated by gaps the strip spans —
266/// one crop carries `FROM: J.Smith   DATE: 3/4/92   REF: 1138` and the
267/// recognizer glues or drops words (`OLDCOLDMENTHOLUIGHTS&ULTRA`). RapidOCR,
268/// docling's engine, recognizes the detector's own boxes, each a single text
269/// run with the detector's margin around it. On FUNSD this is the difference
270/// between 0.61 and 0.85 word recall (see the issue's measurements).
271pub fn prep_region_lines_det(
272    img: &RgbImage,
273    regions: &[crate::layout::Region],
274    scale: f32,
275    detected: &[crate::ocr_det::DetBox],
276) -> (Vec<LineBox>, Vec<PrepLine>) {
277    let (iw, ih) = img.dimensions();
278    let mut bboxes = Vec::new();
279    let mut lines = Vec::new();
280    for region in regions {
281        if !is_text_label(region.label) {
282            continue;
283        }
284        let l = (region.l * scale).max(0.0) as u32;
285        let t = (region.t * scale).max(0.0) as u32;
286        let r = ((region.r * scale).max(0.0) as u32).min(iw);
287        let b = ((region.b * scale).max(0.0) as u32).min(ih);
288        if r <= l || b <= t {
289            continue;
290        }
291        let inside = det_boxes_inside(detected, l, t, r, b, iw, ih);
292        if inside.is_empty() {
293            let (bb, pl) = prep_region_lines(img, std::slice::from_ref(region), scale);
294            bboxes.extend(bb);
295            lines.extend(pl);
296            continue;
297        }
298        for (dl, dt, dr, db) in inside {
299            let line = imageops::crop_imm(img, dl, dt, dr - dl, db - dt).to_image();
300            let Some(pl) = prep_line(&line) else {
301                continue;
302            };
303            bboxes.push((
304                dl as f32 / scale,
305                dt as f32 / scale,
306                dr as f32 / scale,
307                db as f32 / scale,
308            ));
309            lines.push(pl);
310        }
311    }
312    (bboxes, lines)
313}
314
315/// Split a text line into word tokens by a vertical ink-projection profile:
316/// runs of ink columns separated by whitespace wider than ~0.6× the line
317/// height (an inter-word/-column gap, not an inter-character one). Returns tight
318/// `(l, 0, r, h)` boxes in line-crop pixels. The browser table path recognizes
319/// these individually so each word carries its own box for column matching —
320/// the native pipeline gets word boxes from the pdfium text layer instead.
321pub fn segment_words(line: &RgbImage) -> Vec<(u32, u32, u32, u32)> {
322    let (w, h) = line.dimensions();
323    if w == 0 || h == 0 {
324        return Vec::new();
325    }
326    let mean: f32 = line.pixels().map(luma).sum::<f32>() / (w * h) as f32;
327    let thresh = mean * 0.7;
328    let mut col_ink = vec![0u32; w as usize];
329    for y in 0..h {
330        for x in 0..w {
331            if luma(line.get_pixel(x, y)) < thresh {
332                col_ink[x as usize] += 1;
333            }
334        }
335    }
336    let min_gap = ((h as f32) * 0.6).max(4.0) as u32;
337    let mut words = Vec::new();
338    let mut start: Option<u32> = None;
339    let mut last_ink = 0u32;
340    let mut gap = 0u32;
341    for x in 0..w {
342        if col_ink[x as usize] > 0 {
343            if start.is_none() {
344                start = Some(x);
345            }
346            last_ink = x;
347            gap = 0;
348        } else if let Some(s) = start {
349            gap += 1;
350            if gap >= min_gap {
351                words.push((s, 0, last_ink + 1, h));
352                start = None;
353            }
354        }
355    }
356    if let Some(s) = start {
357        words.push((s, 0, last_ink + 1, h));
358    }
359    words
360}
361
362/// Gather word crops from a page's *table* regions (browser table path, #157
363/// stage 3): crop each table region, split it into lines, split each line into
364/// words ([`segment_words`]), prep each word for recognition, and keep the
365/// word's page-point bbox. Recognizing table interiors is what gives the cell
366/// matcher the word boxes it needs — the native pipeline reads those from
367/// pdfium's text layer on digital pages, and calls this on scanned ones
368/// (`OcrModel::ocr_table_words`, #173), same as the browser path.
369pub fn prep_table_words(
370    img: &RgbImage,
371    regions: &[crate::layout::Region],
372    scale: f32,
373) -> (Vec<LineBox>, Vec<PrepLine>) {
374    let (iw, ih) = img.dimensions();
375    let mut bboxes = Vec::new();
376    let mut lines = Vec::new();
377    for region in regions {
378        if !crate::assemble::is_table_like(region.label) {
379            continue;
380        }
381        let l = (region.l * scale).max(0.0) as u32;
382        let t = (region.t * scale).max(0.0) as u32;
383        let r = ((region.r * scale).max(0.0) as u32).min(iw);
384        let b = ((region.b * scale).max(0.0) as u32).min(ih);
385        if r <= l || b <= t {
386            continue;
387        }
388        let crop = imageops::crop_imm(img, l, t, r - l, b - t).to_image();
389        for (lx, ly, rx, ry) in segment_lines(&crop) {
390            let line = imageops::crop_imm(&crop, lx, ly, rx - lx, ry - ly).to_image();
391            for (wx0, _, wx1, _) in segment_words(&line) {
392                let word = imageops::crop_imm(&line, wx0, 0, wx1 - wx0, ry - ly).to_image();
393                let Some(pl) = prep_line(&word) else {
394                    continue;
395                };
396                bboxes.push((
397                    (l + lx + wx0) as f32 / scale,
398                    (t + ly) as f32 / scale,
399                    (l + lx + wx1) as f32 / scale,
400                    (t + ry) as f32 / scale,
401                ));
402                lines.push(pl);
403            }
404        }
405    }
406    (bboxes, lines)
407}
408
409/// [`prep_table_words`] with the text detector's boxes as the word source
410/// (#570): inside each table region the detected boxes whose center falls in
411/// it are the crops (whole) — RapidOCR's cells, which is what
412/// docling's TableFormer matcher receives on a scanned table; a table the
413/// detector found nothing in falls back to the projection line/word split.
414pub fn prep_table_words_det(
415    img: &RgbImage,
416    regions: &[crate::layout::Region],
417    scale: f32,
418    detected: &[crate::ocr_det::DetBox],
419) -> (Vec<LineBox>, Vec<PrepLine>) {
420    let (iw, ih) = img.dimensions();
421    let mut bboxes = Vec::new();
422    let mut lines = Vec::new();
423    for region in regions {
424        if !crate::assemble::is_table_like(region.label) {
425            continue;
426        }
427        let l = (region.l * scale).max(0.0) as u32;
428        let t = (region.t * scale).max(0.0) as u32;
429        let r = ((region.r * scale).max(0.0) as u32).min(iw);
430        let b = ((region.b * scale).max(0.0) as u32).min(ih);
431        if r <= l || b <= t {
432            continue;
433        }
434        let inside = det_boxes_inside(detected, l, t, r, b, iw, ih);
435        if inside.is_empty() {
436            let (bb, pl) = prep_table_words(img, std::slice::from_ref(region), scale);
437            bboxes.extend(bb);
438            lines.extend(pl);
439            continue;
440        }
441        for (dl, dt, dr, db) in inside {
442            let word = imageops::crop_imm(img, dl, dt, dr - dl, db - dt).to_image();
443            let Some(pl) = prep_line(&word) else {
444                continue;
445            };
446            bboxes.push((
447                dl as f32 / scale,
448                dt as f32 / scale,
449                dr as f32 / scale,
450                db as f32 / scale,
451            ));
452            lines.push(pl);
453        }
454    }
455    (bboxes, lines)
456}
457
458/// The detected boxes whose center lies in the `l..r × t..b` image-pixel
459/// rectangle, in the detector's reading order. A box is taken whole (clipped
460/// to the image, not to the region): the layout box is often tight on the
461/// ink while the detector's unclip margin reaches past it, and cutting that
462/// margin off clips the first and last glyphs — RapidOCR recognizes the
463/// detector's box as drawn. The crop may thus poke a little outside the
464/// region; the cell keeps the box geometry and assembly assigns it by
465/// overlap, as for any cell.
466fn det_boxes_inside(
467    detected: &[crate::ocr_det::DetBox],
468    l: u32,
469    t: u32,
470    r: u32,
471    b: u32,
472    iw: u32,
473    ih: u32,
474) -> Vec<(u32, u32, u32, u32)> {
475    detected
476        .iter()
477        .filter(|d| {
478            let (cx, cy) = ((d.l + d.r) / 2.0, (d.t + d.b) / 2.0);
479            cx >= l as f32 && cx < r as f32 && cy >= t as f32 && cy < b as f32
480        })
481        .filter_map(|d| {
482            let dl = d.l.max(0.0) as u32;
483            let dt = d.t.max(0.0) as u32;
484            let dr = (d.r.max(0.0) as u32).min(iw);
485            let db = (d.b.max(0.0) as u32).min(ih);
486            (dr > dl && db > dt).then_some((dl, dt, dr, db))
487        })
488        .collect()
489}
490
491/// Normalize an image to the scan polarity every stage assumes — dark ink
492/// on light paper (the segmentation threshold and the recognition model's
493/// training data both bake it in): a predominantly dark page (mean luma
494/// below mid-gray — a dark-mode screenshot, an inverted scan) is inverted.
495/// Browser-path helper; the native pipeline never calls it (its input is
496/// scanned paper, and the conformance baseline stays untouched).
497pub fn normalize_polarity(mut img: RgbImage) -> RgbImage {
498    let (w, h) = img.dimensions();
499    if w == 0 || h == 0 {
500        return img;
501    }
502    let mean: f32 = img.pixels().map(luma).sum::<f32>() / (w * h) as f32;
503    if mean < 128.0 {
504        for px in img.pixels_mut() {
505            px.0 = [255 - px.0[0], 255 - px.0[1], 255 - px.0[2]];
506        }
507    }
508    img
509}
510
511/// Whole-image line preparation for the browser OCR path (no layout model:
512/// the page itself is the single text region). Returns page-order prepared
513/// lines; callers that need geometry use [`segment_lines`] directly.
514pub fn prep_page_lines(img: &RgbImage) -> Vec<PrepLine> {
515    segment_lines(img)
516        .into_iter()
517        .filter_map(|(l, t, r, b)| {
518            let line = imageops::crop_imm(img, l, t, r - l, b - t).to_image();
519            prep_line(&line)
520        })
521        .collect()
522}
523
524/// RapidOCR's `text_score` (docling's `RapidOcrOptions.text_score`, 0.5): a
525/// recognized line whose mean emitted-character confidence falls below it is
526/// dropped — a shaded bar, a halftone, a rule read as a run of letters
527/// (`TUANEURHENSTGOFUPHYCLTOSHSWOYT` for a form's grey `GEOGRAPHY` band)
528/// would otherwise land in the document as a heading. `DOCLING_RS_OCR_TEXT_SCORE`
529/// overrides it (`0` keeps every line, the pre-#570 behavior). Cached.
530pub fn text_score() -> f32 {
531    static SCORE: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
532    *SCORE.get_or_init(|| {
533        docling_core::env::parse::<f32>("DOCLING_RS_OCR_TEXT_SCORE")
534            .filter(|s| s.is_finite() && (0.0..=1.0).contains(s))
535            .unwrap_or(0.5)
536    })
537}
538
539/// Deterministic recognition batching: page-order line indices grouped by
540/// exact width (equal widths share a run — bit-identical to one-at-a-time
541/// recognition, see `ocr.rs`), each group split into [`REC_BATCH`] chunks.
542pub fn width_batches(lines: &[PrepLine]) -> Vec<(usize, Vec<usize>)> {
543    let mut by_width: std::collections::BTreeMap<usize, Vec<usize>> =
544        std::collections::BTreeMap::new();
545    for (ix, pl) in lines.iter().enumerate() {
546        by_width.entry(pl.w).or_default().push(ix);
547    }
548    let mut out = Vec::new();
549    for (w, ixs) in by_width {
550        for chunk in ixs.chunks(REC_BATCH) {
551            out.push((w, chunk.to_vec()));
552        }
553    }
554    out
555}
556
557/// Pack one width-batch into the model's `(N, 3, H, W)` input buffer.
558pub fn batch_input(w: usize, chunk: &[usize], lines: &[PrepLine]) -> Vec<f32> {
559    let hw = REC_HEIGHT as usize * w;
560    let mut data = vec![0f32; chunk.len() * 3 * hw];
561    for (i, &ix) in chunk.iter().enumerate() {
562        data[i * 3 * hw..(i + 1) * 3 * hw].copy_from_slice(&lines[ix].data);
563    }
564    data
565}
566
567#[cfg(test)]
568mod tests {
569    use super::*;
570
571    /// A synthetic "page": white background with two black text bars.
572    fn page() -> RgbImage {
573        let mut img = RgbImage::from_pixel(200, 100, Rgb([255, 255, 255]));
574        for y in 20..30 {
575            for x in 10..190 {
576                img.put_pixel(x, y, Rgb([0, 0, 0]));
577            }
578        }
579        for y in 60..72 {
580            for x in 10..120 {
581                img.put_pixel(x, y, Rgb([0, 0, 0]));
582            }
583        }
584        img
585    }
586
587    #[test]
588    fn segments_and_preps_page_lines() {
589        let lines = prep_page_lines(&page());
590        assert_eq!(lines.len(), 2);
591        for pl in &lines {
592            assert_eq!(pl.data.len(), 3 * REC_HEIGHT as usize * pl.w);
593        }
594        // Different aspect ratios → different widths → separate batches.
595        let batches = width_batches(&lines);
596        assert_eq!(batches.len(), 2);
597        let (w0, chunk0) = &batches[0];
598        assert_eq!(
599            batch_input(*w0, chunk0, &lines).len(),
600            3 * REC_HEIGHT as usize * w0
601        );
602    }
603
604    /// #570: inside a text region the detector's boxes are the crops — one
605    /// per box, clipped to the region, in the detector's order — and a region
606    /// the detector saw nothing in falls back to the projection strips. The
607    /// two-bar page as one region: two detected boxes → two det-shaped lines;
608    /// the region's second bar split by the detector into two words → three
609    /// lines where the projection found two; no boxes → the projection's two.
610    #[test]
611    fn detector_boxes_are_the_region_lines_with_projection_fallback() {
612        let img = page();
613        let region = crate::layout::Region {
614            label: "text",
615            score: 0.9,
616            l: 0.0,
617            t: 0.0,
618            r: 200.0,
619            b: 100.0,
620        };
621        let bx = |l: f32, t: f32, r: f32, b: f32| crate::ocr_det::DetBox {
622            l,
623            t,
624            r,
625            b,
626            score: 0.9,
627        };
628        let det = [bx(8.0, 18.0, 192.0, 32.0), bx(8.0, 58.0, 122.0, 74.0)];
629        let (boxes, lines) = prep_region_lines_det(&img, std::slice::from_ref(&region), 1.0, &det);
630        assert_eq!(
631            boxes,
632            vec![(8.0, 18.0, 192.0, 32.0), (8.0, 58.0, 122.0, 74.0)]
633        );
634        assert_eq!(lines.len(), 2);
635        // A box reaching past the region is taken whole (clipped only to the
636        // image); one centered outside the region is not its line.
637        let narrow = crate::layout::Region {
638            r: 150.0,
639            ..region.clone()
640        };
641        let det = [bx(-10.0, 18.0, 250.0, 32.0), bx(300.0, 60.0, 400.0, 70.0)];
642        let (boxes, _) = prep_region_lines_det(&img, std::slice::from_ref(&narrow), 1.0, &det);
643        assert_eq!(boxes, vec![(0.0, 18.0, 200.0, 32.0)]);
644        // The detector split the second bar into two words.
645        let det = [
646            bx(8.0, 18.0, 192.0, 32.0),
647            bx(8.0, 58.0, 60.0, 74.0),
648            bx(70.0, 58.0, 122.0, 74.0),
649        ];
650        let (boxes, _) = prep_region_lines_det(&img, std::slice::from_ref(&region), 1.0, &det);
651        assert_eq!(boxes.len(), 3);
652        // No detected box in the region → the projection segmentation.
653        let (fallback, _) = prep_region_lines_det(&img, std::slice::from_ref(&region), 1.0, &[]);
654        let (projection, _) = prep_region_lines(&img, std::slice::from_ref(&region), 1.0);
655        assert_eq!(fallback, projection);
656        assert_eq!(projection.len(), 2);
657        // Table regions take the same boxes as their word crops.
658        let table = crate::layout::Region {
659            label: "table",
660            ..region.clone()
661        };
662        let det = [bx(8.0, 18.0, 92.0, 32.0), bx(108.0, 18.0, 192.0, 32.0)];
663        let (words, _) = prep_table_words_det(&img, std::slice::from_ref(&table), 1.0, &det);
664        assert_eq!(
665            words,
666            vec![(8.0, 18.0, 92.0, 32.0), (108.0, 18.0, 192.0, 32.0)]
667        );
668        let (fallback, _) = prep_table_words_det(&img, std::slice::from_ref(&table), 1.0, &[]);
669        assert_eq!(
670            fallback,
671            prep_table_words(&img, std::slice::from_ref(&table), 1.0).0
672        );
673    }
674
675    /// RapidOCR's `text_score` default, the pre-#570 "keep everything" with
676    /// `0`, and garbage values ignored.
677    #[test]
678    fn text_score_default_is_rapidocrs() {
679        assert_eq!(text_score(), 0.5);
680    }
681
682    #[test]
683    fn dark_mode_pages_normalize_to_scan_polarity() {
684        // The same two-bar page, inverted (light text on dark) — the raw
685        // segmentation misfires (it thresholds the dark *background* as ink,
686        // so the "lines" it finds are the inter-bar gaps, not the bars);
687        // polarity normalization recovers the true structure.
688        let mut dark = page();
689        for px in dark.pixels_mut() {
690            px.0 = [255 - px.0[0], 255 - px.0[1], 255 - px.0[2]];
691        }
692        assert_ne!(segment_lines(&dark), segment_lines(&page()));
693        let fixed = normalize_polarity(dark);
694        assert_eq!(segment_lines(&fixed), segment_lines(&page()));
695        assert_eq!(prep_page_lines(&fixed).len(), 2);
696        // A light page passes through untouched.
697        let light = page();
698        assert_eq!(normalize_polarity(light.clone()), light);
699    }
700
701    #[test]
702    fn ctc_decode_collapses_repeats_and_blanks() {
703        // 3 classes: blank, "a", "b"; timesteps a a blank b b → "ab".
704        let chars = dict_chars("a\nb");
705        assert_eq!(chars.len(), 4); // blank, a, b, space
706        let probs = [
707            0.1, 0.8, 0.1, 0.0, // a
708            0.1, 0.8, 0.1, 0.0, // a (repeat collapses)
709            0.9, 0.05, 0.05, 0.0, // blank
710            0.1, 0.1, 0.8, 0.0, // b
711            0.1, 0.1, 0.8, 0.0, // b (repeat collapses)
712        ];
713        assert_eq!(decode_row(&chars, &probs, 4), "ab");
714    }
715}
716
717#[cfg(test)]
718mod word_segmentation {
719    use image::{Rgb, RgbImage};
720
721    /// A white line carrying two ink blocks separated by `gap` pixels.
722    fn line_with_gap(h: u32, gap: u32) -> RgbImage {
723        let w = 30 + gap + 30 + 10;
724        let mut img = RgbImage::from_pixel(w, h, Rgb([255, 255, 255]));
725        for (x0, x1) in [(5u32, 35u32), (35 + gap, 65 + gap)] {
726            for x in x0..x1.min(w) {
727                for y in h / 4..(3 * h / 4) {
728                    img.put_pixel(x, y, Rgb([0, 0, 0]));
729                }
730            }
731        }
732        img
733    }
734
735    /// [`segment_words`] splits on a gap of `0.6 x line height`, which is the
736    /// property the rest of the browser table path inherits: an inter-word
737    /// space never splits, but a column gap wider than that does. Anything
738    /// narrower stays a single box — and a single box can only ever land in one
739    /// TableFormer cell, so this threshold is the floor on how tight a table's
740    /// columns may be before its cells merge.
741    #[test]
742    fn words_split_only_on_gaps_above_six_tenths_of_the_line_height() {
743        for h in [16u32, 24, 32, 40] {
744            let split_at = (1..=40u32)
745                .find(|&gap| super::segment_words(&line_with_gap(h, gap)).len() >= 2)
746                .expect("some gap splits");
747            let ratio = split_at as f32 / h as f32;
748            assert!(
749                (0.5..=0.65).contains(&ratio),
750                "h={h}: split at {split_at}px ({ratio:.2} x height)"
751            );
752            // Just below the threshold the two blocks are one word.
753            assert_eq!(
754                super::segment_words(&line_with_gap(h, split_at - 1)).len(),
755                1,
756                "h={h}: a narrower gap must not split"
757            );
758        }
759    }
760}