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, TaskSchema, 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 schema(&self) -> TaskSchema {
109        TaskSchema::new(
110            TaskType::TextRecognition,
111            vec!["text_boxes".to_string()],
112            vec!["text_strings".to_string(), "scores".to_string()],
113        )
114    }
115
116    fn validate_input(&self, input: &Self::Input) -> Result<(), OCRError> {
117        ensure_non_empty_images(&input.images, "No images provided for text recognition")?;
118
119        Ok(())
120    }
121
122    fn validate_output(&self, output: &Self::Output) -> Result<(), OCRError> {
123        // Validate score ranges
124        let validator = ScoreValidator::new_unit_range("score");
125        validator.validate_scores_with(&output.scores, |idx| format!("Text {}", idx))?;
126
127        Ok(())
128    }
129
130    fn empty_output(&self) -> Self::Output {
131        TextRecognitionOutput::empty()
132    }
133}
134
135#[cfg(test)]
136mod tests {
137    use super::*;
138    use image::RgbImage;
139
140    #[test]
141    fn test_text_recognition_task_creation() {
142        let task = TextRecognitionTask;
143        assert_eq!(task.task_type(), TaskType::TextRecognition);
144    }
145
146    #[test]
147    fn test_input_validation() {
148        let task = TextRecognitionTask;
149
150        // Empty images should fail
151        let empty_input = ImageTaskInput::new(vec![]);
152        assert!(task.validate_input(&empty_input).is_err());
153
154        // Valid images should pass
155        let valid_input = ImageTaskInput::new(vec![RgbImage::new(100, 32)]);
156        assert!(task.validate_input(&valid_input).is_ok());
157    }
158
159    #[test]
160    fn test_output_validation() {
161        let task = TextRecognitionTask;
162
163        // Matching texts and scores should pass
164        let output = TextRecognitionOutput {
165            texts: vec!["Hello".to_string()],
166            scores: vec![0.95],
167            ..Default::default()
168        };
169        assert!(task.validate_output(&output).is_ok());
170
171        // Invalid score should fail
172        let bad_score = TextRecognitionOutput {
173            texts: vec!["Hello".to_string()],
174            scores: vec![1.5],
175            ..Default::default()
176        };
177        assert!(task.validate_output(&bad_score).is_err());
178    }
179
180    #[test]
181    fn test_schema() {
182        let task = TextRecognitionTask;
183        let schema = task.schema();
184        assert_eq!(schema.task_type, TaskType::TextRecognition);
185        assert!(schema.input_types.contains(&"text_boxes".to_string()));
186        assert!(schema.output_types.contains(&"text_strings".to_string()));
187    }
188}
oar_ocr_core/domain/tasks/text_recognition.rs

oar_ocr_core/domain/tasks/
text_recognition.rs
