docling_pdf/picture.rs
1//! OCR for a standalone picture (#645): the text the recognizer reads off an
2//! image embedded in a non-PDF document — a DOCX/PPTX screenshot, an HTML
3//! figure, a sampled video frame — so the picture-OCR enrichment in
4//! `docling` can attach it to the picture as docling's description
5//! annotation. The models are the ML pipeline's own (the recognizer pair,
6//! the text detector, Tesseract under `ocr_engine`), loaded lazily on the
7//! first picture exactly as they are for a scanned page, so a converter that
8//! OCRs pictures and scans alike holds one set of sessions.
9//!
10//! The picture is treated as the image input `Pipeline::convert_image` sees:
11//! its own scale-1.0 page whose one text region is the whole image, read at
12//! docling's effective OCR resolution for images (3 px/pt shrunk to
13//! RapidOCR's 2000 px longer side, #570) — no layout pass, since the caller
14//! already knows the whole image is the picture. The detector's boxes are
15//! the line crops when `ocr_det.onnx` is installed (#570); the ink-projection
16//! strips otherwise. Lines under `DOCLING_RS_OCR_TEXT_SCORE` are dropped by
17//! the recognizer, as on a page.
18
19use image::RgbImage;
20
21use crate::layout::Region;
22use crate::pdfium_backend::TextCell;
23use crate::{ocr, ocr_input, page_ocr_scale, EnrichSlot, PdfError, Pipeline, RAPIDOCR_MAX_SIDE};
24
25/// What the OCR read off one picture: the lines in reading order (top to
26/// bottom, left to right within a row, rows joined by `\n`) and the engine
27/// that read them — docling's `PictureDescriptionData.provenance`.
28#[derive(Debug, Clone, PartialEq, Eq)]
29pub struct PictureText {
30 pub text: String,
31 /// `ppocr` or `tesseract` — the [`ocr::OcrEngine`] that ran.
32 pub provenance: &'static str,
33}
34
35impl Pipeline {
36 /// Decode an embedded picture's bytes (PNG/JPEG/GIF/BMP/TIFF/WebP) under
37 /// the standalone-image pixel cap (`DOCLING_RS_MAX_IMAGE_PIXELS`, 30000
38 /// per side) — a crafted header never allocates a multi-GB bitmap.
39 pub fn decode_picture(&self, bytes: &[u8]) -> Result<RgbImage, PdfError> {
40 crate::decode_image_limited(bytes)
41 }
42
43 /// The DocumentFigureClassifier's predictions for a picture (descending
44 /// confidence), loading the model on first use. `None` when the model is
45 /// not installed (warned once by the loader) or inference failed — the
46 /// caller decides what an unclassifiable picture means (the picture-OCR
47 /// class filter lets it through).
48 pub fn classify_picture(&mut self, img: &RgbImage) -> Option<Vec<docling_core::PictureClass>> {
49 let mut guard = self.classifier.lock().unwrap_or_else(|p| p.into_inner());
50 if matches!(*guard, EnrichSlot::Unloaded) {
51 *guard = match crate::enrich::PictureClassifier::load_with(crate::intra_threads()) {
52 Some(m) => EnrichSlot::Ready(m),
53 None => EnrichSlot::Missing,
54 };
55 }
56 let EnrichSlot::Ready(model) = &mut *guard else {
57 return None;
58 };
59 match model.classify(img) {
60 Ok(classes) => Some(classes),
61 Err(e) => {
62 eprintln!("docling-pdf: picture classifier: {e}");
63 None
64 }
65 }
66 }
67
68 /// OCR a picture: `Ok(None)` when the engine read no text (or could not
69 /// load — the recognizer's one-time warning says so, and `skip_ocr`
70 /// short-circuits without one), `Err` only for a failed inference.
71 pub fn ocr_picture(&mut self, img: &RgbImage) -> Result<Option<PictureText>, PdfError> {
72 let provenance = match self.ocr_engine {
73 ocr::OcrEngine::PpOcr => "ppocr",
74 ocr::OcrEngine::Tesseract => "tesseract",
75 };
76 let (w, h) = (img.width() as f32, img.height() as f32);
77 if w < 1.0 || h < 1.0 {
78 return Ok(None);
79 }
80 // The image is its own page at 1 px per point; the OCR reads it at
81 // `picture_ocr_scale` under RapidOCR's 2000 px cap (#666), else at
82 // docling's image resolution unless `ocr_scale` says otherwise.
83 let ocr_scale = match self.picture_ocr_scale {
84 Some(p) => picture_ocr_scale(p, w, h),
85 None => page_ocr_scale(self.ocr_scale, w, h, 1.0),
86 };
87 let mut cache = None;
88 let (view, scale) = ocr_input(&mut cache, img, 1.0, ocr_scale);
89 let region = Region {
90 label: "text",
91 score: 1.0,
92 l: 0.0,
93 t: 0.0,
94 r: w,
95 b: h,
96 };
97 let worker = self.primary()?;
98 let detected = match worker.det_model() {
99 Some(det) => Some(det.detect(view).map_err(PdfError::Ocr)?),
100 None => None,
101 };
102 let Some(model) = worker.ocr_model()? else {
103 return Ok(None);
104 };
105 let cells = model
106 .ocr_page_with(
107 view,
108 std::slice::from_ref(®ion),
109 scale,
110 detected.as_deref(),
111 )
112 .map_err(PdfError::Ocr)?;
113 let text = lines_text(cells.into_iter().map(|(c, _)| c).collect());
114 Ok((!text.is_empty()).then_some(PictureText { text, provenance }))
115 }
116}
117
118/// The px/px a picture reads at under `picture_ocr_scale` = `p` (#666):
119/// `p`, shrunk so the longer side stays within RapidOCR's 2000 px
120/// (`max_side_len`) — at `1.0`, RapidOCR's own rule: native pixels, only
121/// shrinking what exceeds the cap. `None` when that is the picture's own
122/// resolution (no resample).
123fn picture_ocr_scale(p: f32, w: f32, h: f32) -> Option<f32> {
124 let s = p.min(RAPIDOCR_MAX_SIDE / w.max(h));
125 ((s - 1.0).abs() > 1e-3).then_some(s)
126}
127
128/// Join OCR line cells into text in reading order: rows top to bottom, cells
129/// left to right within a row. A cell joins the row whose vertical extent
130/// contains its centre — two columns of a screenshot's dialog read across,
131/// the way a human scans it; a line that sits lower starts a new row. Empty
132/// cells are dropped.
133pub(crate) fn lines_text(mut cells: Vec<TextCell>) -> String {
134 cells.retain(|c| !c.text.trim().is_empty());
135 cells.sort_by(|a, b| a.t.total_cmp(&b.t).then(a.l.total_cmp(&b.l)));
136 let mut rows: Vec<(f32, f32, Vec<TextCell>)> = Vec::new();
137 for cell in cells {
138 let mid = (cell.t + cell.b) / 2.0;
139 match rows.last_mut() {
140 Some((t, b, row)) if mid >= *t && mid <= *b => {
141 *t = t.min(cell.t);
142 *b = b.max(cell.b);
143 row.push(cell);
144 }
145 _ => rows.push((cell.t, cell.b, vec![cell])),
146 }
147 }
148 rows.iter_mut()
149 .map(|(_, _, row)| {
150 row.sort_by(|a, b| a.l.total_cmp(&b.l));
151 row.iter()
152 .map(|c| c.text.trim())
153 .collect::<Vec<_>>()
154 .join(" ")
155 })
156 .collect::<Vec<_>>()
157 .join("\n")
158}
159
160#[cfg(test)]
161mod tests {
162 use super::*;
163
164 fn cell(text: &str, l: f32, t: f32, r: f32, b: f32) -> TextCell {
165 TextCell {
166 text: text.into(),
167 l,
168 t,
169 r,
170 b,
171 }
172 }
173
174 /// Rows read across, then down; a cell whose centre sits below the
175 /// current row starts a new one, and empty cells vanish.
176 #[test]
177 fn lines_text_reads_rows_across_then_down() {
178 let cells = vec![
179 cell("right", 200.0, 10.0, 260.0, 30.0),
180 cell("below", 10.0, 40.0, 80.0, 60.0),
181 cell("left", 10.0, 12.0, 60.0, 32.0),
182 cell(" ", 10.0, 70.0, 20.0, 80.0),
183 ];
184 assert_eq!(lines_text(cells), "left right\nbelow");
185 assert_eq!(lines_text(Vec::new()), "");
186 }
187
188 /// #666: `picture_ocr_scale(1.0)` is RapidOCR's rule — a 720 × 180
189 /// screenshot reads at its own 720 × 180 (no resample, where the default
190 /// upscales it to 2000 × 500), a 3000 × 1000 one shrinks to 2000 × 667.
191 #[test]
192 fn picture_scale_reads_native_pixels_under_the_2000_px_cap() {
193 let dims = |p: f32, w: f32, h: f32| {
194 let s = picture_ocr_scale(p, w, h).unwrap_or(1.0);
195 ((w * s).round(), (h * s).round())
196 };
197 assert_eq!(picture_ocr_scale(1.0, 720.0, 180.0), None);
198 assert_eq!(dims(1.0, 3000.0, 1000.0), (2000.0, 667.0));
199 assert_eq!(dims(2.0, 720.0, 180.0), (1440.0, 360.0));
200 assert_eq!(dims(2.0, 1280.0, 720.0), (2000.0, 1125.0));
201 // The default rule, for comparison: 3 px/pt under the cap.
202 let default = page_ocr_scale(None, 720.0, 180.0, 1.0).unwrap();
203 assert_eq!(
204 ((720.0 * default).round(), (180.0 * default).round()),
205 (2000.0, 500.0)
206 );
207 }
208}