pdfpurr 0.4.0

A comprehensive pure-Rust PDF library for reading, writing, rendering, and validating PDF documents
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

//! OCR engine implementation using `ocrs` (pure Rust).
//!
//! Uses the RTen inference engine to run neural network models for
//! text detection and recognition. No C/C++ dependencies.
//!
//! # Model Files
//!
//! Two `.rten` model files are required:
//! - `text-detection.rten` — detects word bounding boxes
//! - `text-recognition.rten` — recognizes text in detected regions
//!
//! Download from: <https://github.com/robertknight/ocrs-models>

use std::path::Path;

use crate::error::{PdfError, PdfResult};

use super::engine::{OcrEngine, OcrImage, OcrResult, OcrWord};

/// OCR engine backed by `ocrs` (pure Rust, RTen inference).
pub struct OcrsEngine {
    engine: ocrs::OcrEngine,
}

impl OcrsEngine {
    /// Creates a new OCR engine by loading model files from disk.
    pub fn new<P: AsRef<Path>>(
        detection_model_path: P,
        recognition_model_path: P,
    ) -> PdfResult<Self> {
        let detection_model = rten::Model::load_file(detection_model_path.as_ref())
            .map_err(|e| PdfError::OcrError(format!("Load detection model: {}", e)))?;

        let recognition_model = rten::Model::load_file(recognition_model_path.as_ref())
            .map_err(|e| PdfError::OcrError(format!("Load recognition model: {}", e)))?;

        let engine = ocrs::OcrEngine::new(ocrs::OcrEngineParams {
            detection_model: Some(detection_model),
            recognition_model: Some(recognition_model),
            ..Default::default()
        })
        .map_err(|e| PdfError::OcrError(format!("Init OCR engine: {}", e)))?;

        Ok(Self { engine })
    }
}

impl OcrEngine for OcrsEngine {
    fn recognize(&self, image: &OcrImage) -> PdfResult<OcrResult> {
        let width = image.width;
        let height = image.height;

        // ocrs ImageSource::from_bytes expects HWC u8 RGB data.
        // Convert grayscale to RGB by tripling each byte.
        let mut rgb = Vec::with_capacity((width * height * 3) as usize);
        for &gray in &image.data {
            rgb.push(gray);
            rgb.push(gray);
            rgb.push(gray);
        }

        let img_source = ocrs::ImageSource::from_bytes(&rgb, (width, height))
            .map_err(|e| PdfError::OcrError(format!("OCR image source: {}", e)))?;

        let input = self
            .engine
            .prepare_input(img_source)
            .map_err(|e| PdfError::OcrError(format!("OCR prepare: {}", e)))?;

        let word_rects = self
            .engine
            .detect_words(&input)
            .map_err(|e| PdfError::OcrError(format!("OCR detect: {}", e)))?;

        let lines = self.engine.find_text_lines(&input, &word_rects);

        let text_lines = self
            .engine
            .recognize_text(&input, &lines)
            .map_err(|e| PdfError::OcrError(format!("OCR recognize: {}", e)))?;

        let mut words = Vec::new();
        for (line_rects, text_line) in lines.iter().zip(text_lines.iter()) {
            if let Some(text_line) = text_line {
                for (rect, word_text) in line_rects.iter().zip(text_line.words()) {
                    let text = word_text.to_string();
                    if text.trim().is_empty() {
                        continue;
                    }

                    // Get axis-aligned bounding box from the rotated rect corners
                    let corners = rect.corners();
                    let min_x = corners.iter().map(|c| c.x).fold(f32::MAX, f32::min);
                    let min_y = corners.iter().map(|c| c.y).fold(f32::MAX, f32::min);
                    let max_x = corners.iter().map(|c| c.x).fold(f32::MIN, f32::max);
                    let max_y = corners.iter().map(|c| c.y).fold(f32::MIN, f32::max);

                    words.push(OcrWord {
                        text,
                        x: min_x.max(0.0) as u32,
                        y: min_y.max(0.0) as u32,
                        width: (max_x - min_x).max(1.0) as u32,
                        height: (max_y - min_y).max(1.0) as u32,
                        confidence: 0.9,
                    });
                }
            }
        }

        Ok(OcrResult {
            words,
            image_width: image.width,
            image_height: image.height,
        })
    }
}