oar_ocr_core/predictors/
layout_detection.rs

1//! Layout Detection Predictor
2//!
3//! This module provides a high-level API for document layout detection.
4
5use super::builder::PredictorBuilderState;
6use crate::TaskPredictorBuilder;
7use crate::core::OcrResult;
8use crate::core::errors::OCRError;
9use crate::core::traits::OrtConfigurable;
10use crate::core::traits::task::ImageTaskInput;
11use crate::domain::adapters::LayoutDetectionAdapterBuilder;
12use crate::domain::tasks::layout_detection::{LayoutDetectionConfig, LayoutDetectionTask};
13use crate::predictors::TaskPredictorCore;
14use image::RgbImage;
15
16/// Layout detection prediction result
17#[derive(Debug, Clone)]
18pub struct LayoutDetectionResult {
19    /// Detected layout elements for each input image
20    pub elements: Vec<Vec<crate::domain::tasks::layout_detection::LayoutDetectionElement>>,
21    /// Whether elements are already sorted by reading order (e.g., from PP-DocLayoutV2)
22    ///
23    /// When `true`, downstream consumers can skip reading order sorting algorithms
24    /// as the elements are already in the correct reading order based on model output.
25    pub is_reading_order_sorted: bool,
26}
27
28/// Layout detection predictor
29pub struct LayoutDetectionPredictor {
30    core: TaskPredictorCore<LayoutDetectionTask>,
31}
32
33impl LayoutDetectionPredictor {
34    pub fn builder() -> LayoutDetectionPredictorBuilder {
35        LayoutDetectionPredictorBuilder::new()
36    }
37
38    /// Predict layout elements in the given images.
39    pub fn predict(&self, images: Vec<RgbImage>) -> OcrResult<LayoutDetectionResult> {
40        let input = ImageTaskInput::new(images);
41        let output = self.core.predict(input)?;
42        Ok(LayoutDetectionResult {
43            elements: output.elements,
44            is_reading_order_sorted: output.is_reading_order_sorted,
45        })
46    }
47}
48
49#[derive(TaskPredictorBuilder)]
50#[builder(config = LayoutDetectionConfig)]
51pub struct LayoutDetectionPredictorBuilder {
52    state: PredictorBuilderState<LayoutDetectionConfig>,
53    model_name: Option<String>,
54}
55
56impl LayoutDetectionPredictorBuilder {
57    pub fn new() -> Self {
58        Self {
59            state: PredictorBuilderState::new(LayoutDetectionConfig::default()),
60            model_name: None,
61        }
62    }
63
64    /// Creates a builder with PP-StructureV3 default class thresholds.
65    pub fn with_pp_structurev3_thresholds() -> Self {
66        Self {
67            state: PredictorBuilderState::new(
68                LayoutDetectionConfig::with_pp_structurev3_thresholds(),
69            ),
70            model_name: None,
71        }
72    }
73
74    pub fn model_name(mut self, name: impl Into<String>) -> Self {
75        self.model_name = Some(name.into());
76        self
77    }
78
79    pub fn score_threshold(mut self, threshold: f32) -> Self {
80        self.state.config_mut().score_threshold = threshold;
81        self
82    }
83
84    pub fn build(
85        self,
86        model_source: impl Into<crate::core::ModelSource>,
87    ) -> OcrResult<LayoutDetectionPredictor> {
88        let (config, ort_config) = self.state.into_parts();
89        let mut adapter_builder = LayoutDetectionAdapterBuilder::new().task_config(config.clone());
90
91        // Set model configuration if model_name was provided
92        if let Some(model_name) = self.model_name {
93            let model_config = Self::get_model_config(&model_name)?;
94            adapter_builder = adapter_builder.model_config(model_config);
95        }
96
97        if let Some(ort_cfg) = ort_config {
98            adapter_builder = adapter_builder.with_ort_config(ort_cfg);
99        }
100
101        let adapter = super::build_adapter(adapter_builder, model_source)?;
102        let task = LayoutDetectionTask::new(config.clone());
103        Ok(LayoutDetectionPredictor {
104            core: TaskPredictorCore::new(adapter, task, config),
105        })
106    }
107
108    /// Supported layout model names
109    const SUPPORTED_MODELS: &'static [&'static str] = &[
110        "picodet_layout_1x",
111        "picodet_layout_1x_table",
112        "picodet_s_layout_3cls",
113        "picodet_l_layout_3cls",
114        "picodet_s_layout_17cls",
115        "picodet_l_layout_17cls",
116        "rtdetr_h_layout_3cls",
117        "rt_detr_h_layout_3cls",
118        "rtdetr_h_layout_17cls",
119        "rt_detr_h_layout_17cls",
120        "pp_docblocklayout",
121        "pp_doclayout_s",
122        "pp_doclayout_m",
123        "pp_doclayout_l",
124        "pp_doclayout_plus_l",
125        "pp_doclayoutv2",
126        "pp_doclayout_v2",
127        "pp_doclayoutv3",
128        "pp_doclayout_v3",
129    ];
130
131    fn get_model_config(model_name: &str) -> OcrResult<crate::domain::adapters::LayoutModelConfig> {
132        use crate::domain::adapters::LayoutModelConfig;
133
134        let normalized = model_name.to_lowercase().replace('-', "_");
135        let config = match normalized.as_str() {
136            "picodet_layout_1x" => LayoutModelConfig::picodet_layout_1x(),
137            "picodet_layout_1x_table" => LayoutModelConfig::picodet_layout_1x_table(),
138            "picodet_s_layout_3cls" => LayoutModelConfig::picodet_s_layout_3cls(),
139            "picodet_l_layout_3cls" => LayoutModelConfig::picodet_l_layout_3cls(),
140            "picodet_s_layout_17cls" => LayoutModelConfig::picodet_s_layout_17cls(),
141            "picodet_l_layout_17cls" => LayoutModelConfig::picodet_l_layout_17cls(),
142            "rtdetr_h_layout_3cls" | "rt_detr_h_layout_3cls" => {
143                LayoutModelConfig::rtdetr_h_layout_3cls()
144            }
145            "rtdetr_h_layout_17cls" | "rt_detr_h_layout_17cls" => {
146                LayoutModelConfig::rtdetr_h_layout_17cls()
147            }
148            "pp_docblocklayout" => LayoutModelConfig::pp_docblocklayout(),
149            "pp_doclayout_s" => LayoutModelConfig::pp_doclayout_s(),
150            "pp_doclayout_m" => LayoutModelConfig::pp_doclayout_m(),
151            "pp_doclayout_l" => LayoutModelConfig::pp_doclayout_l(),
152            "pp_doclayout_plus_l" => LayoutModelConfig::pp_doclayout_plus_l(),
153            "pp_doclayoutv2" | "pp_doclayout_v2" => LayoutModelConfig::pp_doclayoutv2(),
154            "pp_doclayoutv3" | "pp_doclayout_v3" => LayoutModelConfig::pp_doclayoutv3(),
155            _ => {
156                return Err(OCRError::ConfigError {
157                    message: format!(
158                        "Unknown model name: '{}'. Supported models: {}",
159                        model_name,
160                        Self::SUPPORTED_MODELS.join(", ")
161                    ),
162                });
163            }
164        };
165
166        Ok(config)
167    }
168}
169
170impl Default for LayoutDetectionPredictorBuilder {
171    fn default() -> Self {
172        Self::new()
173    }
174}
oar_ocr_core/predictors/layout_detection.rs

oar_ocr_core/predictors/
layout_detection.rs
