oar-ocr-core 0.6.3

Core types and predictors for oar-ocr
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

//! Table Classification Predictor
//!
//! This module provides a high-level API for table classification (wired vs wireless tables).

use super::builder::PredictorBuilderState;
use crate::TaskPredictorBuilder;
use crate::core::OcrResult;
use crate::core::traits::OrtConfigurable;
use crate::core::traits::adapter::AdapterBuilder;
use crate::core::traits::task::ImageTaskInput;
use crate::domain::adapters::TableClassificationAdapterBuilder;
use crate::domain::tasks::document_orientation::Classification;
use crate::domain::tasks::table_classification::{
    TableClassificationConfig, TableClassificationTask,
};
use crate::predictors::TaskPredictorCore;
use image::RgbImage;
use std::path::Path;

/// Table classification prediction result
#[derive(Debug, Clone)]
pub struct TableClassificationResult {
    /// Classification results for each input image
    pub classifications: Vec<Vec<Classification>>,
}

/// Table classification predictor
pub struct TableClassificationPredictor {
    core: TaskPredictorCore<TableClassificationTask>,
}

impl TableClassificationPredictor {
    /// Create a new builder for the table classification predictor
    pub fn builder() -> TableClassificationPredictorBuilder {
        TableClassificationPredictorBuilder::new()
    }

    /// Predict table classifications in the given images.
    pub fn predict(&self, images: Vec<RgbImage>) -> OcrResult<TableClassificationResult> {
        let input = ImageTaskInput::new(images);
        let output = self.core.predict(input)?;
        Ok(TableClassificationResult {
            classifications: output.classifications,
        })
    }
}

/// Builder for table classification predictor
#[derive(TaskPredictorBuilder)]
#[builder(config = TableClassificationConfig)]
pub struct TableClassificationPredictorBuilder {
    state: PredictorBuilderState<TableClassificationConfig>,
    input_shape: (u32, u32),
}

impl TableClassificationPredictorBuilder {
    /// Create a new builder with default configuration
    pub fn new() -> Self {
        Self {
            state: PredictorBuilderState::new(TableClassificationConfig {
                score_threshold: 0.5,
                topk: 2,
            }),
            input_shape: (224, 224),
        }
    }

    /// Set the score threshold
    pub fn score_threshold(mut self, threshold: f32) -> Self {
        self.state.config_mut().score_threshold = threshold;
        self
    }

    /// Set the top-k predictions to return
    pub fn topk(mut self, k: usize) -> Self {
        self.state.config_mut().topk = k;
        self
    }

    /// Set the model input shape (height, width)
    pub fn input_shape(mut self, shape: (u32, u32)) -> Self {
        self.input_shape = shape;
        self
    }

    /// Build the table classification predictor
    pub fn build<P: AsRef<Path>>(self, model_path: P) -> OcrResult<TableClassificationPredictor> {
        let Self { state, input_shape } = self;
        let (config, ort_config) = state.into_parts();
        let mut adapter_builder = TableClassificationAdapterBuilder::new()
            .with_config(config.clone())
            .input_shape(input_shape);

        if let Some(ort_cfg) = ort_config {
            adapter_builder = adapter_builder.with_ort_config(ort_cfg);
        }

        let adapter = Box::new(adapter_builder.build(model_path.as_ref())?);
        let task = TableClassificationTask::new(config.clone());

        Ok(TableClassificationPredictor {
            core: TaskPredictorCore::new(adapter, task, config),
        })
    }
}

impl Default for TableClassificationPredictorBuilder {
    fn default() -> Self {
        Self::new()
    }
}