oar_ocr_core/predictors/
text_detection.rs

1//! Text Detection Predictor
2//!
3//! This module provides a high-level API for text detection in images.
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::TextDetectionAdapterBuilder;
11use crate::domain::tasks::text_detection::{TextDetectionConfig, TextDetectionTask};
12use crate::predictors::TaskPredictorCore;
13use image::RgbImage;
14
15/// Text detection prediction result
16#[derive(Debug, Clone)]
17pub struct TextDetectionResult {
18    /// Detected text regions for each input image
19    pub detections: Vec<Vec<crate::domain::tasks::text_detection::Detection>>,
20}
21
22/// Text detection predictor
23pub struct TextDetectionPredictor {
24    core: TaskPredictorCore<TextDetectionTask>,
25}
26
27impl TextDetectionPredictor {
28    /// Create a new builder for the text detection predictor
29    pub fn builder() -> TextDetectionPredictorBuilder {
30        TextDetectionPredictorBuilder::new()
31    }
32
33    /// Predict text regions in the given images.
34    pub fn predict(&self, images: Vec<RgbImage>) -> OcrResult<TextDetectionResult> {
35        // Create task input
36        let input = ImageTaskInput::new(images);
37
38        // Use core predictor for validation and execution
39        let output = self.core.predict(input)?;
40
41        Ok(TextDetectionResult {
42            detections: output.detections,
43        })
44    }
45}
46
47/// Builder for text detection predictor
48#[derive(TaskPredictorBuilder)]
49#[builder(config = TextDetectionConfig)]
50pub struct TextDetectionPredictorBuilder {
51    state: PredictorBuilderState<TextDetectionConfig>,
52}
53
54impl TextDetectionPredictorBuilder {
55    /// Create a new builder with default configuration
56    pub fn new() -> Self {
57        Self {
58            state: PredictorBuilderState::new(TextDetectionConfig {
59                score_threshold: 0.3,
60                box_threshold: 0.6,
61                unclip_ratio: 1.5,
62                max_candidates: 1000,
63                limit_side_len: None,
64                limit_type: None,
65                max_side_len: None,
66            }),
67        }
68    }
69
70    /// Set the score threshold
71    pub fn score_threshold(mut self, threshold: f32) -> Self {
72        self.state.config_mut().score_threshold = threshold;
73        self
74    }
75
76    /// Set the box threshold
77    pub fn box_threshold(mut self, threshold: f32) -> Self {
78        self.state.config_mut().box_threshold = threshold;
79        self
80    }
81
82    /// Set the unclip ratio
83    pub fn unclip_ratio(mut self, ratio: f32) -> Self {
84        self.state.config_mut().unclip_ratio = ratio;
85        self
86    }
87
88    /// Set the maximum candidates
89    pub fn max_candidates(mut self, max: usize) -> Self {
90        self.state.config_mut().max_candidates = max;
91        self
92    }
93
94    /// Build the text detection predictor
95    pub fn build(
96        self,
97        model_source: impl Into<crate::core::ModelSource>,
98    ) -> OcrResult<TextDetectionPredictor> {
99        let (config, ort_config) = self.state.into_parts();
100        let mut adapter_builder = TextDetectionAdapterBuilder::new().with_config(config.clone());
101
102        if let Some(ort_cfg) = ort_config {
103            adapter_builder = adapter_builder.with_ort_config(ort_cfg);
104        }
105
106        let adapter = super::build_adapter(adapter_builder, model_source)?;
107        let task = TextDetectionTask::new(config.clone());
108
109        Ok(TextDetectionPredictor {
110            core: TaskPredictorCore::new(adapter, task, config),
111        })
112    }
113}
114
115impl Default for TextDetectionPredictorBuilder {
116    fn default() -> Self {
117        Self::new()
118    }
119}
oar_ocr_core/predictors/text_detection.rs

oar_ocr_core/predictors/
text_detection.rs
