Skip to main content

oar_ocr_core/models/detection/
pp_doclayout.rs

1//! PP-DocLayout Detection Model
2//!
3//! This module provides PP-DocLayout-specific types and configurations.
4//! PP-DocLayout is specialized for document layout detection.
5//! The actual implementation is shared via the generic ScaleAwareDetectorModel.
6
7use super::scale_aware_detector::{
8    ScaleAwareDetectorInferenceMode, ScaleAwareDetectorModel, ScaleAwareDetectorModelBuilder,
9    ScaleAwareDetectorModelOutput, ScaleAwareDetectorPostprocessConfig,
10    ScaleAwareDetectorPreprocessConfig,
11};
12use crate::core::OCRError;
13use crate::core::inference::OrtInfer;
14
15/// Preprocessing configuration for PP-DocLayout model.
16///
17/// This is a type alias for the generic configuration with PP-DocLayout defaults.
18pub type PPDocLayoutPreprocessConfig = ScaleAwareDetectorPreprocessConfig;
19
20/// Postprocessing configuration for PP-DocLayout model.
21pub type PPDocLayoutPostprocessConfig = ScaleAwareDetectorPostprocessConfig;
22
23/// Output from PP-DocLayout model.
24pub type PPDocLayoutModelOutput = ScaleAwareDetectorModelOutput;
25
26/// PP-DocLayout document layout detection model.
27///
28/// This is a specialized configuration of the generic ScaleAwareDetectorModel
29/// that uses ScaleFactorAndImageShape inference mode (requires both scale_factor and im_shape).
30pub type PPDocLayoutModel = ScaleAwareDetectorModel;
31
32/// Builder for PP-DocLayout model.
33#[derive(Debug, Default)]
34pub struct PPDocLayoutModelBuilder {
35    inner: ScaleAwareDetectorModelBuilder,
36}
37
38impl PPDocLayoutModelBuilder {
39    /// Creates a new PP-DocLayout builder with default settings.
40    pub fn new() -> Self {
41        Self {
42            inner: ScaleAwareDetectorModelBuilder::pp_doclayout(),
43        }
44    }
45
46    /// Sets the preprocessing configuration.
47    pub fn preprocess_config(mut self, config: PPDocLayoutPreprocessConfig) -> Self {
48        self.inner = self.inner.preprocess_config(config);
49        self
50    }
51
52    /// Sets the image shape.
53    pub fn image_shape(mut self, height: u32, width: u32) -> Self {
54        self.inner = self.inner.image_shape(height, width);
55        self
56    }
57
58    /// Builds the PP-DocLayout model.
59    pub fn build(self, inference: OrtInfer) -> Result<PPDocLayoutModel, OCRError> {
60        self.inner
61            .inference_mode(ScaleAwareDetectorInferenceMode::ScaleFactorAndImageShape)
62            .build(inference)
63    }
64}