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, 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 validate_input(&self, input: &Self::Input) -> Result<(), OCRError> {
108        ensure_non_empty_images(
109            &input.images,
110            "No images provided for document rectification",
111        )?;
112
113        Ok(())
114    }
115
116    fn validate_output(&self, output: &Self::Output) -> Result<(), OCRError> {
117        // Validate that we have rectified images
118        ensure_images_with(
119            &output.rectified_images,
120            "No rectified images in output",
121            |idx, width, height| {
122                format!(
123                    "Invalid rectified image dimensions for item {idx}: width={width}, height={height} must be positive. Please verify the rectification output."
124                )
125            },
126        )?;
127
128        Ok(())
129    }
130
131    fn empty_output(&self) -> Self::Output {
132        DocumentRectificationOutput::empty()
133    }
134}
135
136#[cfg(test)]
137mod tests {
138    use super::*;
139
140    #[test]
141    fn test_document_rectification_task_creation() {
142        let task = DocumentRectificationTask::default();
143        assert_eq!(task.task_type(), TaskType::DocumentRectification);
144    }
145
146    #[test]
147    fn test_input_validation() {
148        let task = DocumentRectificationTask::default();
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, 100)]);
156        assert!(task.validate_input(&valid_input).is_ok());
157    }
158
159    #[test]
160    fn test_output_validation() {
161        let task = DocumentRectificationTask::default();
162
163        // Valid output should pass
164        let output = DocumentRectificationOutput {
165            rectified_images: vec![RgbImage::new(100, 100)],
166        };
167        assert!(task.validate_output(&output).is_ok());
168
169        // Empty output should fail
170        let empty_output = DocumentRectificationOutput::empty();
171        assert!(task.validate_output(&empty_output).is_err());
172    }
173}
oar_ocr_core/domain/tasks/document_rectification.rs

oar_ocr_core/domain/tasks/
document_rectification.rs
