docling_pdf/
redact_ocr.rs

1//! `ImageRedaction::BoxOut` for the PII redaction pass (#621): OCR an
2//! embedded image with the pipeline's own recognizer (the text detector's
3//! line crops when `ocr_det.onnx` is installed, Tesseract under
4//! `ocr_engine`), ask the pass which lines carry personal data, and paint a
5//! solid box over each such line before the image goes out in the JSON /
6//! embedded Markdown. The image is read the way a standalone image input
7//! is — its own scale-1.0 page, one text region covering it, docling's OCR
8//! resolution for images (#570) — so what the scan pipeline would read off
9//! it is what gets covered.
10
11use image::{Rgb, RgbImage};
12
13use crate::layout::Region;
14use crate::pdfium_backend::{PdfPage, TextCell};
15use crate::{decode_image_limited, ocr_input, page_ocr_scale, PdfError, Pipeline};
16
17/// Padding around a covered line, in image pixels: the recognizer's line
18/// box hugs the ink, a few pixels more hide the ascender/descender fringe.
19const PAD: u32 = 2;
20
21impl Pipeline {
22    /// Box out the lines of `bytes` (an embedded PNG/JPEG/…) for which
23    /// `has_pii` says so. `Ok(None)` when no line matched — the image is
24    /// kept as it is; `Ok(Some(png))` the repainted image as PNG bytes;
25    /// `Err` when the OCR could not run (no recognizer: the caller degrades
26    /// to dropping the image).
27    pub fn box_out(
28        &mut self,
29        bytes: &[u8],
30        has_pii: &dyn Fn(&str) -> bool,
31    ) -> Result<Option<Vec<u8>>, PdfError> {
32        let mut image = decode_image_limited(bytes)?;
33        let (w, h) = image.dimensions();
34        if w < 2 || h < 2 {
35            return Ok(None);
36        }
37        let cells = self.ocr_image_lines(&image)?;
38        let hit: Vec<&TextCell> = cells.iter().filter(|c| has_pii(&c.text)).collect();
39        if hit.is_empty() {
40            return Ok(None);
41        }
42        for c in hit {
43            let l = (c.l.max(0.0) as u32).saturating_sub(PAD);
44            let t = (c.t.max(0.0) as u32).saturating_sub(PAD);
45            let r = ((c.r.max(0.0) as u32).saturating_add(PAD)).min(w);
46            let b = ((c.b.max(0.0) as u32).saturating_add(PAD)).min(h);
47            for y in t..b {
48                for x in l..r {
49                    image.put_pixel(x, y, Rgb([0, 0, 0]));
50                }
51            }
52        }
53        let mut png = Vec::new();
54        image
55            .write_to(&mut std::io::Cursor::new(&mut png), image::ImageFormat::Png)
56            .map_err(|e| PdfError::Pdfium(format!("box_out: png encode: {e}")))?;
57        Ok(Some(png))
58    }
59
60    /// The recognizer's line cells over a whole image, in image pixels.
61    /// `Err` when the recognizer is unavailable (`skip_ocr`, model missing).
62    fn ocr_image_lines(&mut self, image: &RgbImage) -> Result<Vec<TextCell>, PdfError> {
63        let (w, h) = image.dimensions();
64        let page = PdfPage {
65            width: w as f32,
66            height: h as f32,
67            scale: 1.0,
68            cells: Vec::new(),
69            code_cells: Vec::new(),
70            checkboxes: Vec::new(),
71            word_cells: Vec::new(),
72            image_layout: None,
73            image: image.clone(),
74            links: Vec::new(),
75            rotation: 0,
76        };
77        let ocr_scale = page_ocr_scale(self.ocr_scale, page.width, page.height, page.scale);
78        let mut cache = None;
79        let (view, scale) = ocr_input(&mut cache, &page.image, 1.0, ocr_scale);
80        let region = Region {
81            label: "text",
82            score: 1.0,
83            l: 0.0,
84            t: 0.0,
85            r: w as f32,
86            b: h as f32,
87        };
88        let worker = self.primary()?;
89        let detected = match worker.det_model() {
90            Some(det) => Some(det.detect(view).map_err(PdfError::Ocr)?),
91            None => None,
92        };
93        let Some(model) = worker.ocr_model()? else {
94            return Err(PdfError::Ocr("OCR model unavailable".into()));
95        };
96        let cells = model
97            .ocr_page_with(
98                view,
99                std::slice::from_ref(&region),
100                scale,
101                detected.as_deref(),
102            )
103            .map_err(PdfError::Ocr)?;
104        Ok(cells.into_iter().map(|(c, _)| c).collect())
105    }
106}
docling_pdf/redact_ocr.rs

docling_pdf/
redact_ocr.rs
