oar-ocr-core 0.6.3

Core types and predictors for oar-ocr
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

//! Text Recognition Predictor
//!
//! This module provides a high-level API for text recognition from cropped text images.

use super::builder::PredictorBuilderState;
use crate::TaskPredictorBuilder;
use crate::core::OcrResult;
use crate::core::errors::OCRError;
use crate::core::traits::OrtConfigurable;
use crate::core::traits::adapter::AdapterBuilder;
use crate::core::traits::task::ImageTaskInput;
use crate::domain::adapters::TextRecognitionAdapterBuilder;
use crate::domain::tasks::text_recognition::{TextRecognitionConfig, TextRecognitionTask};
use crate::predictors::TaskPredictorCore;
use image::RgbImage;
use std::path::{Path, PathBuf};

/// Text recognition prediction result
#[derive(Debug, Clone)]
pub struct TextRecognitionResult {
    /// Recognized text for each input image
    pub texts: Vec<String>,
    /// Confidence scores for each recognition
    pub scores: Vec<f32>,
}

/// Text recognition predictor
pub struct TextRecognitionPredictor {
    core: TaskPredictorCore<TextRecognitionTask>,
}

impl TextRecognitionPredictor {
    /// Create a new builder for the text recognition predictor
    pub fn builder() -> TextRecognitionPredictorBuilder {
        TextRecognitionPredictorBuilder::new()
    }

    /// Predict text in the given images.
    pub fn predict(&self, images: Vec<RgbImage>) -> OcrResult<TextRecognitionResult> {
        let input = ImageTaskInput::new(images);
        let output = self.core.predict(input)?;
        Ok(TextRecognitionResult {
            texts: output.texts,
            scores: output.scores,
        })
    }
}

/// Builder for text recognition predictor
#[derive(TaskPredictorBuilder)]
#[builder(config = TextRecognitionConfig)]
pub struct TextRecognitionPredictorBuilder {
    state: PredictorBuilderState<TextRecognitionConfig>,
    dict_path: Option<PathBuf>,
}

impl TextRecognitionPredictorBuilder {
    pub fn new() -> Self {
        Self {
            state: PredictorBuilderState::new(TextRecognitionConfig {
                score_threshold: 0.0,
                max_text_length: 100,
            }),
            dict_path: None,
        }
    }

    pub fn score_threshold(mut self, threshold: f32) -> Self {
        self.state.config_mut().score_threshold = threshold;
        self
    }

    pub fn max_text_length(mut self, max_length: usize) -> Self {
        self.state.config_mut().max_text_length = max_length;
        self
    }

    pub fn dict_path<P: AsRef<Path>>(mut self, path: P) -> Self {
        self.dict_path = Some(path.as_ref().to_path_buf());
        self
    }

    pub fn build<P: AsRef<Path>>(self, model_path: P) -> OcrResult<TextRecognitionPredictor> {
        let Self { state, dict_path } = self;
        let (config, ort_config) = state.into_parts();

        let dict_path = dict_path
            .ok_or_else(|| OCRError::missing_field("dict_path", "TextRecognitionPredictor"))?;

        // Load character dictionary from file
        let character_dict = std::fs::read_to_string(&dict_path)?
            .lines()
            .map(|s| s.to_string())
            .collect::<Vec<String>>();

        let mut adapter_builder = TextRecognitionAdapterBuilder::new()
            .with_config(config.clone())
            .character_dict(character_dict);

        if let Some(ort_cfg) = ort_config {
            adapter_builder = adapter_builder.with_ort_config(ort_cfg);
        }

        let adapter = Box::new(adapter_builder.build(model_path.as_ref())?);
        let task = TextRecognitionTask::new(config.clone());
        Ok(TextRecognitionPredictor {
            core: TaskPredictorCore::new(adapter, task, config),
        })
    }
}

impl Default for TextRecognitionPredictorBuilder {
    fn default() -> Self {
        Self::new()
    }
}