Skip to main content

pdf_ocr/
engine.rs

1//! Pluggable OCR engine trait and built-in implementations.
2//!
3//! Provides the `OcrEngine` trait for integrating different OCR backends,
4//! and a `NoOpEngine` for testing.
5
6// Copyright (c) 2026 Innovation Trigger B.V.
7//
8// PDFluent is available under two licences, at your option: the GNU AGPLv3, or
9// the PDFluent Commercial Licence. See the LICENSE file in this repository --
10// that file travels with the copy you received, which a URL does not.
11
12/// A single word recognized by OCR.
13#[derive(Debug, Clone)]
14pub struct OcrWord {
15    /// The recognized text.
16    pub text: String,
17    /// Bounding box in pixel coordinates [x0, y0, x1, y1].
18    pub bbox_px: [u32; 4],
19    /// Recognition confidence (0.0 to 1.0).
20    pub confidence: f32,
21}
22
23/// Result of OCR processing on a single page.
24#[derive(Debug, Clone)]
25pub struct OcrPageResult {
26    /// Recognized words.
27    pub words: Vec<OcrWord>,
28    /// Overall confidence for the page.
29    pub confidence: f32,
30    /// Width of the source image in pixels.
31    pub image_width: u32,
32    /// Height of the source image in pixels.
33    pub image_height: u32,
34}
35
36impl OcrPageResult {
37    /// Get the full text of the page by joining all words with spaces.
38    pub fn full_text(&self) -> String {
39        self.words
40            .iter()
41            .map(|w| w.text.as_str())
42            .collect::<Vec<_>>()
43            .join(" ")
44    }
45}
46
47/// Trait for pluggable OCR engines.
48///
49/// Implementors provide character recognition on rasterized page images.
50pub trait OcrEngine: Send + Sync {
51    /// Recognize text in an image.
52    ///
53    /// # Arguments
54    /// * `image_data` - Raw pixel data (RGB, row-major).
55    /// * `width` - Image width in pixels.
56    /// * `height` - Image height in pixels.
57    /// * `dpi` - Resolution in dots per inch.
58    fn recognize(
59        &self,
60        image_data: &[u8],
61        width: u32,
62        height: u32,
63        dpi: u32,
64    ) -> std::result::Result<OcrPageResult, String>;
65
66    /// Return the list of supported languages.
67    fn supported_languages(&self) -> Vec<String>;
68}
69
70/// RGB row-major pixels from an RGB or RGBA buffer.
71///
72/// Backends accept both shapes and drop an alpha channel; anything else is a
73/// caller error naming the expected sizes.
74///
75/// Hidden: an internal helper the Python binding also uses, not API.
76#[doc(hidden)]
77pub fn rgb_pixels(
78    image_data: &[u8],
79    width: u32,
80    height: u32,
81) -> std::result::Result<Vec<u8>, String> {
82    let pixels = width as usize * height as usize;
83    if image_data.len() == pixels * 3 {
84        Ok(image_data.to_vec())
85    } else if image_data.len() == pixels * 4 {
86        Ok(image_data
87            .chunks_exact(4)
88            .flat_map(|px| px[..3].to_vec())
89            .collect())
90    } else {
91        Err(format!(
92            "unexpected image size: {} bytes for {}x{} (expected {} RGB or {} RGBA)",
93            image_data.len(),
94            width,
95            height,
96            pixels * 3,
97            pixels * 4,
98        ))
99    }
100}
101
102/// A no-op OCR engine that always returns empty results.
103///
104/// Useful for testing the pipeline without a real OCR backend.
105#[derive(Debug, Default)]
106pub struct NoOpEngine;
107
108impl OcrEngine for NoOpEngine {
109    fn recognize(
110        &self,
111        _image_data: &[u8],
112        width: u32,
113        height: u32,
114        _dpi: u32,
115    ) -> std::result::Result<OcrPageResult, String> {
116        Ok(OcrPageResult {
117            words: Vec::new(),
118            confidence: 0.0,
119            image_width: width,
120            image_height: height,
121        })
122    }
123
124    fn supported_languages(&self) -> Vec<String> {
125        Vec::new()
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132
133    #[test]
134    fn noop_engine_returns_empty() {
135        let engine = NoOpEngine;
136        let result = engine.recognize(&[], 100, 100, 300).unwrap();
137        assert!(result.words.is_empty());
138        assert_eq!(result.confidence, 0.0);
139        assert_eq!(result.image_width, 100);
140        assert_eq!(result.image_height, 100);
141        assert!(engine.supported_languages().is_empty());
142    }
143
144    #[test]
145    fn ocr_page_result_full_text() {
146        let result = OcrPageResult {
147            words: vec![
148                OcrWord {
149                    text: "Hello".to_string(),
150                    bbox_px: [0, 0, 50, 20],
151                    confidence: 0.95,
152                },
153                OcrWord {
154                    text: "World".to_string(),
155                    bbox_px: [60, 0, 110, 20],
156                    confidence: 0.90,
157                },
158            ],
159            confidence: 0.92,
160            image_width: 200,
161            image_height: 100,
162        };
163        assert_eq!(result.full_text(), "Hello World");
164    }
165
166    #[test]
167    fn rgb_pixels_keeps_rgb_drops_alpha_and_rejects_the_rest() {
168        assert_eq!(
169            rgb_pixels(&[1, 2, 3, 4, 5, 6], 2, 1).unwrap(),
170            vec![1, 2, 3, 4, 5, 6]
171        );
172        assert_eq!(
173            rgb_pixels(&[1, 2, 3, 255, 4, 5, 6, 0], 2, 1).unwrap(),
174            vec![1, 2, 3, 4, 5, 6]
175        );
176        assert!(rgb_pixels(&[1, 2, 3, 4, 5], 2, 1).is_err());
177    }
178}