edgeparse_core/pdf/
image_region.rs

1//! Image-region classification for OCR routing (SRP / Strategy).
2//!
3//! Decides whether a raster should go through table OCR, prose OCR, or be
4//! skipped (charts / photos). Keeps the gate in one place so native and
5//! in-memory recovery share the same policy.
6
7use image::GrayImage;
8
9/// What to do with a page image region.
10#[derive(Debug, Clone, Copy, PartialEq, Eq)]
11pub enum ImageRegionKind {
12    /// Ruled or numeric grid — run table OCR.
13    Table,
14    /// Dense prose screenshot — run text OCR.
15    Prose,
16    /// Chart, photo, or UI chrome — skip OCR.
17    Skip,
18}
19
20/// Strategy for classifying a grayscale page region.
21pub trait ImageRegionClassifier: Send + Sync {
22    /// Return the OCR routing decision for `gray`.
23    fn classify(&self, gray: &GrayImage) -> ImageRegionKind;
24}
25
26/// Optional ML / table-structure model plugged into the `Table` branch.
27///
28/// Implementations (TableFormer, SLANet, …) refine cell geometry after the
29/// heuristic classifier says [`ImageRegionKind::Table`]. The default is a
30/// no-op that leaves lattice recovery to classical line detection.
31pub trait TableStructureModel: Send + Sync {
32    /// Refine a table region. Return `true` when the model produced usable
33    /// structure that downstream OCR should prefer over pure heuristics.
34    fn refine_table(&self, gray: &GrayImage) -> bool {
35        let _ = gray;
36        false
37    }
38}
39
40/// No-op table-structure model (classical bordered-grid path only).
41#[derive(Debug, Default, Clone, Copy)]
42pub struct NullTableStructureModel;
43
44impl TableStructureModel for NullTableStructureModel {}
45
46/// Classifier that runs heuristics, then optionally consults a table model.
47#[derive(Debug, Clone)]
48pub struct ClassifyingRegionRouter<C, M> {
49    /// Region kind classifier.
50    pub classifier: C,
51    /// Table-structure model for the `Table` branch.
52    pub table_model: M,
53}
54
55impl<C: ImageRegionClassifier, M: TableStructureModel> ClassifyingRegionRouter<C, M> {
56    /// Create a router.
57    pub fn new(classifier: C, table_model: M) -> Self {
58        Self {
59            classifier,
60            table_model,
61        }
62    }
63
64    /// Classify then, for tables, invoke the structure model.
65    pub fn route(&self, gray: &GrayImage) -> (ImageRegionKind, bool) {
66        let kind = self.classifier.classify(gray);
67        let model_hit =
68            matches!(kind, ImageRegionKind::Table) && self.table_model.refine_table(gray);
69        (kind, model_hit)
70    }
71}
72
73/// Cheap heuristic classifier (no ML). Wraps the existing raster heuristics in
74/// [`crate::pdf::raster_table_ocr`] so policy lives behind one trait.
75#[derive(Debug, Default, Clone, Copy)]
76pub struct HeuristicImageRegionClassifier;
77
78impl ImageRegionClassifier for HeuristicImageRegionClassifier {
79    fn classify(&self, gray: &GrayImage) -> ImageRegionKind {
80        #[cfg(not(target_arch = "wasm32"))]
81        {
82            use crate::pdf::raster_table_ocr::classify_raster_region;
83            classify_raster_region(gray)
84        }
85        #[cfg(target_arch = "wasm32")]
86        {
87            let _ = gray;
88            ImageRegionKind::Table
89        }
90    }
91}
92
93/// Shared deadline tracker for OCR work across a document.
94#[derive(Debug, Clone)]
95pub struct OcrBudget {
96    deadline: Option<std::time::Instant>,
97}
98
99impl OcrBudget {
100    /// `None` = unlimited. `Some(0)` = already exhausted.
101    pub fn from_millis(ms: Option<u64>) -> Self {
102        let deadline = ms.map(|m| std::time::Instant::now() + std::time::Duration::from_millis(m));
103        Self { deadline }
104    }
105
106    /// No deadline.
107    pub fn unlimited() -> Self {
108        Self { deadline: None }
109    }
110
111    /// True when the wall-clock budget has been spent.
112    pub fn exhausted(&self) -> bool {
113        self.deadline
114            .is_some_and(|d| std::time::Instant::now() >= d)
115    }
116
117    /// True when more OCR work is still allowed.
118    pub fn allow(&self) -> bool {
119        !self.exhausted()
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126
127    struct AlwaysTable;
128    impl TableStructureModel for AlwaysTable {
129        fn refine_table(&self, _: &GrayImage) -> bool {
130            true
131        }
132    }
133
134    #[test]
135    fn router_invokes_table_model_on_table_kind() {
136        let router = ClassifyingRegionRouter::new(HeuristicImageRegionClassifier, AlwaysTable);
137        let gray = GrayImage::new(8, 8);
138        let (kind, hit) = router.route(&gray);
139        // Heuristic on tiny blank image may Skip or Table depending on platform.
140        if kind == ImageRegionKind::Table {
141            assert!(hit);
142        } else {
143            assert!(!hit);
144        }
145    }
146}
edgeparse_core/pdf/image_region.rs

edgeparse_core/pdf/
image_region.rs
