oar_ocr_core/core/traits/
task.rs

1//! Task trait definitions for the OCR pipeline.
2//!
3//! This module defines the `Task` trait and related types that represent
4//! different OCR tasks (text detection, recognition, layout analysis, etc.).
5//! Tasks define input/output schemas and execution contracts.
6
7use crate::core::OCRError;
8use image::RgbImage;
9use serde::{Deserialize, Serialize};
10use std::fmt::Debug;
11use std::sync::Arc;
12
13// Generate TaskType enum from the central task registry
14crate::with_task_registry!(crate::impl_task_type_enum);
15
16/// Schema definition for task inputs and outputs.
17///
18/// This allows for validation that models produce outputs compatible
19/// with what downstream tasks expect.
20#[derive(Debug, Clone, Serialize, Deserialize)]
21pub struct TaskSchema {
22    /// Task type
23    pub task_type: TaskType,
24    /// Expected input types (e.g., "image", "text_boxes")
25    pub input_types: Vec<String>,
26    /// Expected output types (e.g., "text_boxes", "text_strings")
27    pub output_types: Vec<String>,
28    /// Optional metadata schema
29    pub metadata_schema: Option<String>,
30}
31
32impl TaskSchema {
33    /// Creates a new task schema.
34    pub fn new(task_type: TaskType, input_types: Vec<String>, output_types: Vec<String>) -> Self {
35        Self {
36            task_type,
37            input_types,
38            output_types,
39            metadata_schema: None,
40        }
41    }
42
43    /// Validates that this schema is compatible with another schema.
44    ///
45    /// Returns true if the output types of this schema match the input types
46    /// of the target schema.
47    pub fn is_compatible_with(&self, target: &TaskSchema) -> bool {
48        // Check if any of our output types match any of target's input types
49        self.output_types
50            .iter()
51            .any(|output| target.input_types.contains(output))
52    }
53}
54
55/// Core trait for OCR tasks.
56///
57/// Tasks represent distinct operations in the OCR pipeline (detection, recognition, etc.).
58/// Each task defines its input/output schema and can be executed with various model adapters.
59pub trait Task: Send + Sync + Debug {
60    /// Configuration type for this task
61    type Config: Send + Sync + Debug + Clone;
62
63    /// Input type for this task
64    type Input: Send + Sync + Debug;
65
66    /// Output type from this task
67    type Output: Send + Sync + Debug;
68
69    /// Returns the task type identifier.
70    fn task_type(&self) -> TaskType;
71
72    /// Returns the schema defining inputs and outputs for this task.
73    fn schema(&self) -> TaskSchema;
74
75    /// Validates that the given input is suitable for this task.
76    ///
77    /// # Arguments
78    ///
79    /// * `input` - The input to validate
80    ///
81    /// # Returns
82    ///
83    /// Result indicating success or validation error
84    fn validate_input(&self, input: &Self::Input) -> Result<(), OCRError>;
85
86    /// Validates that the given output matches the expected schema.
87    ///
88    /// # Arguments
89    ///
90    /// * `output` - The output to validate
91    ///
92    /// # Returns
93    ///
94    /// Result indicating success or validation error
95    fn validate_output(&self, output: &Self::Output) -> Result<(), OCRError>;
96
97    /// Returns an empty output instance for when no valid results are produced.
98    fn empty_output(&self) -> Self::Output;
99
100    /// Returns a human-readable description of this task.
101    fn description(&self) -> String {
102        format!("Task: {}", self.task_type().name())
103    }
104}
105
106/// A task runner that executes tasks using a model adapter.
107///
108/// This struct coordinates the execution of a task with a specific model,
109/// handling validation and error propagation.
110#[derive(Debug)]
111pub struct TaskRunner<T: Task> {
112    /// The task to execute
113    task: T,
114    /// Configuration for the task
115    config: T::Config,
116}
117
118impl<T: Task> TaskRunner<T> {
119    /// Creates a new task runner.
120    ///
121    /// # Arguments
122    ///
123    /// * `task` - The task to execute
124    /// * `config` - Configuration for the task
125    pub fn new(task: T, config: T::Config) -> Self {
126        Self { task, config }
127    }
128
129    /// Returns a reference to the task.
130    pub fn task(&self) -> &T {
131        &self.task
132    }
133
134    /// Returns a reference to the configuration.
135    pub fn config(&self) -> &T::Config {
136        &self.config
137    }
138
139    /// Returns the task type.
140    pub fn task_type(&self) -> TaskType {
141        self.task.task_type()
142    }
143
144    /// Validates input before execution.
145    pub fn validate_input(&self, input: &T::Input) -> Result<(), OCRError> {
146        self.task.validate_input(input)
147    }
148
149    /// Validates output after execution.
150    pub fn validate_output(&self, output: &T::Output) -> Result<(), OCRError> {
151        self.task.validate_output(output)
152    }
153}
154
155/// Common input type for image-based tasks.
156#[derive(Debug, Clone)]
157pub struct ImageTaskInput {
158    /// Input images
159    pub images: Vec<Arc<RgbImage>>,
160    /// Optional metadata per image
161    pub metadata: Vec<Option<String>>,
162}
163
164impl ImageTaskInput {
165    /// Creates a new image task input from owned images.
166    pub fn new(images: Vec<RgbImage>) -> Self {
167        let count = images.len();
168        Self {
169            images: images.into_iter().map(Arc::new).collect(),
170            metadata: vec![None; count],
171        }
172    }
173
174    /// Creates a new image task input from shared images.
175    pub fn from_arc_images(images: Vec<Arc<RgbImage>>) -> Self {
176        let count = images.len();
177        Self {
178            images,
179            metadata: vec![None; count],
180        }
181    }
182
183    /// Creates a new image task input with metadata.
184    pub fn with_metadata(images: Vec<RgbImage>, metadata: Vec<Option<String>>) -> Self {
185        Self {
186            images: images.into_iter().map(Arc::new).collect(),
187            metadata,
188        }
189    }
190
191    /// Converts shared images into owned images for model APIs that still take ownership.
192    ///
193    /// This avoids a copy when the image is uniquely owned and clones only when another
194    /// pipeline stage still holds the same image.
195    pub fn into_owned_images(self) -> Vec<RgbImage> {
196        self.images
197            .into_iter()
198            .map(|img| Arc::try_unwrap(img).unwrap_or_else(|img| (*img).clone()))
199            .collect()
200    }
201}
202
203#[cfg(test)]
204mod tests {
205    use super::*;
206    use std::sync::Arc;
207
208    #[test]
209    fn test_task_type_name() {
210        assert_eq!(TaskType::TextDetection.name(), "text_detection");
211        assert_eq!(TaskType::TextRecognition.name(), "text_recognition");
212    }
213
214    #[test]
215    fn test_schema_compatibility() {
216        let detection_schema = TaskSchema::new(
217            TaskType::TextDetection,
218            vec!["image".to_string()],
219            vec!["text_boxes".to_string()],
220        );
221
222        let recognition_schema = TaskSchema::new(
223            TaskType::TextRecognition,
224            vec!["text_boxes".to_string()],
225            vec!["text_strings".to_string()],
226        );
227
228        // Detection output (text_boxes) should be compatible with recognition input (text_boxes)
229        assert!(detection_schema.is_compatible_with(&recognition_schema));
230
231        // Recognition output (text_strings) is not compatible with detection input (image)
232        assert!(!recognition_schema.is_compatible_with(&detection_schema));
233    }
234
235    #[test]
236    fn test_image_task_input_creation() {
237        let images = vec![RgbImage::new(100, 100), RgbImage::new(200, 200)];
238        let input = ImageTaskInput::new(images.clone());
239
240        assert_eq!(input.images.len(), 2);
241        assert_eq!(input.metadata.len(), 2);
242        assert!(input.metadata.iter().all(|m| m.is_none()));
243    }
244
245    #[test]
246    fn test_image_task_input_from_owned() {
247        let images = vec![RgbImage::new(100, 100), RgbImage::new(200, 200)];
248        let input = ImageTaskInput::new(images);
249
250        assert_eq!(input.images.len(), 2);
251        assert_eq!(input.metadata.len(), 2);
252        assert!(input.metadata.iter().all(|m| m.is_none()));
253    }
254
255    #[test]
256    fn test_into_owned_images_reuses_unique_arc() {
257        let mut image = RgbImage::new(2, 1);
258        image.put_pixel(0, 0, image::Rgb([1, 2, 3]));
259        let input = ImageTaskInput::from_arc_images(vec![Arc::new(image)]);
260
261        let owned = input.into_owned_images();
262
263        assert_eq!(owned.len(), 1);
264        assert_eq!(owned[0].get_pixel(0, 0).0, [1, 2, 3]);
265    }
266
267    #[test]
268    fn test_into_owned_images_clones_when_arc_is_shared() {
269        let mut image = RgbImage::new(2, 1);
270        image.put_pixel(1, 0, image::Rgb([9, 8, 7]));
271        let shared = Arc::new(image);
272        let input = ImageTaskInput::from_arc_images(vec![Arc::clone(&shared)]);
273
274        let owned = input.into_owned_images();
275
276        assert_eq!(Arc::strong_count(&shared), 1);
277        assert_eq!(owned.len(), 1);
278        assert_eq!(owned[0].get_pixel(1, 0).0, [9, 8, 7]);
279        assert_eq!(shared.get_pixel(1, 0).0, [9, 8, 7]);
280    }
281}
oar_ocr_core/core/traits/task.rs

oar_ocr_core/core/traits/
task.rs
