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 typed inputs, outputs, and validation contracts.
6
7use crate::core::OCRError;
8use image::RgbImage;
9use std::fmt::Debug;
10use std::sync::Arc;
11
12// Generate TaskType enum from the central task registry
13crate::with_task_registry!(crate::impl_task_type_enum);
14
15/// Core trait for OCR tasks.
16///
17/// Tasks represent distinct operations in the OCR pipeline (detection, recognition, etc.).
18/// Each task defines typed inputs and outputs and can be executed with model adapters.
19pub trait Task: Send + Sync + Debug {
20    /// Configuration type for this task
21    type Config: Send + Sync + Debug + Clone;
22
23    /// Input type for this task
24    type Input: Send + Sync + Debug;
25
26    /// Output type from this task
27    type Output: Send + Sync + Debug;
28
29    /// Returns the task type identifier.
30    fn task_type(&self) -> TaskType;
31
32    /// Validates that the given input is suitable for this task.
33    ///
34    /// # Arguments
35    ///
36    /// * `input` - The input to validate
37    ///
38    /// # Returns
39    ///
40    /// Result indicating success or validation error
41    fn validate_input(&self, input: &Self::Input) -> Result<(), OCRError>;
42
43    /// Validates that the given output is suitable for this task.
44    ///
45    /// # Arguments
46    ///
47    /// * `output` - The output to validate
48    ///
49    /// # Returns
50    ///
51    /// Result indicating success or validation error
52    fn validate_output(&self, output: &Self::Output) -> Result<(), OCRError>;
53
54    /// Returns an empty output instance for when no valid results are produced.
55    fn empty_output(&self) -> Self::Output;
56
57    /// Returns a human-readable description of this task.
58    fn description(&self) -> String {
59        format!("Task: {}", self.task_type().name())
60    }
61}
62
63/// Common input type for image-based tasks.
64#[derive(Debug, Clone)]
65pub struct ImageTaskInput {
66    /// Input images
67    pub images: Vec<Arc<RgbImage>>,
68    /// Optional metadata per image
69    pub metadata: Vec<Option<String>>,
70}
71
72impl ImageTaskInput {
73    /// Creates a new image task input from owned images.
74    pub fn new(images: Vec<RgbImage>) -> Self {
75        let count = images.len();
76        Self {
77            images: images.into_iter().map(Arc::new).collect(),
78            metadata: vec![None; count],
79        }
80    }
81
82    /// Creates a new image task input from shared images.
83    pub fn from_arc_images(images: Vec<Arc<RgbImage>>) -> Self {
84        let count = images.len();
85        Self {
86            images,
87            metadata: vec![None; count],
88        }
89    }
90
91    /// Creates a new image task input with metadata.
92    pub fn with_metadata(images: Vec<RgbImage>, metadata: Vec<Option<String>>) -> Self {
93        Self {
94            images: images.into_iter().map(Arc::new).collect(),
95            metadata,
96        }
97    }
98
99    /// Converts shared images into owned images for model APIs that still take ownership.
100    ///
101    /// This avoids a copy when the image is uniquely owned and clones only when another
102    /// pipeline stage still holds the same image.
103    pub fn into_owned_images(self) -> Vec<RgbImage> {
104        self.images
105            .into_iter()
106            .map(|img| Arc::try_unwrap(img).unwrap_or_else(|img| (*img).clone()))
107            .collect()
108    }
109}
110
111#[cfg(test)]
112mod tests {
113    use super::*;
114    use std::sync::Arc;
115
116    #[test]
117    fn test_task_type_name() {
118        assert_eq!(TaskType::TextDetection.name(), "text_detection");
119        assert_eq!(TaskType::TextRecognition.name(), "text_recognition");
120    }
121
122    #[test]
123    fn test_image_task_input_creation() {
124        let images = vec![RgbImage::new(100, 100), RgbImage::new(200, 200)];
125        let input = ImageTaskInput::new(images.clone());
126
127        assert_eq!(input.images.len(), 2);
128        assert_eq!(input.metadata.len(), 2);
129        assert!(input.metadata.iter().all(|m| m.is_none()));
130    }
131
132    #[test]
133    fn test_image_task_input_from_owned() {
134        let images = vec![RgbImage::new(100, 100), RgbImage::new(200, 200)];
135        let input = ImageTaskInput::new(images);
136
137        assert_eq!(input.images.len(), 2);
138        assert_eq!(input.metadata.len(), 2);
139        assert!(input.metadata.iter().all(|m| m.is_none()));
140    }
141
142    #[test]
143    fn test_into_owned_images_reuses_unique_arc() {
144        let mut image = RgbImage::new(2, 1);
145        image.put_pixel(0, 0, image::Rgb([1, 2, 3]));
146        let input = ImageTaskInput::from_arc_images(vec![Arc::new(image)]);
147
148        let owned = input.into_owned_images();
149
150        assert_eq!(owned.len(), 1);
151        assert_eq!(owned[0].get_pixel(0, 0).0, [1, 2, 3]);
152    }
153
154    #[test]
155    fn test_into_owned_images_clones_when_arc_is_shared() {
156        let mut image = RgbImage::new(2, 1);
157        image.put_pixel(1, 0, image::Rgb([9, 8, 7]));
158        let shared = Arc::new(image);
159        let input = ImageTaskInput::from_arc_images(vec![Arc::clone(&shared)]);
160
161        let owned = input.into_owned_images();
162
163        assert_eq!(Arc::strong_count(&shared), 1);
164        assert_eq!(owned.len(), 1);
165        assert_eq!(owned[0].get_pixel(1, 0).0, [9, 8, 7]);
166        assert_eq!(shared.get_pixel(1, 0).0, [9, 8, 7]);
167    }
168}
oar_ocr_core/core/traits/task.rs

oar_ocr_core/core/traits/
task.rs
