oar_ocr_core/predictors/
table_classification.rs

1//! Table Classification Predictor
2//!
3//! This module provides a high-level API for table classification (wired vs wireless tables).
4
5use super::builder::PredictorBuilderState;
6use crate::TaskPredictorBuilder;
7use crate::core::OcrResult;
8use crate::core::traits::OrtConfigurable;
9use crate::core::traits::task::ImageTaskInput;
10use crate::domain::adapters::TableClassificationAdapterBuilder;
11use crate::domain::tasks::document_orientation::Classification;
12use crate::domain::tasks::table_classification::{
13    TableClassificationConfig, TableClassificationTask,
14};
15use crate::predictors::TaskPredictorCore;
16use image::RgbImage;
17
18/// Table classification prediction result
19#[derive(Debug, Clone)]
20pub struct TableClassificationResult {
21    /// Classification results for each input image
22    pub classifications: Vec<Vec<Classification>>,
23}
24
25/// Table classification predictor
26pub struct TableClassificationPredictor {
27    core: TaskPredictorCore<TableClassificationTask>,
28}
29
30impl TableClassificationPredictor {
31    /// Create a new builder for the table classification predictor
32    pub fn builder() -> TableClassificationPredictorBuilder {
33        TableClassificationPredictorBuilder::new()
34    }
35
36    /// Predict table classifications in the given images.
37    pub fn predict(&self, images: Vec<RgbImage>) -> OcrResult<TableClassificationResult> {
38        let input = ImageTaskInput::new(images);
39        let output = self.core.predict(input)?;
40        Ok(TableClassificationResult {
41            classifications: output.classifications,
42        })
43    }
44}
45
46/// Builder for table classification predictor
47#[derive(TaskPredictorBuilder)]
48#[builder(config = TableClassificationConfig)]
49pub struct TableClassificationPredictorBuilder {
50    state: PredictorBuilderState<TableClassificationConfig>,
51    input_shape: (u32, u32),
52}
53
54impl TableClassificationPredictorBuilder {
55    /// Create a new builder with default configuration
56    pub fn new() -> Self {
57        Self {
58            state: PredictorBuilderState::new(TableClassificationConfig {
59                score_threshold: 0.5,
60                topk: 2,
61            }),
62            input_shape: (224, 224),
63        }
64    }
65
66    /// Set the score threshold
67    pub fn score_threshold(mut self, threshold: f32) -> Self {
68        self.state.config_mut().score_threshold = threshold;
69        self
70    }
71
72    /// Set the top-k predictions to return
73    pub fn topk(mut self, k: usize) -> Self {
74        self.state.config_mut().topk = k;
75        self
76    }
77
78    /// Set the model input shape (height, width)
79    pub fn input_shape(mut self, shape: (u32, u32)) -> Self {
80        self.input_shape = shape;
81        self
82    }
83
84    /// Build the table classification predictor
85    pub fn build(
86        self,
87        model_source: impl Into<crate::core::ModelSource>,
88    ) -> OcrResult<TableClassificationPredictor> {
89        let Self { state, input_shape } = self;
90        let (config, ort_config) = state.into_parts();
91        let mut adapter_builder = TableClassificationAdapterBuilder::new()
92            .with_config(config.clone())
93            .input_shape(input_shape);
94
95        if let Some(ort_cfg) = ort_config {
96            adapter_builder = adapter_builder.with_ort_config(ort_cfg);
97        }
98
99        let adapter = super::build_adapter(adapter_builder, model_source)?;
100        let task = TableClassificationTask::new(config.clone());
101
102        Ok(TableClassificationPredictor {
103            core: TaskPredictorCore::new(adapter, task, config),
104        })
105    }
106}
107
108impl Default for TableClassificationPredictorBuilder {
109    fn default() -> Self {
110        Self::new()
111    }
112}
oar_ocr_core/predictors/table_classification.rs

oar_ocr_core/predictors/
table_classification.rs
