oar_ocr/core/traits/
granular.rs

1//! Granular traits for composable predictor components.
2//!
3//! This module provides granular traits that separate the concerns of the StandardPredictor
4//! trait, making it easier to compose, test, and extend individual pipeline components.
5//!
6//! The design focuses on practical composability and clean separation of concerns:
7//! - **ImageReader**: Handles I/O operations (loading images from files/memory)
8//! - **Preprocessor**: Handles image preprocessing (resize, normalize, tensor conversion)
9//! - **InferenceEngine**: Handles model inference (ONNX, TensorRT, etc.)
10//! - **Postprocessor**: Handles result processing (decoding, formatting, filtering)
11//!
12//! # Architecture
13//!
14//! ```text
15//! ┌─────────────┐    ┌─────────────┐    ┌─────────────┐    ┌─────────────-┐
16//! │ImageReader  │───▶│Preprocessor │───▶│InferenceEng │───▶│Postprocessor │
17//! │             │    │             │    │             │    │              │
18//! │• read_images│    │• preprocess │    │• infer      │    │• postprocess │
19//! │• validate   │    │• validate   │    │• engine_info│    │• empty_result│
20//! └─────────────┘    └─────────────┘    └─────────────┘    └─────────────-┘
21//! ```
22//!
23//! # Examples
24//!
25//! ```rust,no_run
26//! use oar_ocr::core::traits::granular::Preprocessor;
27//! use image::RgbImage;
28//!
29//! // Example of how you would implement and use a custom preprocessor
30//! // (This is a conceptual example - actual implementations would be in separate modules)
31//!
32//! # #[derive(Debug)]
33//! # struct MyPreprocessor;
34//! # #[derive(Debug)]
35//! # struct MyConfig { brightness_factor: f32 }
36//! # impl Preprocessor for MyPreprocessor {
37//! #     type Config = MyConfig;
38//! #     type Output = Vec<RgbImage>;
39//! #     fn preprocess(&self, images: Vec<RgbImage>, config: Option<&Self::Config>) -> Result<Self::Output, oar_ocr::core::OCRError> {
40//! #         Ok(images)
41//! #     }
42//! #     fn preprocessing_info(&self) -> String { "MyPreprocessor".to_string() }
43//! # }
44//!
45//! let preprocessor = MyPreprocessor;
46//! let test_image = RgbImage::new(32, 32);
47//! let images = vec![test_image];
48//! let config = MyConfig { brightness_factor: 1.5 };
49//!
50//! // Apply preprocessing
51//! let result = preprocessor.preprocess(images, Some(&config));
52//! assert!(result.is_ok());
53//! ```
54
55use crate::core::traits::StandardPredictor;
56use crate::core::{BatchData, OCRError};
57use image::RgbImage;
58use std::fmt::Debug;
59
60/// Trait for image reading and I/O operations.
61///
62/// This trait handles loading images from various sources (file paths, URLs, memory)
63/// and converting them to a standard RGB format for processing.
64pub trait ImageReader: Send + Sync + Debug {
65    /// Read images from file paths.
66    ///
67    /// # Arguments
68    ///
69    /// * `paths` - Iterator over file paths to read
70    ///
71    /// # Returns
72    ///
73    /// Vector of loaded RGB images or an error
74    fn read_images<'a>(
75        &self,
76        paths: impl Iterator<Item = &'a str>,
77    ) -> Result<Vec<RgbImage>, OCRError>;
78
79    /// Validate that images can be read from the given paths.
80    ///
81    /// # Arguments
82    ///
83    /// * `paths` - Iterator over file paths to validate
84    ///
85    /// # Returns
86    ///
87    /// Result indicating success or validation error
88    fn validate_paths<'a>(&self, paths: impl Iterator<Item = &'a str>) -> Result<(), OCRError> {
89        // Default implementation: try to read first path
90        if let Some(path) = paths.into_iter().next() {
91            let images = self.read_images(std::iter::once(path))?;
92            if images.is_empty() {
93                return Err(OCRError::InvalidInput {
94                    message: "Failed to read any images".to_string(),
95                });
96            }
97        }
98        Ok(())
99    }
100}
101
102/// Trait for image preprocessing operations.
103///
104/// This trait handles transforming raw images into the format required by
105/// the inference engine (resizing, normalization, tensor conversion, etc.).
106pub trait Preprocessor: Send + Sync + Debug {
107    /// Configuration type for preprocessing
108    type Config: Send + Sync + Debug;
109
110    /// Output type after preprocessing
111    type Output: Send + Sync + Debug;
112
113    /// Preprocess input images into inference-ready format.
114    ///
115    /// # Arguments
116    ///
117    /// * `images` - Input images to preprocess
118    /// * `config` - Optional configuration for preprocessing
119    ///
120    /// # Returns
121    ///
122    /// Preprocessed output ready for inference or an error
123    fn preprocess(
124        &self,
125        images: Vec<RgbImage>,
126        config: Option<&Self::Config>,
127    ) -> Result<Self::Output, OCRError>;
128
129    /// Get information about the preprocessing requirements.
130    ///
131    /// # Returns
132    ///
133    /// String describing preprocessing requirements (input size, format, etc.)
134    fn preprocessing_info(&self) -> String {
135        "Generic preprocessing".to_string()
136    }
137
138    /// Validate that the input images are suitable for preprocessing.
139    ///
140    /// # Arguments
141    ///
142    /// * `images` - Input images to validate
143    ///
144    /// # Returns
145    ///
146    /// Result indicating success or validation error
147    fn validate_input(&self, images: &[RgbImage]) -> Result<(), OCRError> {
148        if images.is_empty() {
149            return Err(OCRError::InvalidInput {
150                message: "No images provided for preprocessing".to_string(),
151            });
152        }
153        Ok(())
154    }
155}
156
157/// Trait for inference engine operations.
158///
159/// This trait handles running the actual model inference, whether through
160/// ONNX Runtime, TensorRT, PyTorch, or other backends.
161pub trait InferenceEngine: Send + Sync + Debug {
162    /// Input type for inference (typically a tensor)
163    type Input: Send + Sync + Debug;
164
165    /// Output type from inference (typically a tensor)
166    type Output: Send + Sync + Debug;
167
168    /// Perform inference on preprocessed input.
169    ///
170    /// # Arguments
171    ///
172    /// * `input` - Preprocessed input ready for inference
173    ///
174    /// # Returns
175    ///
176    /// Raw inference output or an error
177    fn infer(&self, input: &Self::Input) -> Result<Self::Output, OCRError>;
178
179    /// Get information about the inference engine.
180    ///
181    /// # Returns
182    ///
183    /// String describing the inference engine (model type, backend, etc.)
184    fn engine_info(&self) -> String;
185
186    /// Validate that the input is suitable for inference.
187    ///
188    /// # Arguments
189    ///
190    /// * `input` - Input to validate
191    ///
192    /// # Returns
193    ///
194    /// Result indicating success or validation error
195    fn validate_inference_input(&self, _input: &Self::Input) -> Result<(), OCRError> {
196        // Default implementation - basic validation
197        Ok(())
198    }
199}
200
201/// Trait for postprocessing operations.
202///
203/// This trait handles converting raw inference outputs into meaningful results
204/// (decoding, filtering, formatting, etc.).
205pub trait Postprocessor: Send + Sync + Debug {
206    /// Configuration type for postprocessing
207    type Config: Send + Sync + Debug;
208
209    /// Input type from inference engine
210    type InferenceOutput: Send + Sync + Debug;
211
212    /// Preprocessed data type for context
213    type PreprocessOutput: Send + Sync + Debug;
214
215    /// Final result type after postprocessing
216    type Result: Send + Sync + Debug;
217
218    /// Postprocess inference output into final results.
219    ///
220    /// # Arguments
221    ///
222    /// * `inference_output` - Raw output from inference engine
223    /// * `preprocess_output` - Optional preprocessed data for context
224    /// * `batch_data` - Batch metadata
225    /// * `raw_images` - Original input images
226    /// * `config` - Optional configuration for postprocessing
227    ///
228    /// # Returns
229    ///
230    /// Final processed result or an error
231    fn postprocess(
232        &self,
233        inference_output: Self::InferenceOutput,
234        preprocess_output: Option<&Self::PreprocessOutput>,
235        batch_data: &BatchData,
236        raw_images: Vec<RgbImage>,
237        config: Option<&Self::Config>,
238    ) -> crate::core::OcrResult<Self::Result>;
239
240    /// Create an empty result for when no input is provided.
241    ///
242    /// # Returns
243    ///
244    /// Empty result instance
245    fn empty_result(&self) -> Result<Self::Result, OCRError>;
246
247    /// Get information about the postprocessing operations.
248    ///
249    /// # Returns
250    ///
251    /// String describing postprocessing operations
252    fn postprocessing_info(&self) -> String {
253        "Generic postprocessing".to_string()
254    }
255}
256
257/// A modular predictor that composes granular components.
258///
259/// This struct demonstrates how to build a complete predictor using the granular traits.
260/// It provides the same interface as StandardPredictor but with composable components.
261#[derive(Debug)]
262pub struct ModularPredictor<R, P, I, O> {
263    /// Image reader component
264    pub image_reader: R,
265    /// Preprocessor component
266    pub preprocessor: P,
267    /// Inference engine component
268    pub inference_engine: I,
269    /// Postprocessor component
270    pub postprocessor: O,
271}
272
273impl<R, P, I, O> ModularPredictor<R, P, I, O>
274where
275    R: ImageReader,
276    P: Preprocessor,
277    I: InferenceEngine<Input = P::Output>,
278    O: Postprocessor<InferenceOutput = I::Output, PreprocessOutput = P::Output>,
279{
280    /// Create a new modular predictor with the given components.
281    pub fn new(image_reader: R, preprocessor: P, inference_engine: I, postprocessor: O) -> Self {
282        Self {
283            image_reader,
284            preprocessor,
285            inference_engine,
286            postprocessor,
287        }
288    }
289}
290
291impl<R, P, I, O> StandardPredictor for ModularPredictor<R, P, I, O>
292where
293    R: ImageReader,
294    P: Preprocessor,
295    I: InferenceEngine<Input = P::Output>,
296    O: Postprocessor<InferenceOutput = I::Output, PreprocessOutput = P::Output, Config = P::Config>,
297{
298    type Config = P::Config;
299    type Result = O::Result;
300    type PreprocessOutput = P::Output;
301    type InferenceOutput = I::Output;
302
303    fn read_images<'a>(
304        &self,
305        paths: impl Iterator<Item = &'a str>,
306    ) -> Result<Vec<RgbImage>, OCRError> {
307        self.image_reader.read_images(paths)
308    }
309
310    fn preprocess(
311        &self,
312        images: Vec<RgbImage>,
313        config: Option<&Self::Config>,
314    ) -> Result<Self::PreprocessOutput, OCRError> {
315        self.preprocessor.preprocess(images, config)
316    }
317
318    fn infer(&self, input: &Self::PreprocessOutput) -> Result<Self::InferenceOutput, OCRError> {
319        self.inference_engine.infer(input)
320    }
321
322    fn postprocess(
323        &self,
324        output: Self::InferenceOutput,
325        preprocessed: &Self::PreprocessOutput,
326        batch_data: &BatchData,
327        raw_images: Vec<RgbImage>,
328        config: Option<&Self::Config>,
329    ) -> crate::core::OcrResult<Self::Result> {
330        self.postprocessor
331            .postprocess(output, Some(preprocessed), batch_data, raw_images, config)
332    }
333
334    fn empty_result(&self) -> crate::core::OcrResult<Self::Result> {
335        self.postprocessor.empty_result()
336    }
337}
oar_ocr/core/traits/granular.rs

oar_ocr/core/traits/
granular.rs
