oar_ocr_core/domain/tasks/
text_recognition.rs

1//! Concrete task implementations for text recognition.
2//!
3//! This module provides the text recognition task that converts text regions to strings.
4
5use super::validation::ensure_non_empty_images;
6use crate::ConfigValidator;
7use crate::core::OCRError;
8use crate::core::traits::TaskDefinition;
9use crate::core::traits::task::{ImageTaskInput, Task, TaskType};
10use crate::utils::ScoreValidator;
11use serde::{Deserialize, Serialize};
12
13/// Configuration for text recognition task.
14///
15/// Default values are aligned with PP-StructureV3.
16#[derive(Debug, Clone, Serialize, Deserialize, ConfigValidator)]
17pub struct TextRecognitionConfig {
18    /// Score threshold for recognition (default: 0.0, no filtering)
19    #[validate(range(min = 0.0, max = 1.0))]
20    pub score_threshold: f32,
21}
22
23impl Default for TextRecognitionConfig {
24    fn default() -> Self {
25        Self {
26            score_threshold: 0.0,
27        }
28    }
29}
30
31/// Output from text recognition task.
32#[derive(Debug, Clone)]
33pub struct TextRecognitionOutput {
34    /// Recognized text strings
35    pub texts: Vec<String>,
36    /// Confidence scores for each text
37    pub scores: Vec<f32>,
38    /// Character/word positions within each text line (optional)
39    /// Each inner vector contains normalized x-positions (0.0-1.0) for characters
40    /// Only populated when word box detection is enabled
41    pub char_positions: Vec<Vec<f32>>,
42    /// Column indices for each character in the CTC output
43    /// Used for accurate word box generation with compatible approach
44    pub char_col_indices: Vec<Vec<usize>>,
45    /// Total number of columns (sequence length) in the CTC output for each text line
46    pub sequence_lengths: Vec<usize>,
47}
48
49impl TextRecognitionOutput {
50    /// Creates an empty text recognition output.
51    pub fn empty() -> Self {
52        Self {
53            texts: Vec::new(),
54            scores: Vec::new(),
55            char_positions: Vec::new(),
56            char_col_indices: Vec::new(),
57            sequence_lengths: Vec::new(),
58        }
59    }
60
61    /// Creates a text recognition output with the given capacity.
62    pub fn with_capacity(capacity: usize) -> Self {
63        Self {
64            texts: Vec::with_capacity(capacity),
65            scores: Vec::with_capacity(capacity),
66            char_positions: Vec::with_capacity(capacity),
67            char_col_indices: Vec::with_capacity(capacity),
68            sequence_lengths: Vec::with_capacity(capacity),
69        }
70    }
71}
72
73impl Default for TextRecognitionOutput {
74    fn default() -> Self {
75        Self::empty()
76    }
77}
78
79impl TaskDefinition for TextRecognitionOutput {
80    const TASK_NAME: &'static str = "text_recognition";
81    const TASK_DOC: &'static str = "Text recognition - converting text regions to strings";
82
83    fn empty() -> Self {
84        TextRecognitionOutput::empty()
85    }
86}
87
88/// Text recognition task implementation.
89#[derive(Debug, Default)]
90pub struct TextRecognitionTask;
91
92impl TextRecognitionTask {
93    /// Creates a new text recognition task.
94    pub fn new(_config: TextRecognitionConfig) -> Self {
95        Self
96    }
97}
98
99impl Task for TextRecognitionTask {
100    type Config = TextRecognitionConfig;
101    type Input = ImageTaskInput;
102    type Output = TextRecognitionOutput;
103
104    fn task_type(&self) -> TaskType {
105        TaskType::TextRecognition
106    }
107
108    fn validate_input(&self, input: &Self::Input) -> Result<(), OCRError> {
109        ensure_non_empty_images(&input.images, "No images provided for text recognition")?;
110
111        Ok(())
112    }
113
114    fn validate_output(&self, output: &Self::Output) -> Result<(), OCRError> {
115        // Validate score ranges
116        let validator = ScoreValidator::new_unit_range("score");
117        validator.validate_scores_with(&output.scores, |idx| format!("Text {}", idx))?;
118
119        Ok(())
120    }
121
122    fn empty_output(&self) -> Self::Output {
123        TextRecognitionOutput::empty()
124    }
125}
126
127#[cfg(test)]
128mod tests {
129    use super::*;
130    use image::RgbImage;
131
132    #[test]
133    fn test_text_recognition_task_creation() {
134        let task = TextRecognitionTask;
135        assert_eq!(task.task_type(), TaskType::TextRecognition);
136    }
137
138    #[test]
139    fn test_input_validation() {
140        let task = TextRecognitionTask;
141
142        // Empty images should fail
143        let empty_input = ImageTaskInput::new(vec![]);
144        assert!(task.validate_input(&empty_input).is_err());
145
146        // Valid images should pass
147        let valid_input = ImageTaskInput::new(vec![RgbImage::new(100, 32)]);
148        assert!(task.validate_input(&valid_input).is_ok());
149    }
150
151    #[test]
152    fn test_output_validation() {
153        let task = TextRecognitionTask;
154
155        // Matching texts and scores should pass
156        let output = TextRecognitionOutput {
157            texts: vec!["Hello".to_string()],
158            scores: vec![0.95],
159            ..Default::default()
160        };
161        assert!(task.validate_output(&output).is_ok());
162
163        // Invalid score should fail
164        let bad_score = TextRecognitionOutput {
165            texts: vec!["Hello".to_string()],
166            scores: vec![1.5],
167            ..Default::default()
168        };
169        assert!(task.validate_output(&bad_score).is_err());
170    }
171}
oar_ocr_core/domain/tasks/text_recognition.rs

oar_ocr_core/domain/tasks/
text_recognition.rs
