oar_ocr_core/predictors/
document_rectification.rs

1//! Document Rectification Predictor
2//!
3//! This module provides a high-level API for document rectification (dewarp).
4
5use super::builder::PredictorBuilderState;
6use crate::TaskPredictorBuilder;
7use crate::core::OcrResult;
8use crate::core::traits::OrtConfigurable;
9use crate::core::traits::adapter::AdapterBuilder;
10use crate::core::traits::task::ImageTaskInput;
11use crate::domain::adapters::UVDocRectifierAdapterBuilder;
12use crate::domain::tasks::document_rectification::{
13    DocumentRectificationConfig, DocumentRectificationTask,
14};
15use crate::predictors::TaskPredictorCore;
16use image::RgbImage;
17
18/// Document rectification prediction result
19#[derive(Debug, Clone)]
20pub struct DocumentRectificationResult {
21    /// Rectified images
22    pub images: Vec<RgbImage>,
23}
24
25/// Document rectification predictor
26pub struct DocumentRectificationPredictor {
27    core: TaskPredictorCore<DocumentRectificationTask>,
28}
29
30impl DocumentRectificationPredictor {
31    pub fn builder() -> DocumentRectificationPredictorBuilder {
32        DocumentRectificationPredictorBuilder::new()
33    }
34
35    /// Predict document rectification for the given images.
36    pub fn predict(&self, images: Vec<RgbImage>) -> OcrResult<DocumentRectificationResult> {
37        let input = ImageTaskInput::new(images);
38        let output = self.core.predict(input)?;
39        Ok(DocumentRectificationResult {
40            images: output.rectified_images,
41        })
42    }
43}
44
45#[derive(TaskPredictorBuilder)]
46#[builder(config = DocumentRectificationConfig)]
47pub struct DocumentRectificationPredictorBuilder {
48    state: PredictorBuilderState<DocumentRectificationConfig>,
49}
50
51impl DocumentRectificationPredictorBuilder {
52    pub fn new() -> Self {
53        Self {
54            state: PredictorBuilderState::new(DocumentRectificationConfig {
55                rec_image_shape: [3, 0, 0],
56            }),
57        }
58    }
59
60    pub fn build(
61        self,
62        model_source: impl Into<crate::core::ModelSource>,
63    ) -> OcrResult<DocumentRectificationPredictor> {
64        let (config, ort_config) = self.state.into_parts();
65        let mut adapter_builder = UVDocRectifierAdapterBuilder::new().with_config(config.clone());
66
67        if let Some(ort_cfg) = ort_config {
68            adapter_builder = adapter_builder.with_ort_config(ort_cfg);
69        }
70
71        let adapter = super::build_adapter(adapter_builder, model_source)?;
72        let task = DocumentRectificationTask::new(config.clone());
73        Ok(DocumentRectificationPredictor {
74            core: TaskPredictorCore::new(adapter, task, config),
75        })
76    }
77}
78
79impl Default for DocumentRectificationPredictorBuilder {
80    fn default() -> Self {
81        Self::new()
82    }
83}
oar_ocr_core/predictors/document_rectification.rs

oar_ocr_core/predictors/
document_rectification.rs
