oar_ocr_core/domain/tasks/
document_rectification.rs

1//! Concrete task implementations for document rectification.
2//!
3//! This module provides the document rectification task that corrects distortions in document images.
4
5use super::validation::{ensure_images_with, ensure_non_empty_images};
6use crate::core::OCRError;
7use crate::core::config::{ConfigError, ConfigValidator};
8use crate::core::traits::TaskDefinition;
9use crate::core::traits::task::{ImageTaskInput, Task, TaskSchema, TaskType};
10use image::RgbImage;
11use serde::{Deserialize, Serialize};
12
13/// Configuration for document rectification task.
14#[derive(Debug, Clone, Serialize, Deserialize)]
15pub struct DocumentRectificationConfig {
16    /// Input shape for the rectification model [channels, height, width]
17    pub rec_image_shape: [usize; 3],
18}
19
20impl Default for DocumentRectificationConfig {
21    fn default() -> Self {
22        Self {
23            rec_image_shape: [3, 0, 0],
24        }
25    }
26}
27
28impl ConfigValidator for DocumentRectificationConfig {
29    fn validate(&self) -> Result<(), ConfigError> {
30        if self.rec_image_shape[0] == 0 {
31            return Err(ConfigError::InvalidConfig {
32                message: "rec_image_shape channels must be greater than 0".to_string(),
33            });
34        }
35
36        let height = self.rec_image_shape[1];
37        let width = self.rec_image_shape[2];
38        if (height == 0) ^ (width == 0) {
39            return Err(ConfigError::InvalidConfig {
40                message: "rec_image_shape height and width must both be set or both be zero"
41                    .to_string(),
42            });
43        }
44
45        Ok(())
46    }
47
48    fn get_defaults() -> Self {
49        Self::default()
50    }
51}
52
53/// Output from document rectification task.
54#[derive(Debug, Clone)]
55pub struct DocumentRectificationOutput {
56    /// Rectified images
57    pub rectified_images: Vec<RgbImage>,
58}
59
60impl DocumentRectificationOutput {
61    /// Creates an empty document rectification output.
62    pub fn empty() -> Self {
63        Self {
64            rectified_images: Vec::new(),
65        }
66    }
67
68    /// Creates a document rectification output with the given capacity.
69    pub fn with_capacity(capacity: usize) -> Self {
70        Self {
71            rectified_images: Vec::with_capacity(capacity),
72        }
73    }
74}
75
76impl TaskDefinition for DocumentRectificationOutput {
77    const TASK_NAME: &'static str = "document_rectification";
78    const TASK_DOC: &'static str = "Document rectification/unwarp";
79
80    fn empty() -> Self {
81        DocumentRectificationOutput::empty()
82    }
83}
84
85/// Document rectification task implementation.
86#[derive(Debug, Default)]
87pub struct DocumentRectificationTask {
88    _config: DocumentRectificationConfig,
89}
90
91impl DocumentRectificationTask {
92    /// Creates a new document rectification task.
93    pub fn new(config: DocumentRectificationConfig) -> Self {
94        Self { _config: config }
95    }
96}
97
98impl Task for DocumentRectificationTask {
99    type Config = DocumentRectificationConfig;
100    type Input = ImageTaskInput;
101    type Output = DocumentRectificationOutput;
102
103    fn task_type(&self) -> TaskType {
104        TaskType::DocumentRectification
105    }
106
107    fn schema(&self) -> TaskSchema {
108        TaskSchema::new(
109            TaskType::DocumentRectification,
110            vec!["image".to_string()],
111            vec!["rectified_image".to_string()],
112        )
113    }
114
115    fn validate_input(&self, input: &Self::Input) -> Result<(), OCRError> {
116        ensure_non_empty_images(
117            &input.images,
118            "No images provided for document rectification",
119        )?;
120
121        Ok(())
122    }
123
124    fn validate_output(&self, output: &Self::Output) -> Result<(), OCRError> {
125        // Validate that we have rectified images
126        ensure_images_with(
127            &output.rectified_images,
128            "No rectified images in output",
129            |idx, width, height| {
130                format!(
131                    "Invalid rectified image dimensions for item {idx}: width={width}, height={height} must be positive. Please verify the rectification output."
132                )
133            },
134        )?;
135
136        Ok(())
137    }
138
139    fn empty_output(&self) -> Self::Output {
140        DocumentRectificationOutput::empty()
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use super::*;
147
148    #[test]
149    fn test_document_rectification_task_creation() {
150        let task = DocumentRectificationTask::default();
151        assert_eq!(task.task_type(), TaskType::DocumentRectification);
152    }
153
154    #[test]
155    fn test_input_validation() {
156        let task = DocumentRectificationTask::default();
157
158        // Empty images should fail
159        let empty_input = ImageTaskInput::new(vec![]);
160        assert!(task.validate_input(&empty_input).is_err());
161
162        // Valid images should pass
163        let valid_input = ImageTaskInput::new(vec![RgbImage::new(100, 100)]);
164        assert!(task.validate_input(&valid_input).is_ok());
165    }
166
167    #[test]
168    fn test_output_validation() {
169        let task = DocumentRectificationTask::default();
170
171        // Valid output should pass
172        let output = DocumentRectificationOutput {
173            rectified_images: vec![RgbImage::new(100, 100)],
174        };
175        assert!(task.validate_output(&output).is_ok());
176
177        // Empty output should fail
178        let empty_output = DocumentRectificationOutput::empty();
179        assert!(task.validate_output(&empty_output).is_err());
180    }
181
182    #[test]
183    fn test_schema() {
184        let task = DocumentRectificationTask::default();
185        let schema = task.schema();
186        assert_eq!(schema.task_type, TaskType::DocumentRectification);
187        assert!(schema.input_types.contains(&"image".to_string()));
188        assert!(schema.output_types.contains(&"rectified_image".to_string()));
189    }
190}
oar_ocr_core/domain/tasks/document_rectification.rs

oar_ocr_core/domain/tasks/
document_rectification.rs
