1use image::{imageops, imageops::FilterType, Rgb, RgbImage};
13
14pub const REC_HEIGHT: u32 = 48;
16
17pub const REC_BATCH: usize = 16;
21
22pub struct PrepLine {
25 pub w: usize,
27 pub data: Vec<f32>,
29}
30
31pub 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 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
55pub fn dict_chars(dict: &str) -> Vec<String> {
58 let mut chars = vec![String::new()]; chars.extend(dict.lines().map(|s| s.to_string()));
60 chars.push(" ".to_string());
61 chars
62}
63
64pub fn decode_row(chars: &[String], probs: &[f32], nc: usize) -> String {
66 decode_row_scored(chars, probs, nc).0
67}
68
69pub 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
108pub 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; let min_ink = ((w as f32) * 0.005).max(1.0) as u32;
118
119 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 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 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 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
197pub 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
212pub type LineBox = (f32, f32, f32, f32);
214
215pub 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
257pub 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
315pub 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
362pub 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
409pub 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
458fn 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
491pub 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
511pub 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
524pub 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
539pub 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
557pub 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 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 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 #[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(®ion), 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 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 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(®ion), 1.0, &det);
651 assert_eq!(boxes.len(), 3);
652 let (fallback, _) = prep_region_lines_det(&img, std::slice::from_ref(®ion), 1.0, &[]);
654 let (projection, _) = prep_region_lines(&img, std::slice::from_ref(®ion), 1.0);
655 assert_eq!(fallback, projection);
656 assert_eq!(projection.len(), 2);
657 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 #[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 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 let light = page();
698 assert_eq!(normalize_polarity(light.clone()), light);
699 }
700
701 #[test]
702 fn ctc_decode_collapses_repeats_and_blanks() {
703 let chars = dict_chars("a\nb");
705 assert_eq!(chars.len(), 4); let probs = [
707 0.1, 0.8, 0.1, 0.0, 0.1, 0.8, 0.1, 0.0, 0.9, 0.05, 0.05, 0.0, 0.1, 0.1, 0.8, 0.0, 0.1, 0.1, 0.8, 0.0, ];
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 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 #[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 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}