Skip to main content

oar_ocr/oarocr/
structure.rs

1//! High-level builder API for document structure analysis.
2//!
3//! This module provides a fluent builder interface for constructing document structure
4//! analysis pipelines that can detect layout elements, recognize tables, extract formulas,
5//! and optionally integrate OCR for text extraction.
6
7use super::builder_utils::{
8    apply_gpu_memory_budget, build_optional_adapter, default_cpu_region_batch_size,
9    resolve_device_batch_sizes, resolve_model_path, resolve_model_source,
10};
11use oar_ocr_core::core::config::OrtSessionConfig;
12use oar_ocr_core::core::traits::OrtConfigurable;
13use oar_ocr_core::core::traits::adapter::{AdapterBuilder, ModelAdapter};
14use oar_ocr_core::core::{ModelSource, OCRError};
15use oar_ocr_core::domain::adapters::{
16    DocumentOrientationAdapter, DocumentOrientationAdapterBuilder, FormulaRecognitionAdapter,
17    LayoutDetectionAdapter, LayoutDetectionAdapterBuilder, PPFormulaNetAdapterBuilder,
18    SLANetWiredAdapterBuilder, SLANetWirelessAdapterBuilder, SealTextDetectionAdapter,
19    SealTextDetectionAdapterBuilder, TableCellDetectionAdapter, TableCellDetectionAdapterBuilder,
20    TableClassificationAdapter, TableClassificationAdapterBuilder,
21    TableStructureRecognitionAdapter, TextDetectionAdapter, TextDetectionAdapterBuilder,
22    TextLineOrientationAdapter, TextLineOrientationAdapterBuilder, TextRecognitionAdapter,
23    TextRecognitionAdapterBuilder, UVDocRectifierAdapter, UVDocRectifierAdapterBuilder,
24    UniMERNetAdapterBuilder,
25};
26use oar_ocr_core::domain::structure::{StructureResult, TableResult, TableType};
27use oar_ocr_core::domain::tasks::{
28    FormulaRecognitionConfig, LayoutDetectionConfig, TableCellDetectionConfig,
29    TableClassificationConfig, TableStructureRecognitionConfig, TextDetectionConfig,
30    TextRecognitionConfig,
31};
32use oar_ocr_core::predictors::FormulaModelKind;
33use std::path::PathBuf;
34use std::sync::Arc;
35use std::time::Instant;
36
37/// IoU threshold for removing overlapping layout elements (0.5 = 50% overlap).
38const LAYOUT_OVERLAP_IOU_THRESHOLD: f32 = 0.5;
39
40/// IoU threshold for determining if an OCR box overlaps with table cells.
41const CELL_OVERLAP_IOU_THRESHOLD: f32 = 0.5;
42
43/// IoA threshold for assigning layout elements to region blocks during reading order.
44/// A low threshold (0.1 = 10%) allows elements near region boundaries to be included.
45const REGION_MEMBERSHIP_IOA_THRESHOLD: f32 = 0.1;
46
47/// IoA threshold for splitting text boxes that intersect with container elements.
48/// A moderate threshold (0.3 = 30%) balances precision with avoiding over-splitting.
49const TEXT_BOX_SPLIT_IOA_THRESHOLD: f32 = 0.3;
50
51/// Internal structure holding the structure analysis pipeline adapters.
52#[derive(Debug)]
53struct StructurePipeline {
54    // Document preprocessing (optional)
55    document_orientation_adapter: Option<DocumentOrientationAdapter>,
56    rectification_adapter: Option<UVDocRectifierAdapter>,
57
58    // Layout analysis (required)
59    layout_detection_adapter: LayoutDetectionAdapter,
60
61    // Region detection for hierarchical ordering (optional, PP-DocBlockLayout)
62    region_detection_adapter: Option<LayoutDetectionAdapter>,
63
64    // Table analysis (optional)
65    table_classification_adapter: Option<TableClassificationAdapter>,
66    table_orientation_adapter: Option<DocumentOrientationAdapter>, // Reuses doc orientation model
67    table_cell_detection_adapter: Option<TableCellDetectionAdapter>,
68    table_structure_recognition_adapter: Option<TableStructureRecognitionAdapter>,
69    // PP-StructureV3 auto-switch: separate adapters for wired/wireless tables
70    wired_table_structure_adapter: Option<TableStructureRecognitionAdapter>,
71    wireless_table_structure_adapter: Option<TableStructureRecognitionAdapter>,
72    wired_table_cell_adapter: Option<TableCellDetectionAdapter>,
73    wireless_table_cell_adapter: Option<TableCellDetectionAdapter>,
74    // E2E mode: when true, skip cell detection and use only structure model output
75    use_e2e_wired_table_rec: bool,
76    use_e2e_wireless_table_rec: bool,
77    // PaddleX compatibility: build table HTML from cell detection boxes
78    use_wired_table_cells_trans_to_html: bool,
79    use_wireless_table_cells_trans_to_html: bool,
80
81    formula_recognition_adapter: Option<FormulaRecognitionAdapter>,
82
83    seal_text_detection_adapter: Option<SealTextDetectionAdapter>,
84
85    // OCR integration (optional)
86    text_detection_adapter: Option<TextDetectionAdapter>,
87    text_line_orientation_adapter: Option<TextLineOrientationAdapter>,
88    text_recognition_adapter: Option<TextRecognitionAdapter>,
89
90    // Batch sizes for image-level and region-level processing.
91    image_batch_size: Option<usize>,
92    region_batch_size: Option<usize>,
93}
94
95/// High-level builder for document structure analysis pipelines.
96///
97/// This builder provides a fluent API for constructing document structure analysis
98/// pipelines with various components:
99/// - Document preprocessing (optional): orientation detection and rectification
100/// - Layout detection (required)
101/// - Table classification (optional)
102/// - Table cell detection (optional)
103/// - Table structure recognition (optional)
104/// - Formula recognition (optional)
105/// - Seal text detection (optional)
106/// - OCR integration (optional)
107///
108/// # Example
109///
110/// ```no_run
111/// use oar_ocr::oarocr::structure::OARStructureBuilder;
112///
113/// # fn main() -> Result<(), Box<dyn std::error::Error>> {
114/// let structure = OARStructureBuilder::new("layout.onnx")
115///     .with_table_classification("table_cls.onnx")
116///     .with_table_cell_detection("table_cell.onnx", "wired")
117///     .with_table_structure_recognition("table_struct.onnx", "wired")
118///     .with_formula_recognition(
119///         "formula.onnx",
120///         "tokenizer.json",
121///         "pp_formulanet"
122///     )
123///     .build()?;
124/// # let _ = structure;
125/// # Ok(())
126/// # }
127/// ```
128#[derive(Debug, Clone)]
129pub struct OARStructureBuilder {
130    // Required models
131    layout_detection_model: ModelSource,
132    layout_model_name: Option<String>,
133
134    // Optional document preprocessing
135    document_orientation_model: Option<ModelSource>,
136    document_rectification_model: Option<ModelSource>,
137
138    // Optional region detection for hierarchical ordering (PP-DocBlockLayout)
139    region_detection_model: Option<ModelSource>,
140
141    // Optional table analysis models
142    table_classification_model: Option<ModelSource>,
143    table_orientation_model: Option<ModelSource>, // Reuses doc orientation model for rotated tables
144    table_cell_detection_model: Option<ModelSource>,
145    table_cell_detection_type: Option<Result<TableType, String>>,
146    table_structure_recognition_model: Option<ModelSource>,
147    table_structure_recognition_type: Option<Result<TableType, String>>,
148    table_structure_dict_path: Option<PathBuf>,
149
150    wired_table_structure_model: Option<ModelSource>,
151    wireless_table_structure_model: Option<ModelSource>,
152    wired_table_cell_model: Option<ModelSource>,
153    wireless_table_cell_model: Option<ModelSource>,
154    // E2E mode: when true, skip cell detection and use only structure model output
155    // Defaults: wired=false, wireless=true
156    use_e2e_wired_table_rec: bool,
157    use_e2e_wireless_table_rec: bool,
158    // PaddleX compatibility: build table HTML from cell detection boxes
159    // Defaults: wired=false, wireless=false
160    use_wired_table_cells_trans_to_html: bool,
161    use_wireless_table_cells_trans_to_html: bool,
162
163    // Optional formula recognition
164    formula_recognition_model: Option<ModelSource>,
165    formula_recognition_type: Option<Result<FormulaModelKind, String>>,
166    formula_tokenizer_path: Option<PathBuf>,
167    formula_ort_session_config: Option<OrtSessionConfig>,
168
169    // Optional seal text detection
170    seal_text_detection_model: Option<ModelSource>,
171
172    // Optional OCR integration
173    text_detection_model: Option<ModelSource>,
174    text_line_orientation_model: Option<ModelSource>,
175    text_recognition_model: Option<ModelSource>,
176    character_dict_path: Option<PathBuf>,
177
178    // Model name presets for loading correct pre/post processors
179    region_model_name: Option<String>,
180    wired_table_structure_model_name: Option<String>,
181    wireless_table_structure_model_name: Option<String>,
182    wired_table_cell_model_name: Option<String>,
183    wireless_table_cell_model_name: Option<String>,
184    text_detection_model_name: Option<String>,
185    text_recognition_model_name: Option<String>,
186
187    // Configuration
188    ort_session_config: Option<OrtSessionConfig>,
189    layout_detection_config: Option<LayoutDetectionConfig>,
190    table_classification_config: Option<TableClassificationConfig>,
191    table_cell_detection_config: Option<TableCellDetectionConfig>,
192    table_structure_recognition_config: Option<TableStructureRecognitionConfig>,
193    formula_recognition_config: Option<FormulaRecognitionConfig>,
194    text_detection_config: Option<TextDetectionConfig>,
195    text_recognition_config: Option<TextRecognitionConfig>,
196
197    // Batch sizes
198    image_batch_size: Option<usize>,
199    region_batch_size: Option<usize>,
200    gpu_memory_budget: Option<usize>,
201}
202
203/// Detection thresholds and merge rules for a layout preset without an explicit
204/// config; PP-StructureV3 defaults cover the earlier 20/23-class models.
205fn default_layout_detection_config(model_name: &str) -> LayoutDetectionConfig {
206    match model_name {
207        "pp-doclayoutv2" => LayoutDetectionConfig::with_pp_doclayoutv2_defaults(),
208        "pp-doclayoutv3" => LayoutDetectionConfig::with_pp_doclayoutv3_defaults(),
209        _ => LayoutDetectionConfig::with_pp_structurev3_defaults(),
210    }
211}
212
213/// Resolves a `layout_model_name` preset, falling back to PP-DocLayout_plus-L.
214fn layout_model_config_for(name: &str) -> oar_ocr_core::domain::adapters::LayoutModelConfig {
215    use oar_ocr_core::domain::adapters::LayoutModelConfig;
216    // Match presets case- and separator-insensitively so the documented
217    // forms (e.g. `PicoDet-L_layout_17cls`, `RT-DETR-H_layout_17cls`,
218    // `PP-DocLayout_plus-L`) resolve correctly. Mirrors the normalization
219    // used by `region_model_name` below.
220    match name.to_lowercase().replace('-', "_").as_str() {
221        "picodet_layout_1x" => LayoutModelConfig::picodet_layout_1x(),
222        "picodet_layout_1x_table" => LayoutModelConfig::picodet_layout_1x_table(),
223        "picodet_s_layout_3cls" => LayoutModelConfig::picodet_s_layout_3cls(),
224        "picodet_l_layout_3cls" => LayoutModelConfig::picodet_l_layout_3cls(),
225        "picodet_s_layout_17cls" => LayoutModelConfig::picodet_s_layout_17cls(),
226        "picodet_l_layout_17cls" => LayoutModelConfig::picodet_l_layout_17cls(),
227        "rt_detr_h_layout_3cls" => LayoutModelConfig::rtdetr_h_layout_3cls(),
228        "rt_detr_h_layout_17cls" => LayoutModelConfig::rtdetr_h_layout_17cls(),
229        "pp_docblocklayout" => LayoutModelConfig::pp_docblocklayout(),
230        "pp_doclayout_s" => LayoutModelConfig::pp_doclayout_s(),
231        "pp_doclayout_m" => LayoutModelConfig::pp_doclayout_m(),
232        "pp_doclayout_l" => LayoutModelConfig::pp_doclayout_l(),
233        "pp_doclayout_plus_l" => LayoutModelConfig::pp_doclayout_plus_l(),
234        "pp_doclayoutv2" | "pp_doclayout_v2" => LayoutModelConfig::pp_doclayoutv2(),
235        "pp_doclayoutv3" | "pp_doclayout_v3" => LayoutModelConfig::pp_doclayoutv3(),
236        _ => {
237            tracing::warn!(
238                requested = %name,
239                "Unknown --layout-model-name preset; falling back to PP-DocLayout_plus-L. \
240                 This may apply the wrong class labels/preprocessing for your model."
241            );
242            LayoutModelConfig::pp_doclayout_plus_l()
243        }
244    }
245}
246
247/// Turns CUDA arena shrinkage on for a CUDA session config unless the caller
248/// already chose a value; CPU-only configs are left untouched.
249fn default_arena_shrinkage(config: &mut OrtSessionConfig) {
250    use oar_ocr_core::core::config::OrtExecutionProvider;
251    let uses_cuda = config
252        .execution_providers
253        .iter()
254        .flatten()
255        .any(|ep| matches!(ep, OrtExecutionProvider::CUDA { .. }));
256    if uses_cuda && config.arena_shrinkage.is_none() {
257        config.arena_shrinkage = Some(true);
258    }
259}
260
261impl OARStructureBuilder {
262    const MAX_BATCH_SIZE: usize = 4096;
263
264    /// Creates a new structure builder with the required layout detection model.
265    ///
266    /// # Arguments
267    ///
268    /// * `layout_detection_model` - Path to the layout detection model file
269    pub fn new(layout_detection_model: impl Into<ModelSource>) -> Self {
270        Self {
271            layout_detection_model: layout_detection_model.into(),
272            layout_model_name: None,
273            document_orientation_model: None,
274            document_rectification_model: None,
275            region_detection_model: None,
276            table_classification_model: None,
277            table_orientation_model: None,
278            table_cell_detection_model: None,
279            table_cell_detection_type: None,
280            table_structure_recognition_model: None,
281            table_structure_recognition_type: None,
282            table_structure_dict_path: None,
283            wired_table_structure_model: None,
284            wireless_table_structure_model: None,
285            wired_table_cell_model: None,
286            wireless_table_cell_model: None,
287            // Defaults: wired=false (use cell detection), wireless=true (E2E mode)
288            use_e2e_wired_table_rec: false,
289            use_e2e_wireless_table_rec: true,
290            use_wired_table_cells_trans_to_html: false,
291            use_wireless_table_cells_trans_to_html: false,
292            formula_recognition_model: None,
293            formula_recognition_type: None,
294            formula_tokenizer_path: None,
295            formula_ort_session_config: None,
296            seal_text_detection_model: None,
297            text_detection_model: None,
298            text_line_orientation_model: None,
299            text_recognition_model: None,
300            character_dict_path: None,
301            region_model_name: None,
302            wired_table_structure_model_name: None,
303            wireless_table_structure_model_name: None,
304            wired_table_cell_model_name: None,
305            wireless_table_cell_model_name: None,
306            text_detection_model_name: None,
307            text_recognition_model_name: None,
308            ort_session_config: None,
309            layout_detection_config: None,
310            table_classification_config: None,
311            table_cell_detection_config: None,
312            table_structure_recognition_config: None,
313            formula_recognition_config: None,
314            text_detection_config: None,
315            text_recognition_config: None,
316            image_batch_size: None,
317            region_batch_size: None,
318            gpu_memory_budget: None,
319        }
320    }
321
322    /// Creates a PP-StructureV3-style document analysis pipeline.
323    ///
324    /// This is the stack the benchmark tool's `structure-v3` case runs:
325    /// PP-DocLayoutV3 layout detection, PP-OCRv6 Tiny text detection and
326    /// recognition (the structure pipeline OCRs many small region crops, where
327    /// the tiny recognizer's throughput wins and the larger layout and table
328    /// models already dominate memory), PP-LCNet table classification,
329    /// SLANeXt wired and SLANet+ wireless table structure recognition, the
330    /// RT-DETR-L wired cell detector, and the Chinese table structure
331    /// dictionary. With the `auto-download` feature the names resolve through
332    /// the model registry; without it they resolve as local paths like any
333    /// other model source. The returned builder is ordinary, so every setter
334    /// can still override anything.
335    ///
336    /// ```no_run
337    /// use oar_ocr::oarocr::OARStructureBuilder;
338    ///
339    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
340    /// let structure = OARStructureBuilder::pp_structurev3().build()?;
341    /// # let _ = structure;
342    /// # Ok(())
343    /// # }
344    /// ```
345    pub fn pp_structurev3() -> Self {
346        use super::ocr::PpOcrV6Size;
347        let size = PpOcrV6Size::Tiny;
348        Self::new("pp-doclayoutv3.onnx")
349            .layout_model_name("PP-DocLayoutV3")
350            .with_ocr(
351                size.detection_model(),
352                size.recognition_model(),
353                size.character_dict(),
354            )
355            .with_table_classification("pp-lcnet_x1_0_table_cls.onnx")
356            .with_wired_table_structure("slanext_wired.onnx")
357            .with_wireless_table_structure("slanet_plus.onnx")
358            .with_wired_table_cell_detection("rt-detr-l_wired_table_cell_det.onnx")
359            .table_structure_dict_path("table_structure_dict_ch.txt")
360    }
361
362    /// Sets the ONNX Runtime session configuration.
363    ///
364    /// This configuration will be applied to all models in the pipeline.
365    pub fn ort_session(mut self, config: OrtSessionConfig) -> Self {
366        self.ort_session_config = Some(config);
367        self
368    }
369
370    /// Sets a GPU memory budget in bytes without selecting a device.
371    ///
372    /// With `g = floor(bytes / GiB)`, image batches use `clamp(2*g - 6, 1, 8)`
373    /// and text-region batches use `clamp(12*g - 28, 4, 64)`, saturating negative
374    /// values at zero. Thus 4 GiB gives 2/20, and 8 GiB or more gives 8/64.
375    /// Explicit batch sizes take precedence. CUDA sessions cap each arena at half
376    /// the budget and default to idle memory recovery and SameAsRequested growth.
377    /// Existing smaller limits and explicit arena settings are retained. This is a tuning hint, not a device-wide cap:
378    /// other sessions, weights, and driver allocations also consume memory.
379    pub fn gpu_memory_budget(mut self, bytes: usize) -> Self {
380        self.gpu_memory_budget = Some(bytes);
381        self
382    }
383
384    /// Sets the layout detection model configuration.
385    pub fn layout_detection_config(mut self, config: LayoutDetectionConfig) -> Self {
386        self.layout_detection_config = Some(config);
387        self
388    }
389
390    /// Overrides the built-in layout model preset used to configure preprocessing/postprocessing.
391    ///
392    /// This is useful when the ONNX file name alone is not enough to infer the correct
393    /// model family. Preset names are matched case- and separator-insensitively
394    /// (`-` and `_` are interchangeable), so `PP-DocLayout_plus-L` and
395    /// `pp_doclayout_plus_l` are equivalent. Supported presets:
396    /// - `PP-DocLayout_plus-L` (default)
397    /// - `PP-DocLayout-S`, `PP-DocLayout-M`, `PP-DocLayout-L`
398    /// - `PP-DocLayoutV2`, `PP-DocLayoutV3`
399    /// - `PP-DocBlockLayout`
400    /// - `PicoDet_layout_1x`, `PicoDet_layout_1x_table`
401    /// - `PicoDet-S_layout_3cls`, `PicoDet-L_layout_3cls`
402    /// - `PicoDet-S_layout_17cls`, `PicoDet-L_layout_17cls`
403    /// - `RT-DETR-H_layout_3cls`, `RT-DETR-H_layout_17cls`
404    ///
405    /// An unrecognized name logs a warning and falls back to the default preset.
406    pub fn layout_model_name(mut self, name: impl Into<String>) -> Self {
407        self.layout_model_name = Some(name.into());
408        self
409    }
410
411    /// Sets the region detection model name preset.
412    ///
413    /// This is used to load the correct preprocessing/postprocessing for the region
414    /// detection model. Supported presets: `PP-DocBlockLayout`.
415    pub fn region_model_name(mut self, name: impl Into<String>) -> Self {
416        self.region_model_name = Some(name.into());
417        self
418    }
419
420    /// Sets the reported model name for the wired table structure model
421    /// (e.g. `SLANeXt_wired`).
422    ///
423    /// This is identification metadata: it labels the model in logs and error
424    /// messages. The wired/wireless slot and the table-structure config govern
425    /// the actual input shape and decoding.
426    pub fn wired_table_structure_model_name(mut self, name: impl Into<String>) -> Self {
427        self.wired_table_structure_model_name = Some(name.into());
428        self
429    }
430
431    /// Sets the reported model name for the wireless table structure model
432    /// (e.g. `SLANet_plus`). Identification metadata only; see
433    /// [`Self::wired_table_structure_model_name`].
434    pub fn wireless_table_structure_model_name(mut self, name: impl Into<String>) -> Self {
435        self.wireless_table_structure_model_name = Some(name.into());
436        self
437    }
438
439    /// Sets the reported model name for the wired table cell detector
440    /// (e.g. `RT-DETR-L_wired_table_cell_det`). Identification metadata only.
441    pub fn wired_table_cell_model_name(mut self, name: impl Into<String>) -> Self {
442        self.wired_table_cell_model_name = Some(name.into());
443        self
444    }
445
446    /// Sets the reported model name for the wireless table cell detector
447    /// (e.g. `RT-DETR-L_wireless_table_cell_det`). Identification metadata only.
448    pub fn wireless_table_cell_model_name(mut self, name: impl Into<String>) -> Self {
449        self.wireless_table_cell_model_name = Some(name.into());
450        self
451    }
452
453    /// Sets the reported model name for the text detector
454    /// (e.g. `PP-OCRv5_server_det`). Identification metadata only: detection
455    /// behavior is driven by the config and the ONNX model.
456    pub fn text_detection_model_name(mut self, name: impl Into<String>) -> Self {
457        self.text_detection_model_name = Some(name.into());
458        self
459    }
460
461    /// Sets the reported model name for the text recognizer
462    /// (e.g. `PP-OCRv5_server_rec`). Identification metadata only: recognition
463    /// behavior is driven by the config, dictionary, and the ONNX model.
464    pub fn text_recognition_model_name(mut self, name: impl Into<String>) -> Self {
465        self.text_recognition_model_name = Some(name.into());
466        self
467    }
468
469    /// Sets the batch size for image-level processing.
470    ///
471    /// Controls how many pages are processed together by image-level stages such as
472    /// layout detection, region detection, and OCR text detection. When unset, CPU
473    /// execution uses `1`; explicitly configured accelerators retain adapter defaults.
474    pub fn image_batch_size(mut self, size: usize) -> Self {
475        self.image_batch_size = Some(size);
476        self
477    }
478
479    /// Sets the batch size for region-level processing (text recognition).
480    ///
481    /// Controls how many text regions are processed together during OCR recognition.
482    /// When unset, CPU execution uses `16` for PP-OCRv6 Tiny and `4` for other
483    /// recognizers; explicitly configured accelerators retain their larger
484    /// throughput-oriented default.
485    pub fn region_batch_size(mut self, size: usize) -> Self {
486        self.region_batch_size = Some(size);
487        self
488    }
489
490    /// Adds document orientation detection to the pipeline.
491    ///
492    /// This component detects and corrects document rotation (0°, 90°, 180°, 270°).
493    /// Should be run before other processing for best results.
494    pub fn with_document_orientation(mut self, model_source: impl Into<ModelSource>) -> Self {
495        self.document_orientation_model = Some(model_source.into());
496        self
497    }
498
499    /// Adds document rectification to the pipeline.
500    ///
501    /// This component corrects document distortion and perspective issues.
502    /// Should be run after orientation detection if both are enabled.
503    pub fn with_document_rectification(mut self, model_source: impl Into<ModelSource>) -> Self {
504        self.document_rectification_model = Some(model_source.into());
505        self
506    }
507
508    /// Adds region detection to the pipeline (PP-DocBlockLayout).
509    ///
510    /// This component detects document regions (columns, blocks) for hierarchical
511    /// layout ordering. Region blocks provide grouping information for improved
512    /// reading order within multi-column or complex layouts.
513    ///
514    /// # PP-StructureV3 Integration
515    ///
516    /// When enabled, the pipeline uses region detection results to:
517    /// 1. Group layout elements by their parent regions
518    /// 2. Apply XY-cut ordering within each region
519    /// 3. Order regions based on their relative positions
520    pub fn with_region_detection(mut self, model_source: impl Into<ModelSource>) -> Self {
521        self.region_detection_model = Some(model_source.into());
522        self
523    }
524
525    /// Adds seal text detection to the pipeline.
526    ///
527    /// This component detects circular/curved seal and stamp text regions.
528    /// Seal regions will be included in the layout elements.
529    pub fn with_seal_text_detection(mut self, model_source: impl Into<ModelSource>) -> Self {
530        self.seal_text_detection_model = Some(model_source.into());
531        self
532    }
533
534    /// Adds table classification to the pipeline.
535    ///
536    /// This component classifies tables as wired or wireless.
537    pub fn with_table_classification(mut self, model_source: impl Into<ModelSource>) -> Self {
538        self.table_classification_model = Some(model_source.into());
539        self
540    }
541
542    /// Sets the table classification configuration.
543    pub fn table_classification_config(mut self, config: TableClassificationConfig) -> Self {
544        self.table_classification_config = Some(config);
545        self
546    }
547
548    /// Adds table orientation detection to the pipeline.
549    ///
550    /// This component detects if tables are rotated (0°, 90°, 180°, 270°) and corrects them
551    /// before structure recognition. Uses the same model as document orientation detection
552    /// (PP-LCNet_x1_0_doc_ori).
553    ///
554    /// # Arguments
555    ///
556    /// * `model_path` - Path to the orientation classification model (same as document orientation)
557    pub fn with_table_orientation(mut self, model_source: impl Into<ModelSource>) -> Self {
558        self.table_orientation_model = Some(model_source.into());
559        self
560    }
561
562    /// Sets whether to use end-to-end mode for wired table recognition.
563    ///
564    /// When enabled, cell detection model is skipped and only the table structure
565    /// recognition model's cell output is used. When disabled, RT-DETR cell detection
566    /// provides more precise cell bounding boxes.
567    ///
568    /// Default: `false` (use cell detection for wired tables)
569    pub fn use_e2e_wired_table_rec(mut self, enabled: bool) -> Self {
570        self.use_e2e_wired_table_rec = enabled;
571        self
572    }
573
574    /// Sets whether to use end-to-end mode for wireless table recognition.
575    ///
576    /// When enabled, cell detection model is skipped and only the table structure
577    /// recognition model's cell output is used. When disabled, RT-DETR cell detection
578    /// provides more precise cell bounding boxes.
579    ///
580    /// Default: `true` (E2E mode for wireless tables)
581    pub fn use_e2e_wireless_table_rec(mut self, enabled: bool) -> Self {
582        self.use_e2e_wireless_table_rec = enabled;
583        self
584    }
585
586    /// Enables PaddleX-compatible conversion of wired cell detections to HTML structure tokens.
587    ///
588    /// When enabled, wired tables can derive structure tokens directly from detected cell boxes,
589    /// which is useful for parity testing with PaddleX `use_wired_table_cells_trans_to_html`.
590    pub fn use_wired_table_cells_trans_to_html(mut self, enabled: bool) -> Self {
591        self.use_wired_table_cells_trans_to_html = enabled;
592        self
593    }
594
595    /// Enables PaddleX-compatible conversion of wireless cell detections to HTML structure tokens.
596    ///
597    /// When enabled, wireless tables can derive structure tokens directly from detected cell boxes,
598    /// which is useful for parity testing with PaddleX `use_wireless_table_cells_trans_to_html`.
599    pub fn use_wireless_table_cells_trans_to_html(mut self, enabled: bool) -> Self {
600        self.use_wireless_table_cells_trans_to_html = enabled;
601        self
602    }
603
604    /// Adds table cell detection to the pipeline.
605    ///
606    /// # Arguments
607    ///
608    /// * `model_path` - Path to the table cell detection model
609    /// * `cell_type` - `"wired"`, `"wireless"`, or the corresponding [`TableType`]
610    ///
611    /// This replaces any model previously set for the same type with
612    /// [`with_wired_table_cell_detection`](Self::with_wired_table_cell_detection)
613    /// or [`with_wireless_table_cell_detection`](Self::with_wireless_table_cell_detection),
614    /// for example one chosen by [`pp_structurev3`](Self::pp_structurev3).
615    pub fn with_table_cell_detection(
616        mut self,
617        model_source: impl Into<ModelSource>,
618        cell_type: impl AsRef<str>,
619    ) -> Self {
620        let cell_type = cell_type.as_ref().parse();
621        match cell_type {
622            Ok(TableType::Wired) => self.wired_table_cell_model = None,
623            Ok(TableType::Wireless) => self.wireless_table_cell_model = None,
624            _ => {}
625        }
626        self.table_cell_detection_model = Some(model_source.into());
627        self.table_cell_detection_type = Some(cell_type);
628        self
629    }
630
631    /// Sets the table cell detection configuration.
632    pub fn table_cell_detection_config(mut self, config: TableCellDetectionConfig) -> Self {
633        self.table_cell_detection_config = Some(config);
634        self
635    }
636
637    /// Adds table structure recognition to the pipeline.
638    ///
639    /// # Arguments
640    ///
641    /// * `model_path` - Path to the table structure recognition model
642    /// * `table_type` - `"wired"`, `"wireless"`, or the corresponding [`TableType`]
643    ///
644    /// This component recognizes the structure of tables and outputs HTML. It
645    /// replaces any model previously set for the same type with
646    /// [`with_wired_table_structure`](Self::with_wired_table_structure) or
647    /// [`with_wireless_table_structure`](Self::with_wireless_table_structure),
648    /// for example one chosen by [`pp_structurev3`](Self::pp_structurev3).
649    pub fn with_table_structure_recognition(
650        mut self,
651        model_source: impl Into<ModelSource>,
652        table_type: impl AsRef<str>,
653    ) -> Self {
654        let table_type = table_type.as_ref().parse();
655        match table_type {
656            Ok(TableType::Wired) => self.wired_table_structure_model = None,
657            Ok(TableType::Wireless) => self.wireless_table_structure_model = None,
658            _ => {}
659        }
660        self.table_structure_recognition_model = Some(model_source.into());
661        self.table_structure_recognition_type = Some(table_type);
662        self
663    }
664
665    /// Sets the dictionary path for table structure recognition.
666    ///
667    /// The dictionary file should match the model type:
668    /// - `table_structure_dict_ch.txt` for Chinese
669    /// - `table_structure_dict.txt` for English
670    /// - `table_master_structure_dict.txt` for extended tags
671    pub fn table_structure_dict_path(mut self, path: impl Into<PathBuf>) -> Self {
672        self.table_structure_dict_path = Some(path.into());
673        self
674    }
675
676    /// Sets the table structure recognition configuration.
677    pub fn table_structure_recognition_config(
678        mut self,
679        config: TableStructureRecognitionConfig,
680    ) -> Self {
681        self.table_structure_recognition_config = Some(config);
682        self
683    }
684
685    /// Adds wired table structure recognition model.
686    ///
687    /// When both wired and wireless models are configured along with table classification,
688    /// the system automatically selects the appropriate model based on classification results.
689    pub fn with_wired_table_structure(mut self, model_source: impl Into<ModelSource>) -> Self {
690        self.wired_table_structure_model = Some(model_source.into());
691        self
692    }
693
694    /// Adds wireless table structure recognition model.
695    ///
696    /// When both wired and wireless models are configured along with table classification,
697    /// the system automatically selects the appropriate model based on classification results.
698    pub fn with_wireless_table_structure(mut self, model_source: impl Into<ModelSource>) -> Self {
699        self.wireless_table_structure_model = Some(model_source.into());
700        self
701    }
702
703    /// Adds wired table cell detection model.
704    ///
705    /// When both wired and wireless models are configured along with table classification,
706    /// the system automatically selects the appropriate model based on classification results.
707    pub fn with_wired_table_cell_detection(mut self, model_source: impl Into<ModelSource>) -> Self {
708        self.wired_table_cell_model = Some(model_source.into());
709        self
710    }
711
712    /// Adds wireless table cell detection model.
713    ///
714    /// When both wired and wireless models are configured along with table classification,
715    /// the system automatically selects the appropriate model based on classification results.
716    pub fn with_wireless_table_cell_detection(
717        mut self,
718        model_source: impl Into<ModelSource>,
719    ) -> Self {
720        self.wireless_table_cell_model = Some(model_source.into());
721        self
722    }
723
724    /// Adds formula recognition to the pipeline.
725    ///
726    /// # Arguments
727    ///
728    /// * `model_path` - Path to the formula recognition model
729    /// * `tokenizer_path` - Path to the tokenizer JSON file
730    /// * `model_type` - `"pp_formulanet"`, `"unimernet"`, or the corresponding
731    ///   [`FormulaModelKind`]
732    ///
733    /// This component recognizes mathematical formulas and outputs LaTeX.
734    pub fn with_formula_recognition(
735        mut self,
736        model_source: impl Into<ModelSource>,
737        tokenizer_path: impl Into<PathBuf>,
738        model_type: impl AsRef<str>,
739    ) -> Self {
740        self.formula_recognition_model = Some(model_source.into());
741        self.formula_tokenizer_path = Some(tokenizer_path.into());
742        self.formula_recognition_type = Some(model_type.as_ref().parse());
743        self
744    }
745
746    /// Sets the formula recognition configuration.
747    pub fn formula_recognition_config(mut self, config: FormulaRecognitionConfig) -> Self {
748        self.formula_recognition_config = Some(config);
749        self
750    }
751
752    /// Sets an ONNX Runtime session configuration only for formula recognition.
753    pub fn formula_ort_session(mut self, config: OrtSessionConfig) -> Self {
754        self.formula_ort_session_config = Some(config);
755        self
756    }
757
758    /// Integrates OCR into the pipeline for text extraction.
759    ///
760    /// # Arguments
761    ///
762    /// * `text_detection_model` - Text detection model: a path or raw model bytes
763    /// * `text_recognition_model` - Text recognition model: a path or raw model bytes
764    /// * `character_dict_path` - Path to the character dictionary file
765    pub fn with_ocr(
766        mut self,
767        text_detection_model: impl Into<ModelSource>,
768        text_recognition_model: impl Into<ModelSource>,
769        character_dict_path: impl Into<PathBuf>,
770    ) -> Self {
771        self.text_detection_model = Some(text_detection_model.into());
772        self.text_recognition_model = Some(text_recognition_model.into());
773        self.character_dict_path = Some(character_dict_path.into());
774        self
775    }
776
777    /// Adds text line orientation detection to the OCR pipeline.
778    ///
779    /// This component detects whether text lines are upright (0°) or inverted (180°),
780    /// which helps improve OCR accuracy for documents with mixed text orientations.
781    ///
782    /// # PP-StructureV3 Integration
783    ///
784    /// When enabled, detected text lines are classified before recognition:
785    /// - Lines classified as 180° rotated are flipped before OCR
786    /// - This improves accuracy for documents scanned upside-down or with mixed orientations
787    pub fn with_text_line_orientation(mut self, model_source: impl Into<ModelSource>) -> Self {
788        self.text_line_orientation_model = Some(model_source.into());
789        self
790    }
791
792    /// Sets the text detection configuration.
793    pub fn text_detection_config(mut self, config: TextDetectionConfig) -> Self {
794        self.text_detection_config = Some(config);
795        self
796    }
797
798    /// Sets the text recognition configuration.
799    pub fn text_recognition_config(mut self, config: TextRecognitionConfig) -> Self {
800        self.text_recognition_config = Some(config);
801        self
802    }
803
804    /// Builds the structure analyzer runtime.
805    ///
806    /// This method instantiates all adapters and returns a ready-to-use structure analyzer.
807    pub fn build(mut self) -> Result<OARStructure, OCRError> {
808        for config in [
809            self.ort_session_config.as_mut(),
810            self.formula_ort_session_config.as_mut(),
811        ] {
812            apply_gpu_memory_budget(config, self.gpu_memory_budget)?;
813        }
814        for (component, selection) in [
815            ("table_cell_detection", &self.table_cell_detection_type),
816            (
817                "table_structure_recognition",
818                &self.table_structure_recognition_type,
819            ),
820        ] {
821            match selection {
822                Some(Err(message)) => {
823                    return Err(OCRError::config_error_detailed(component, message.clone()));
824                }
825                Some(Ok(TableType::Unknown)) => {
826                    return Err(OCRError::config_error_detailed(
827                        component,
828                        "'unknown' is not valid when selecting a table model".to_string(),
829                    ));
830                }
831                _ => {}
832            }
833        }
834        if let Some(Err(message)) = self.formula_recognition_type.as_ref() {
835            return Err(OCRError::config_error_detailed(
836                "formula_recognition",
837                message.clone(),
838            ));
839        }
840
841        if let Some(size) = self.image_batch_size {
842            Self::validate_batch_size("image_batch_size", size)?;
843        }
844        if let Some(size) = self.region_batch_size {
845            Self::validate_batch_size("region_batch_size", size)?;
846        }
847
848        // The structure pipeline keeps a dozen CUDA sessions resident, each with
849        // its own memory arena sized to the largest crop it has seen; together
850        // they reach ~19 GB on a multi-page PDF and OOM 16 GB cards. Returning
851        // idle arena memory after each run keeps the peak near 9 GB with
852        // byte-identical output, so it is the default here unless the caller
853        // chose explicitly.
854        for config in [
855            self.ort_session_config.as_mut(),
856            self.formula_ort_session_config.as_mut(),
857        ]
858        .into_iter()
859        .flatten()
860        {
861            default_arena_shrinkage(config);
862        }
863
864        // PP-FormulaNet's CUDA autoregressive Loop races on EP arena buffers when
865        // its `session.run()`s interleave with other models' (onnxruntime#4829),
866        // garbling later formulas. The fix is `CUDA_LAUNCH_BLOCKING=1`, but it
867        // only takes effect if set before the first CUDA session is created — and
868        // here the formula adapter is built after ~10 other CUDA models. So set it
869        // up front whenever a formula model will run on CUDA.
870        if self.formula_recognition_model.is_some() {
871            use oar_ocr_core::core::config::OrtExecutionProvider;
872            let uses_cuda = self
873                .formula_ort_session_config
874                .as_ref()
875                .or(self.ort_session_config.as_ref())
876                .and_then(|cfg| cfg.execution_providers.as_ref())
877                .is_some_and(|eps| {
878                    eps.iter().any(|ep| {
879                        matches!(
880                            ep,
881                            OrtExecutionProvider::CUDA { .. }
882                                | OrtExecutionProvider::TensorRT { .. }
883                        )
884                    })
885                });
886            if uses_cuda {
887                oar_ocr_core::core::inference::ensure_cuda_launch_blocking();
888            }
889        }
890
891        self.ort_session_config = self.ort_session_config.map(OrtSessionConfig::resolve_auto);
892        self.formula_ort_session_config = self
893            .formula_ort_session_config
894            .map(OrtSessionConfig::resolve_auto);
895
896        // Resolve every model/dict/tokenizer path through the auto-download
897        // cache when the `auto-download` feature is enabled. With the feature
898        // off these calls are infallible no-ops.
899        self.layout_detection_model = resolve_model_source(&self.layout_detection_model)?;
900        fn resolve_opt_path(p: &mut Option<PathBuf>) -> Result<(), OCRError> {
901            if let Some(path) = p {
902                *path = resolve_model_path(path)?;
903            }
904            Ok(())
905        }
906        fn resolve_opt_source(s: &mut Option<ModelSource>) -> Result<(), OCRError> {
907            if let Some(source) = s {
908                *source = resolve_model_source(source)?;
909            }
910            Ok(())
911        }
912        resolve_opt_source(&mut self.document_orientation_model)?;
913        resolve_opt_source(&mut self.document_rectification_model)?;
914        resolve_opt_source(&mut self.region_detection_model)?;
915        resolve_opt_source(&mut self.table_classification_model)?;
916        resolve_opt_source(&mut self.table_orientation_model)?;
917        resolve_opt_source(&mut self.table_cell_detection_model)?;
918        resolve_opt_source(&mut self.table_structure_recognition_model)?;
919        resolve_opt_path(&mut self.table_structure_dict_path)?;
920        resolve_opt_source(&mut self.wired_table_structure_model)?;
921        resolve_opt_source(&mut self.wireless_table_structure_model)?;
922        resolve_opt_source(&mut self.wired_table_cell_model)?;
923        resolve_opt_source(&mut self.wireless_table_cell_model)?;
924        resolve_opt_source(&mut self.formula_recognition_model)?;
925        resolve_opt_path(&mut self.formula_tokenizer_path)?;
926        resolve_opt_source(&mut self.seal_text_detection_model)?;
927        resolve_opt_source(&mut self.text_detection_model)?;
928        resolve_opt_source(&mut self.text_line_orientation_model)?;
929        resolve_opt_source(&mut self.text_recognition_model)?;
930        resolve_opt_path(&mut self.character_dict_path)?;
931
932        let cpu_region_batch_size = default_cpu_region_batch_size(
933            self.text_recognition_model.as_ref(),
934            self.text_recognition_model_name.as_deref(),
935        );
936        (self.image_batch_size, self.region_batch_size) = resolve_device_batch_sizes(
937            self.ort_session_config.as_ref(),
938            self.image_batch_size,
939            self.region_batch_size,
940            1,
941            cpu_region_batch_size,
942            self.gpu_memory_budget,
943        );
944
945        // Load character dictionary if OCR is enabled
946        let char_dict = if let Some(ref dict_path) = self.character_dict_path {
947            Some(
948                std::fs::read_to_string(dict_path).map_err(|e| OCRError::InvalidInput {
949                    message: format!(
950                        "Failed to read character dictionary from '{}': {}",
951                        dict_path.display(),
952                        e
953                    ),
954                })?,
955            )
956        } else {
957            None
958        };
959
960        // Build document orientation adapter if enabled
961        let document_orientation_adapter = build_optional_adapter(
962            self.document_orientation_model.as_ref(),
963            self.ort_session_config.as_ref(),
964            DocumentOrientationAdapterBuilder::new,
965        )?;
966
967        // Build document rectification adapter if enabled
968        let rectification_adapter = build_optional_adapter(
969            self.document_rectification_model.as_ref(),
970            self.ort_session_config.as_ref(),
971            UVDocRectifierAdapterBuilder::new,
972        )?;
973
974        // Build layout detection adapter (required)
975        let mut layout_builder = LayoutDetectionAdapterBuilder::new();
976
977        // Use explicit model name or default
978        let layout_model_config = match &self.layout_model_name {
979            Some(name) => layout_model_config_for(name),
980            None => crate::domain::adapters::LayoutModelConfig::pp_doclayout_plus_l(),
981        };
982
983        let effective_layout_cfg = self
984            .layout_detection_config
985            .clone()
986            .unwrap_or_else(|| default_layout_detection_config(&layout_model_config.model_name));
987        layout_builder = layout_builder.model_config(layout_model_config);
988        layout_builder = layout_builder.with_config(effective_layout_cfg);
989
990        if let Some(ref ort_config) = self.ort_session_config {
991            layout_builder = layout_builder.with_ort_config(ort_config.clone());
992        }
993
994        let layout_detection_adapter = layout_builder.build(&self.layout_detection_model)?;
995
996        // Build region detection adapter if enabled (PP-DocBlockLayout)
997        let region_detection_adapter = if let Some(ref model_path) = self.region_detection_model {
998            use oar_ocr_core::domain::adapters::LayoutModelConfig;
999            let mut region_builder = LayoutDetectionAdapterBuilder::new();
1000
1001            // Use model name to select configuration, default to PP-DocBlockLayout
1002            let region_model_config = if let Some(ref name) = self.region_model_name {
1003                match name.to_lowercase().replace("-", "_").as_str() {
1004                    "pp_docblocklayout" => LayoutModelConfig::pp_docblocklayout(),
1005                    _ => LayoutModelConfig::pp_docblocklayout(),
1006                }
1007            } else {
1008                LayoutModelConfig::pp_docblocklayout()
1009            };
1010            region_builder = region_builder.model_config(region_model_config);
1011
1012            // PP-StructureV3 region detection uses merge_bboxes_mode="small".
1013            let mut region_cfg = LayoutDetectionConfig::default();
1014            let mut merge_modes = std::collections::HashMap::new();
1015            merge_modes.insert(
1016                "region".to_string(),
1017                crate::domain::tasks::layout_detection::MergeBboxMode::Small,
1018            );
1019            region_cfg.class_merge_modes = Some(merge_modes);
1020            region_builder = region_builder.with_config(region_cfg);
1021
1022            if let Some(ref ort_config) = self.ort_session_config {
1023                region_builder = region_builder.with_ort_config(ort_config.clone());
1024            }
1025
1026            Some(region_builder.build(model_path)?)
1027        } else {
1028            None
1029        };
1030
1031        // Build table classification adapter if enabled
1032        let table_classification_adapter =
1033            if let Some(ref model_path) = self.table_classification_model {
1034                let mut builder = TableClassificationAdapterBuilder::new();
1035
1036                if let Some(ref config) = self.table_classification_config {
1037                    builder = builder.with_config(config.clone());
1038                }
1039
1040                if let Some(ref ort_config) = self.ort_session_config {
1041                    builder = builder.with_ort_config(ort_config.clone());
1042                }
1043
1044                Some(builder.build(model_path)?)
1045            } else {
1046                None
1047            };
1048
1049        // Build table orientation adapter if enabled (reuses document orientation model)
1050        // This detects rotated tables (0°, 90°, 180°, 270°) before structure recognition
1051        let table_orientation_adapter = build_optional_adapter(
1052            self.table_orientation_model.as_ref(),
1053            self.ort_session_config.as_ref(),
1054            DocumentOrientationAdapterBuilder::new,
1055        )?;
1056
1057        // Build table cell detection adapter if enabled
1058        let table_cell_detection_adapter = if let Some(ref model_path) =
1059            self.table_cell_detection_model
1060        {
1061            let cell_type = match self.table_cell_detection_type.as_ref() {
1062                Some(Ok(cell_type)) => *cell_type,
1063                Some(Err(message)) => {
1064                    return Err(OCRError::config_error_detailed(
1065                        "table_cell_detection",
1066                        message.clone(),
1067                    ));
1068                }
1069                None => TableType::Wired,
1070            };
1071
1072            use oar_ocr_core::domain::adapters::table_cell_detection_adapter::TableCellModelConfig;
1073
1074            let model_config = match cell_type {
1075                TableType::Wired => TableCellModelConfig::rtdetr_l_wired_table_cell_det(),
1076                TableType::Wireless => TableCellModelConfig::rtdetr_l_wireless_table_cell_det(),
1077                TableType::Unknown => {
1078                    return Err(OCRError::config_error_detailed(
1079                        "table_cell_detection",
1080                        "TableType::Unknown is not valid for a table cell detector".to_string(),
1081                    ));
1082                }
1083            };
1084
1085            let mut builder = TableCellDetectionAdapterBuilder::new().model_config(model_config);
1086
1087            if let Some(ref config) = self.table_cell_detection_config {
1088                builder = builder.with_config(config.clone());
1089            }
1090
1091            if let Some(ref ort_config) = self.ort_session_config {
1092                builder = builder.with_ort_config(ort_config.clone());
1093            }
1094
1095            Some(builder.build(model_path)?)
1096        } else {
1097            None
1098        };
1099
1100        // Build table structure recognition adapter if enabled
1101        let table_structure_recognition_adapter = if let Some(ref model_path) =
1102            self.table_structure_recognition_model
1103        {
1104            let table_type = match self.table_structure_recognition_type.as_ref() {
1105                Some(Ok(table_type)) => *table_type,
1106                Some(Err(message)) => {
1107                    return Err(OCRError::config_error_detailed(
1108                        "table_structure_recognition",
1109                        message.clone(),
1110                    ));
1111                }
1112                None => TableType::Wired,
1113            };
1114            let dict_path = self
1115                    .table_structure_dict_path
1116                    .clone()
1117                    .ok_or_else(|| {
1118                        OCRError::config_error_detailed(
1119                            "table_structure_recognition",
1120                            "Dictionary path is required. Call table_structure_dict_path() when enabling table structure recognition.".to_string(),
1121                        )
1122                    })?;
1123
1124            let adapter: TableStructureRecognitionAdapter = match table_type {
1125                TableType::Wired => {
1126                    let mut builder = SLANetWiredAdapterBuilder::new().dict_path(dict_path.clone());
1127
1128                    if let Some(ref config) = self.table_structure_recognition_config {
1129                        builder = builder.with_config(config.clone());
1130                    }
1131
1132                    if let Some(ref ort_config) = self.ort_session_config {
1133                        builder = builder.with_ort_config(ort_config.clone());
1134                    }
1135
1136                    builder.build(model_path)?
1137                }
1138                TableType::Wireless => {
1139                    let mut builder =
1140                        SLANetWirelessAdapterBuilder::new().dict_path(dict_path.clone());
1141
1142                    if let Some(ref config) = self.table_structure_recognition_config {
1143                        builder = builder.with_config(config.clone());
1144                    }
1145
1146                    if let Some(ref ort_config) = self.ort_session_config {
1147                        builder = builder.with_ort_config(ort_config.clone());
1148                    }
1149
1150                    builder.build(model_path)?
1151                }
1152                TableType::Unknown => {
1153                    return Err(OCRError::config_error_detailed(
1154                        "table_structure_recognition",
1155                        "TableType::Unknown is not valid for table structure recognition"
1156                            .to_string(),
1157                    ));
1158                }
1159            };
1160
1161            Some(adapter)
1162        } else {
1163            None
1164        };
1165
1166        // Build wired/wireless table structure adapters for auto-switch (PP-StructureV3)
1167        let wired_table_structure_adapter = if let Some(ref model_path) =
1168            self.wired_table_structure_model
1169        {
1170            let dict_path = self.table_structure_dict_path.clone().ok_or_else(|| {
1171                OCRError::config_error_detailed(
1172                    "wired_table_structure",
1173                    "Dictionary path is required. Call table_structure_dict_path() when enabling table structure recognition.".to_string(),
1174                )
1175            })?;
1176
1177            let mut builder = SLANetWiredAdapterBuilder::new().dict_path(dict_path);
1178
1179            // Label the model in logs/errors with the caller-provided preset name
1180            // (e.g. `SLANeXt_wired`). The wired/wireless slot already fixes the
1181            // SLANet variant's input shape, so this is identification metadata.
1182            if let Some(ref name) = self.wired_table_structure_model_name {
1183                builder = builder.model_name(name.clone());
1184            }
1185
1186            if let Some(ref config) = self.table_structure_recognition_config {
1187                builder = builder.with_config(config.clone());
1188            }
1189
1190            if let Some(ref ort_config) = self.ort_session_config {
1191                builder = builder.with_ort_config(ort_config.clone());
1192            }
1193
1194            Some(builder.build(model_path)?)
1195        } else {
1196            None
1197        };
1198
1199        let wireless_table_structure_adapter = if let Some(ref model_path) =
1200            self.wireless_table_structure_model
1201        {
1202            let dict_path = self.table_structure_dict_path.clone().ok_or_else(|| {
1203                OCRError::config_error_detailed(
1204                    "wireless_table_structure",
1205                    "Dictionary path is required. Call table_structure_dict_path() when enabling table structure recognition.".to_string(),
1206                )
1207            })?;
1208
1209            let mut builder = SLANetWirelessAdapterBuilder::new().dict_path(dict_path);
1210
1211            if let Some(ref name) = self.wireless_table_structure_model_name {
1212                builder = builder.model_name(name.clone());
1213            }
1214
1215            if let Some(ref config) = self.table_structure_recognition_config {
1216                builder = builder.with_config(config.clone());
1217            }
1218
1219            if let Some(ref ort_config) = self.ort_session_config {
1220                builder = builder.with_ort_config(ort_config.clone());
1221            }
1222
1223            Some(builder.build(model_path)?)
1224        } else {
1225            None
1226        };
1227
1228        // Build wired/wireless table cell detection adapters for auto-switch
1229        let wired_table_cell_adapter = if let Some(ref model_path) = self.wired_table_cell_model {
1230            use oar_ocr_core::domain::adapters::table_cell_detection_adapter::TableCellModelConfig;
1231
1232            let mut model_config = TableCellModelConfig::rtdetr_l_wired_table_cell_det();
1233            // Honor the caller-provided preset name for model identification.
1234            if let Some(ref name) = self.wired_table_cell_model_name {
1235                model_config.model_name = name.clone();
1236            }
1237            let mut builder = TableCellDetectionAdapterBuilder::new().model_config(model_config);
1238
1239            if let Some(ref config) = self.table_cell_detection_config {
1240                builder = builder.with_config(config.clone());
1241            }
1242
1243            if let Some(ref ort_config) = self.ort_session_config {
1244                builder = builder.with_ort_config(ort_config.clone());
1245            }
1246
1247            Some(builder.build(model_path)?)
1248        } else {
1249            None
1250        };
1251
1252        let wireless_table_cell_adapter = if let Some(ref model_path) =
1253            self.wireless_table_cell_model
1254        {
1255            use oar_ocr_core::domain::adapters::table_cell_detection_adapter::TableCellModelConfig;
1256
1257            let mut model_config = TableCellModelConfig::rtdetr_l_wireless_table_cell_det();
1258            if let Some(ref name) = self.wireless_table_cell_model_name {
1259                model_config.model_name = name.clone();
1260            }
1261            let mut builder = TableCellDetectionAdapterBuilder::new().model_config(model_config);
1262
1263            if let Some(ref config) = self.table_cell_detection_config {
1264                builder = builder.with_config(config.clone());
1265            }
1266
1267            if let Some(ref ort_config) = self.ort_session_config {
1268                builder = builder.with_ort_config(ort_config.clone());
1269            }
1270
1271            Some(builder.build(model_path)?)
1272        } else {
1273            None
1274        };
1275
1276        // Build formula recognition adapter if enabled
1277        let formula_recognition_adapter =
1278            if let Some(ref model_path) = self.formula_recognition_model {
1279                let tokenizer_path = self.formula_tokenizer_path.as_ref().ok_or_else(|| {
1280                    OCRError::config_error_detailed(
1281                        "formula_recognition",
1282                        "Tokenizer path is required for formula recognition".to_string(),
1283                    )
1284                })?;
1285
1286                let model_type = match self.formula_recognition_type.as_ref() {
1287                    Some(Ok(model_type)) => *model_type,
1288                    Some(Err(message)) => {
1289                        return Err(OCRError::config_error_detailed(
1290                            "formula_recognition",
1291                            message.clone(),
1292                        ));
1293                    }
1294                    None => {
1295                        return Err(OCRError::config_error_detailed(
1296                            "formula_recognition",
1297                            "Formula model kind is required".to_string(),
1298                        ));
1299                    }
1300                };
1301
1302                let adapter: FormulaRecognitionAdapter = match model_type {
1303                    FormulaModelKind::PPFormulaNet => {
1304                        let mut builder = PPFormulaNetAdapterBuilder::new();
1305
1306                        builder = builder.tokenizer_path(tokenizer_path);
1307
1308                        if let Some(ref config) = self.formula_recognition_config {
1309                            builder = builder.task_config(config.clone());
1310                        }
1311
1312                        if let Some(ort_config) = self
1313                            .formula_ort_session_config
1314                            .as_ref()
1315                            .or(self.ort_session_config.as_ref())
1316                        {
1317                            builder = builder.with_ort_config(ort_config.clone());
1318                        }
1319
1320                        builder.build(model_path)?
1321                    }
1322                    FormulaModelKind::UniMERNet => {
1323                        let mut builder = UniMERNetAdapterBuilder::new();
1324
1325                        builder = builder.tokenizer_path(tokenizer_path);
1326
1327                        if let Some(ref config) = self.formula_recognition_config {
1328                            builder = builder.task_config(config.clone());
1329                        }
1330
1331                        if let Some(ort_config) = self
1332                            .formula_ort_session_config
1333                            .as_ref()
1334                            .or(self.ort_session_config.as_ref())
1335                        {
1336                            builder = builder.with_ort_config(ort_config.clone());
1337                        }
1338
1339                        builder.build(model_path)?
1340                    }
1341                };
1342
1343                Some(adapter)
1344            } else {
1345                None
1346            };
1347
1348        // Build seal text detection adapter if enabled
1349        let seal_text_detection_adapter =
1350            if let Some(ref model_path) = self.seal_text_detection_model {
1351                let mut builder = SealTextDetectionAdapterBuilder::new();
1352
1353                if let Some(ref ort_config) = self.ort_session_config {
1354                    builder = builder.with_ort_config(ort_config.clone());
1355                }
1356
1357                Some(builder.build(model_path)?)
1358            } else {
1359                None
1360            };
1361
1362        // Build text detection adapter if enabled.
1363        //
1364        // PP-StructureV3 overall OCR uses DB preprocess with:
1365        // - limit_side_len=736
1366        // - limit_type="min"
1367        // - max_side_limit=4000
1368        // We fill these defaults here (only for the structure pipeline) unless the caller
1369        // explicitly overrides them via `text_detection_config`.
1370        let text_detection_adapter = if let Some(ref model_path) = self.text_detection_model {
1371            let mut builder = TextDetectionAdapterBuilder::new();
1372
1373            let mut effective_cfg = self.text_detection_config.clone().unwrap_or_default();
1374
1375            // Table-heavy documents are sensitive to detection fragmentation.
1376            // Match PaddleX's lower table-scene threshold when users don't override config.
1377            let has_table_pipeline = self.table_classification_model.is_some()
1378                || self.table_structure_recognition_model.is_some()
1379                || self.wired_table_structure_model.is_some()
1380                || self.wireless_table_structure_model.is_some()
1381                || self.table_cell_detection_model.is_some()
1382                || self.wired_table_cell_model.is_some()
1383                || self.wireless_table_cell_model.is_some();
1384            if self.text_detection_config.is_none() && has_table_pipeline {
1385                effective_cfg.box_threshold = 0.4;
1386            }
1387
1388            if effective_cfg.limit_side_len.is_none() {
1389                effective_cfg.limit_side_len = Some(736);
1390            }
1391            if effective_cfg.limit_type.is_none() {
1392                effective_cfg.limit_type = Some(crate::processors::LimitType::Min);
1393            }
1394            if effective_cfg.max_side_len.is_none() {
1395                effective_cfg.max_side_len = Some(4000);
1396            }
1397            builder = builder.with_config(effective_cfg);
1398
1399            // Label the detector with the caller-provided preset name (e.g.
1400            // `PP-OCRv5_server_det`) for logs/errors. Detection behavior is
1401            // driven by the config and ONNX model, not the name.
1402            if let Some(ref name) = self.text_detection_model_name {
1403                builder = builder.model_name(name.clone());
1404            }
1405
1406            if let Some(ref ort_config) = self.ort_session_config {
1407                builder = builder.with_ort_config(ort_config.clone());
1408            }
1409
1410            Some(builder.build(model_path)?)
1411        } else {
1412            None
1413        };
1414
1415        // Build text line orientation adapter if enabled (PP-StructureV3)
1416        let text_line_orientation_adapter =
1417            if let Some(ref model_path) = self.text_line_orientation_model {
1418                let mut builder = TextLineOrientationAdapterBuilder::new();
1419
1420                if let Some(ref ort_config) = self.ort_session_config {
1421                    builder = builder.with_ort_config(ort_config.clone());
1422                }
1423
1424                Some(builder.build(model_path)?)
1425            } else {
1426                None
1427            };
1428
1429        // Build text recognition adapter if enabled
1430        let text_recognition_adapter = if let Some(ref model_path) = self.text_recognition_model {
1431            let dict = char_dict.ok_or_else(|| OCRError::InvalidInput {
1432                message: "Character dictionary is required for text recognition".to_string(),
1433            })?;
1434
1435            // Parse dict into Vec<String> - one character per line
1436            let char_vec: Vec<String> = dict.lines().map(|s| s.to_string()).collect();
1437
1438            let mut builder = TextRecognitionAdapterBuilder::new().character_dict(char_vec);
1439
1440            if let Some(ref config) = self.text_recognition_config {
1441                builder = builder.with_config(config.clone());
1442            }
1443
1444            // Label the recognizer with the caller-provided preset name (e.g.
1445            // `PP-OCRv5_server_rec`) for logs/errors. Recognition behavior is
1446            // driven by the config, dictionary, and ONNX model, not the name.
1447            if let Some(ref name) = self.text_recognition_model_name {
1448                builder = builder.model_name(name.clone());
1449            }
1450
1451            if let Some(ref ort_config) = self.ort_session_config {
1452                builder = builder.with_ort_config(ort_config.clone());
1453            }
1454
1455            Some(builder.build(model_path)?)
1456        } else {
1457            None
1458        };
1459
1460        let pipeline = StructurePipeline {
1461            document_orientation_adapter,
1462            rectification_adapter,
1463            layout_detection_adapter,
1464            region_detection_adapter,
1465            table_classification_adapter,
1466            table_orientation_adapter,
1467            table_cell_detection_adapter,
1468            table_structure_recognition_adapter,
1469            wired_table_structure_adapter,
1470            wireless_table_structure_adapter,
1471            wired_table_cell_adapter,
1472            wireless_table_cell_adapter,
1473            use_e2e_wired_table_rec: self.use_e2e_wired_table_rec,
1474            use_e2e_wireless_table_rec: self.use_e2e_wireless_table_rec,
1475            use_wired_table_cells_trans_to_html: self.use_wired_table_cells_trans_to_html,
1476            use_wireless_table_cells_trans_to_html: self.use_wireless_table_cells_trans_to_html,
1477            formula_recognition_adapter,
1478            seal_text_detection_adapter,
1479            text_detection_adapter,
1480            text_line_orientation_adapter,
1481            text_recognition_adapter,
1482            image_batch_size: self.image_batch_size,
1483            region_batch_size: self.region_batch_size,
1484        };
1485
1486        Ok(OARStructure { pipeline })
1487    }
1488
1489    fn validate_batch_size(field: &str, size: usize) -> Result<(), OCRError> {
1490        if size == 0 || size > Self::MAX_BATCH_SIZE {
1491            return Err(OCRError::validation_error(
1492                "OARStructureBuilder",
1493                field,
1494                &format!("1..={}", Self::MAX_BATCH_SIZE),
1495                &size.to_string(),
1496            ));
1497        }
1498
1499        Ok(())
1500    }
1501}
1502
1503/// Runtime for document structure analysis.
1504///
1505/// This struct represents a configured and ready-to-use document structure analyzer.
1506#[derive(Debug)]
1507pub struct OARStructure {
1508    pipeline: StructurePipeline,
1509}
1510
1511/// Intermediate result from preprocessing and layout detection for a single page.
1512/// Produced by `OARStructure::prepare_page` and consumed by `complete_page`.
1513struct PreparedPage {
1514    current_image: std::sync::Arc<image::RgbImage>,
1515    orientation_angle: Option<f32>,
1516    rectified_img: Option<std::sync::Arc<image::RgbImage>>,
1517    rotation: Option<crate::oarocr::preprocess::OrientationCorrection>,
1518    layout_elements: Vec<crate::domain::structure::LayoutElement>,
1519    detected_region_blocks: Option<Vec<crate::domain::structure::RegionBlock>>,
1520    precomputed_text_regions: Option<Vec<crate::oarocr::TextRegion>>,
1521}
1522
1523impl OARStructure {
1524    fn finish_layout_elements(layout_elements: &mut Vec<crate::domain::structure::LayoutElement>) {
1525        if layout_elements.len() > 1 {
1526            let removed = crate::domain::structure::remove_overlapping_layout_elements(
1527                layout_elements,
1528                LAYOUT_OVERLAP_IOU_THRESHOLD,
1529            );
1530            if removed > 0 {
1531                tracing::info!(
1532                    "Removing {} overlapping layout elements (threshold={})",
1533                    removed,
1534                    LAYOUT_OVERLAP_IOU_THRESHOLD
1535                );
1536            }
1537        }
1538
1539        crate::domain::structure::apply_standardized_layout_label_fixes(layout_elements);
1540    }
1541
1542    fn layout_elements_from_detection(
1543        elements: &[oar_ocr_core::domain::tasks::LayoutDetectionElement],
1544    ) -> Vec<crate::domain::structure::LayoutElement> {
1545        use oar_ocr_core::domain::structure::LayoutElementType;
1546
1547        elements
1548            .iter()
1549            .map(|element| {
1550                let element_type_enum = LayoutElementType::from_label(&element.element_type);
1551                crate::domain::structure::LayoutElement::new(
1552                    element.bbox.clone(),
1553                    element_type_enum,
1554                    element.score,
1555                )
1556                .with_label(element.element_type.clone())
1557            })
1558            .collect()
1559    }
1560
1561    /// Refinement of overall OCR results using layout boxes.
1562    ///
1563    /// This mirrors two behaviors in `layout_parsing/pipeline_v2.py`:
1564    /// 1) If a single overall OCR box overlaps multiple layout blocks, re-recognize
1565    ///    the intersection crop per block and replace/append OCR entries.
1566    /// 2) If a non-vision layout block has no matched OCR text, run recognition
1567    ///    on the layout bbox crop as a fallback.
1568    ///
1569    /// We approximate poly handling with AABB intersections. The resulting
1570    /// crops are stored into `TextRegion` with `dt_poly/rec_poly` set to the crop box.
1571    fn refine_overall_ocr_with_layout(
1572        text_regions: &mut Vec<crate::oarocr::TextRegion>,
1573        layout_elements: &[crate::domain::structure::LayoutElement],
1574        region_blocks: Option<&[crate::domain::structure::RegionBlock]>,
1575        page_image: &image::RgbImage,
1576        text_recognition_adapter: &TextRecognitionAdapter,
1577        region_batch_size: usize,
1578    ) -> Result<(), OCRError> {
1579        use oar_ocr_core::core::traits::task::ImageTaskInput;
1580        use oar_ocr_core::domain::structure::LayoutElementType;
1581        use oar_ocr_core::processors::BoundingBox;
1582        use oar_ocr_core::utils::BBoxCrop;
1583
1584        if text_regions.is_empty() || layout_elements.is_empty() {
1585            return Ok(());
1586        }
1587
1588        fn aabb_intersection(b1: &BoundingBox, b2: &BoundingBox) -> Option<BoundingBox> {
1589            let x1 = b1.x_min().max(b2.x_min());
1590            let y1 = b1.y_min().max(b2.y_min());
1591            let x2 = b1.x_max().min(b2.x_max());
1592            let y2 = b1.y_max().min(b2.y_max());
1593            if x2 - x1 <= 1.0 || y2 - y1 <= 1.0 {
1594                None
1595            } else {
1596                Some(BoundingBox::from_coords(x1, y1, x2, y2))
1597            }
1598        }
1599
1600        // Layout boxes that participate in OCR matching (exclude specialized types).
1601        let is_excluded_layout = |t: LayoutElementType| {
1602            matches!(
1603                t,
1604                LayoutElementType::Formula
1605                    | LayoutElementType::FormulaNumber
1606                    | LayoutElementType::Table
1607                    | LayoutElementType::Seal
1608            )
1609        };
1610
1611        // Build overlap maps: ocr_idx -> layout_idxes.
1612        // `get_sub_regions_ocr_res` uses get_overlap_boxes_idx:
1613        // any overlap with intersection width/height >3px counts as a match (no ratio threshold).
1614        let min_pixels = 3.0;
1615        let mut matched_ocr: Vec<Vec<usize>> = vec![Vec::new(); text_regions.len()];
1616        for (ocr_idx, region) in text_regions.iter().enumerate() {
1617            for (layout_idx, elem) in layout_elements.iter().enumerate() {
1618                if is_excluded_layout(elem.element_type) {
1619                    continue;
1620                }
1621                let inter_x_min = region.bounding_box.x_min().max(elem.bbox.x_min());
1622                let inter_y_min = region.bounding_box.y_min().max(elem.bbox.y_min());
1623                let inter_x_max = region.bounding_box.x_max().min(elem.bbox.x_max());
1624                let inter_y_max = region.bounding_box.y_max().min(elem.bbox.y_max());
1625                if inter_x_max - inter_x_min > min_pixels && inter_y_max - inter_y_min > min_pixels
1626                {
1627                    matched_ocr[ocr_idx].push(layout_idx);
1628                }
1629            }
1630        }
1631
1632        // 1) Cross-layout re-recognition for OCR boxes matched to multiple blocks.
1633        let mut appended_regions: Vec<crate::oarocr::TextRegion> = Vec::new();
1634        let original_ocr_len = text_regions.len();
1635        let mut multi_layout_ocr_count = 0usize;
1636        let mut multi_layout_crop_count = 0usize;
1637
1638        for ocr_idx in 0..original_ocr_len {
1639            let layout_ids = matched_ocr[ocr_idx].clone();
1640            if layout_ids.len() <= 1 {
1641                continue;
1642            }
1643            multi_layout_ocr_count += 1;
1644
1645            let ocr_box = text_regions[ocr_idx].bounding_box.clone();
1646
1647            let mut crops: Vec<image::RgbImage> = Vec::new();
1648            let mut crop_boxes: Vec<(BoundingBox, bool)> = Vec::new(); // (bbox, is_first)
1649
1650            for (j, layout_idx) in layout_ids.iter().enumerate() {
1651                let layout_box = &layout_elements[*layout_idx].bbox;
1652                let Some(crop_box) = aabb_intersection(&ocr_box, layout_box) else {
1653                    continue;
1654                };
1655
1656                // Suppress existing OCR text fully covered by this crop (IoU > 0.8).
1657                for (other_idx, other_region) in text_regions.iter_mut().enumerate() {
1658                    if other_idx == ocr_idx {
1659                        continue;
1660                    }
1661                    if other_region.bounding_box.iou(&crop_box) > 0.8 {
1662                        other_region.text = None;
1663                    }
1664                }
1665
1666                if let Ok(crop_img) = BBoxCrop::crop_bounding_box(page_image, &crop_box) {
1667                    crops.push(crop_img);
1668                    crop_boxes.push((crop_box, j == 0));
1669                }
1670            }
1671            multi_layout_crop_count += crop_boxes.len();
1672
1673            if crops.is_empty() {
1674                continue;
1675            }
1676
1677            // Run recognition on all crops (batched).
1678            let mut rec_texts: Vec<String> = Vec::with_capacity(crops.len());
1679            let mut rec_scores: Vec<f32> = Vec::with_capacity(crops.len());
1680
1681            for batch_start in (0..crops.len()).step_by(region_batch_size.max(1)) {
1682                let batch_end = (batch_start + region_batch_size).min(crops.len());
1683                let batch: Vec<_> = crops[batch_start..batch_end].to_vec();
1684                let rec_input = ImageTaskInput::new(batch);
1685                let rec_result = text_recognition_adapter.execute(rec_input, None)?;
1686                rec_texts.extend(rec_result.texts);
1687                rec_scores.extend(rec_result.scores);
1688            }
1689
1690            for ((crop_box, is_first), (text, score)) in crop_boxes
1691                .into_iter()
1692                .zip(rec_texts.into_iter().zip(rec_scores))
1693            {
1694                if text.is_empty() {
1695                    continue;
1696                }
1697                if is_first {
1698                    text_regions[ocr_idx].bounding_box = crop_box.clone();
1699                    text_regions[ocr_idx].dt_poly = Some(crop_box.clone());
1700                    text_regions[ocr_idx].rec_poly = Some(crop_box.clone());
1701                    text_regions[ocr_idx].text = Some(Arc::from(text));
1702                    text_regions[ocr_idx].confidence = Some(score);
1703                } else {
1704                    appended_regions.push(crate::oarocr::TextRegion {
1705                        bounding_box: crop_box.clone(),
1706                        dt_poly: Some(crop_box.clone()),
1707                        rec_poly: Some(crop_box),
1708                        text: Some(Arc::from(text)),
1709                        confidence: Some(score),
1710                        orientation_angle: None,
1711                        word_boxes: None,
1712                        label: None,
1713                    });
1714                }
1715            }
1716        }
1717
1718        if !appended_regions.is_empty() {
1719            text_regions.extend(appended_regions);
1720        }
1721
1722        // 2) Layout-bbox fallback OCR for blocks with no matched text.
1723        // Prefer region blocks for hierarchy if present, but OCR fallback is driven by layout boxes.
1724        let mut fallback_blocks = 0usize;
1725        for elem in layout_elements.iter() {
1726            if is_excluded_layout(elem.element_type) {
1727                continue;
1728            }
1729            if matches!(
1730                elem.element_type,
1731                LayoutElementType::Image | LayoutElementType::Chart
1732            ) {
1733                continue;
1734            }
1735
1736            let mut has_text = false;
1737            for region in text_regions.iter() {
1738                if !region.text.as_ref().map(|t| !t.is_empty()).unwrap_or(false) {
1739                    continue;
1740                }
1741                let inter_x_min = region.bounding_box.x_min().max(elem.bbox.x_min());
1742                let inter_y_min = region.bounding_box.y_min().max(elem.bbox.y_min());
1743                let inter_x_max = region.bounding_box.x_max().min(elem.bbox.x_max());
1744                let inter_y_max = region.bounding_box.y_max().min(elem.bbox.y_max());
1745                if inter_x_max - inter_x_min > min_pixels && inter_y_max - inter_y_min > min_pixels
1746                {
1747                    has_text = true;
1748                    break;
1749                }
1750            }
1751
1752            if has_text {
1753                continue;
1754            }
1755            fallback_blocks += 1;
1756
1757            // Crop layout bbox and run recognition.
1758            if let Ok(crop_img) = BBoxCrop::crop_bounding_box(page_image, &elem.bbox) {
1759                let rec_input = ImageTaskInput::new(vec![crop_img]);
1760                let rec_result = text_recognition_adapter.execute(rec_input, None)?;
1761                if let (Some(text), Some(score)) =
1762                    (rec_result.texts.first(), rec_result.scores.first())
1763                    && !text.is_empty()
1764                {
1765                    let crop_box = elem.bbox.clone();
1766                    text_regions.push(crate::oarocr::TextRegion {
1767                        bounding_box: crop_box.clone(),
1768                        dt_poly: Some(crop_box.clone()),
1769                        rec_poly: Some(crop_box),
1770                        text: Some(Arc::from(text.as_str())),
1771                        confidence: Some(*score),
1772                        orientation_angle: None,
1773                        word_boxes: None,
1774                        label: None,
1775                    });
1776                }
1777            }
1778        }
1779
1780        tracing::info!(
1781            "overall OCR refine: multi-layout OCR boxes={}, crops={}, fallback layout blocks={}",
1782            multi_layout_ocr_count,
1783            multi_layout_crop_count,
1784            fallback_blocks
1785        );
1786
1787        // Region blocks currently do not require special handling here; they are only
1788        // used for ordering later. Kept as a parameter for future parity work.
1789        let _ = region_blocks;
1790
1791        Ok(())
1792    }
1793
1794    /// Split OCR bounding boxes based on table cell boundaries when they span multiple cells.
1795    ///
1796    /// This mirrors `split_ocr_bboxes_by_table_cells`:
1797    /// - For each OCR box that overlaps >= k cells (by intersection / cell_area > 0.5),
1798    ///   split the box vertically at cell boundaries
1799    /// - Re-run text recognition on each split crop
1800    /// - Replace the original OCR box/text with the split boxes/texts
1801    fn split_ocr_bboxes_by_table_cells(
1802        tables: &[TableResult],
1803        text_regions: &mut Vec<crate::oarocr::TextRegion>,
1804        page_image: &image::RgbImage,
1805        text_recognition_adapter: &TextRecognitionAdapter,
1806    ) -> Result<(), OCRError> {
1807        use oar_ocr_core::core::traits::task::ImageTaskInput;
1808        use oar_ocr_core::processors::BoundingBox;
1809
1810        // Collect all cell boxes in [x1, y1, x2, y2] format
1811        let mut cell_boxes: Vec<[f32; 4]> = Vec::new();
1812        for table in tables {
1813            for cell in &table.cells {
1814                let x1 = cell.bbox.x_min();
1815                let y1 = cell.bbox.y_min();
1816                let x2 = cell.bbox.x_max();
1817                let y2 = cell.bbox.y_max();
1818                if x2 > x1 && y2 > y1 {
1819                    cell_boxes.push([x1, y1, x2, y2]);
1820                }
1821            }
1822        }
1823
1824        if cell_boxes.is_empty() || text_regions.is_empty() {
1825            return Ok(());
1826        }
1827
1828        // Calculate intersection / cell_area (matches calculate_iou in split_ocr_bboxes_by_table_cells)
1829        fn overlap_ratio_box_over_cell(box1: &[f32; 4], box2: &[f32; 4]) -> f32 {
1830            let x_left = box1[0].max(box2[0]);
1831            let y_top = box1[1].max(box2[1]);
1832            let x_right = box1[2].min(box2[2]);
1833            let y_bottom = box1[3].min(box2[3]);
1834
1835            if x_right <= x_left || y_bottom <= y_top {
1836                return 0.0;
1837            }
1838
1839            let inter_area = (x_right - x_left) * (y_bottom - y_top);
1840            let cell_area = (box2[2] - box2[0]) * (box2[3] - box2[1]);
1841            if cell_area <= 0.0 {
1842                0.0
1843            } else {
1844                inter_area / cell_area
1845            }
1846        }
1847
1848        // Find cells that significantly overlap with an OCR box
1849        fn get_overlapping_cells(
1850            ocr_box: &[f32; 4],
1851            cells: &[[f32; 4]],
1852            threshold: f32,
1853        ) -> Vec<usize> {
1854            let mut overlapping = Vec::new();
1855            for (idx, cell) in cells.iter().enumerate() {
1856                if overlap_ratio_box_over_cell(ocr_box, cell) > threshold {
1857                    overlapping.push(idx);
1858                }
1859            }
1860            // Sort by cell x1 (left to right)
1861            overlapping.sort_by(|&i, &j| {
1862                cells[i][0]
1863                    .partial_cmp(&cells[j][0])
1864                    .unwrap_or(std::cmp::Ordering::Equal)
1865            });
1866            overlapping
1867        }
1868
1869        // Split an OCR box vertically at cell boundaries.
1870        fn split_box_by_cells(
1871            ocr_box: &[f32; 4],
1872            cell_indices: &[usize],
1873            cells: &[[f32; 4]],
1874        ) -> Vec<[f32; 4]> {
1875            if cell_indices.is_empty() {
1876                return vec![*ocr_box];
1877            }
1878
1879            let mut split_boxes: Vec<[f32; 4]> = Vec::new();
1880            let cells_to_split: Vec<[f32; 4]> = cell_indices.iter().map(|&i| cells[i]).collect();
1881
1882            // Leading segment before first cell
1883            if ocr_box[0] < cells_to_split[0][0] {
1884                split_boxes.push([ocr_box[0], ocr_box[1], cells_to_split[0][0], ocr_box[3]]);
1885            }
1886
1887            // Segments overlapping each cell and gaps between cells
1888            for (i, current_cell) in cells_to_split.iter().enumerate() {
1889                // Cell overlap segment
1890                split_boxes.push([
1891                    ocr_box[0].max(current_cell[0]),
1892                    ocr_box[1],
1893                    ocr_box[2].min(current_cell[2]),
1894                    ocr_box[3],
1895                ]);
1896
1897                // Gap between this cell and the next cell
1898                if i + 1 < cells_to_split.len() {
1899                    let next_cell = cells_to_split[i + 1];
1900                    if current_cell[2] < next_cell[0] {
1901                        split_boxes.push([current_cell[2], ocr_box[1], next_cell[0], ocr_box[3]]);
1902                    }
1903                }
1904            }
1905
1906            // Trailing segment after last cell
1907            let last_cell = cells_to_split[cells_to_split.len() - 1];
1908            if last_cell[2] < ocr_box[2] {
1909                split_boxes.push([last_cell[2], ocr_box[1], ocr_box[2], ocr_box[3]]);
1910            }
1911
1912            // Deduplicate boxes
1913            let mut unique = Vec::new();
1914            let mut seen = std::collections::HashSet::new();
1915            for b in split_boxes {
1916                let key = (
1917                    b[0].to_bits(),
1918                    b[1].to_bits(),
1919                    b[2].to_bits(),
1920                    b[3].to_bits(),
1921                );
1922                if seen.insert(key) {
1923                    unique.push(b);
1924                }
1925            }
1926            unique
1927        }
1928
1929        let k_min_cells = 2usize;
1930        let overlap_threshold = CELL_OVERLAP_IOU_THRESHOLD;
1931
1932        let mut new_regions: Vec<crate::oarocr::TextRegion> =
1933            Vec::with_capacity(text_regions.len());
1934
1935        for region in text_regions.iter() {
1936            let ocr_box = [
1937                region.bounding_box.x_min(),
1938                region.bounding_box.y_min(),
1939                region.bounding_box.x_max(),
1940                region.bounding_box.y_max(),
1941            ];
1942
1943            let overlapping_cells = get_overlapping_cells(&ocr_box, &cell_boxes, overlap_threshold);
1944
1945            // If OCR box does not span multiple cells, keep as-is
1946            if overlapping_cells.len() < k_min_cells {
1947                new_regions.push(region.clone());
1948                continue;
1949            }
1950
1951            let split_boxes = split_box_by_cells(&ocr_box, &overlapping_cells, &cell_boxes);
1952
1953            for box_coords in split_boxes {
1954                // Convert to integer crop coordinates, clamp to image bounds
1955                let img_w = page_image.width() as i32;
1956                let img_h = page_image.height() as i32;
1957
1958                let mut x1 = box_coords[0].floor() as i32;
1959                let mut y1 = box_coords[1].floor() as i32;
1960                let mut x2 = box_coords[2].ceil() as i32;
1961                let mut y2 = box_coords[3].ceil() as i32;
1962
1963                x1 = x1.clamp(0, img_w.saturating_sub(1));
1964                y1 = y1.clamp(0, img_h.saturating_sub(1));
1965                x2 = x2.clamp(0, img_w);
1966                y2 = y2.clamp(0, img_h);
1967
1968                if x2 - x1 <= 1 || y2 - y1 <= 1 {
1969                    continue;
1970                }
1971
1972                let crop_w = (x2 - x1) as u32;
1973                let crop_h = (y2 - y1) as u32;
1974                if crop_w <= 1 || crop_h <= 1 {
1975                    continue;
1976                }
1977
1978                let x1u = x1 as u32;
1979                let y1u = y1 as u32;
1980                if x1u >= page_image.width() || y1u >= page_image.height() {
1981                    continue;
1982                }
1983                let crop_w = crop_w.min(page_image.width() - x1u);
1984                let crop_h = crop_h.min(page_image.height() - y1u);
1985                if crop_w <= 1 || crop_h <= 1 {
1986                    continue;
1987                }
1988
1989                let crop =
1990                    image::imageops::crop_imm(page_image, x1u, y1u, crop_w, crop_h).to_image();
1991
1992                let rec_input = ImageTaskInput::new(vec![crop]);
1993                let rec_result = text_recognition_adapter.execute(rec_input, None)?;
1994                if let (Some(text), Some(score)) =
1995                    (rec_result.texts.first(), rec_result.scores.first())
1996                    && !text.is_empty()
1997                {
1998                    let bbox = BoundingBox::from_coords(
1999                        box_coords[0],
2000                        box_coords[1],
2001                        box_coords[2],
2002                        box_coords[3],
2003                    );
2004                    new_regions.push(crate::oarocr::TextRegion {
2005                        bounding_box: bbox.clone(),
2006                        dt_poly: Some(bbox.clone()),
2007                        rec_poly: Some(bbox),
2008                        text: Some(Arc::from(text.as_str())),
2009                        confidence: Some(*score),
2010                        orientation_angle: None,
2011                        word_boxes: None,
2012                        label: None,
2013                    });
2014                }
2015            }
2016        }
2017
2018        *text_regions = new_regions;
2019        Ok(())
2020    }
2021
2022    fn detect_layout_and_regions(
2023        &self,
2024        page_image: &image::RgbImage,
2025    ) -> Result<
2026        (
2027            Vec<crate::domain::structure::LayoutElement>,
2028            Option<Vec<crate::domain::structure::RegionBlock>>,
2029        ),
2030        OCRError,
2031    > {
2032        use oar_ocr_core::core::traits::task::ImageTaskInput;
2033        use oar_ocr_core::domain::structure::RegionBlock;
2034
2035        let input = ImageTaskInput::new(vec![page_image.clone()]);
2036        let t_layout = Instant::now();
2037        let layout_result = self
2038            .pipeline
2039            .layout_detection_adapter
2040            .execute(input, None)?;
2041        let layout_dur = t_layout.elapsed();
2042
2043        let mut layout_elements = layout_result
2044            .elements
2045            .first()
2046            .map(|elements| Self::layout_elements_from_detection(elements))
2047            .unwrap_or_default();
2048
2049        let mut detected_region_blocks: Option<Vec<RegionBlock>> = None;
2050        if let Some(ref region_adapter) = self.pipeline.region_detection_adapter {
2051            let region_input = ImageTaskInput::new(vec![page_image.clone()]);
2052            let t_region = Instant::now();
2053            if let Ok(region_result) = region_adapter.execute(region_input, None)
2054                && let Some(region_elements) = region_result.elements.first()
2055                && !region_elements.is_empty()
2056            {
2057                let blocks: Vec<RegionBlock> = region_elements
2058                    .iter()
2059                    .map(|e| RegionBlock {
2060                        bbox: e.bbox.clone(),
2061                        confidence: e.score,
2062                        order_index: None,
2063                        element_indices: Vec::new(),
2064                    })
2065                    .collect();
2066                detected_region_blocks = Some(blocks);
2067            }
2068            tracing::debug!(
2069                "structure stage: region detection {:.1} ms, blocks={}",
2070                t_region.elapsed().as_secs_f64() * 1000.0,
2071                detected_region_blocks.as_ref().map_or(0, Vec::len)
2072            );
2073        }
2074
2075        Self::finish_layout_elements(&mut layout_elements);
2076        tracing::debug!(
2077            "structure stage: layout detection {:.1} ms, elements={}",
2078            layout_dur.as_secs_f64() * 1000.0,
2079            layout_elements.len()
2080        );
2081
2082        Ok((layout_elements, detected_region_blocks))
2083    }
2084
2085    fn recognize_formulas(
2086        &self,
2087        page_image: &image::RgbImage,
2088        layout_elements: &[crate::domain::structure::LayoutElement],
2089    ) -> Result<Vec<crate::domain::structure::FormulaResult>, OCRError> {
2090        use oar_ocr_core::core::traits::task::ImageTaskInput;
2091        use oar_ocr_core::domain::structure::FormulaResult;
2092        use oar_ocr_core::utils::BBoxCrop;
2093
2094        let Some(ref formula_adapter) = self.pipeline.formula_recognition_adapter else {
2095            return Ok(Vec::new());
2096        };
2097
2098        let formula_elements: Vec<_> = layout_elements
2099            .iter()
2100            .filter(|e| e.element_type.is_formula())
2101            .collect();
2102
2103        if formula_elements.is_empty() {
2104            tracing::debug!(
2105                "Formula recognition skipped: no formula regions from layout detection"
2106            );
2107            return Ok(Vec::new());
2108        }
2109
2110        let mut crops = Vec::new();
2111        let mut bboxes = Vec::new();
2112
2113        for elem in &formula_elements {
2114            match BBoxCrop::crop_bounding_box(page_image, &elem.bbox) {
2115                Ok(crop) => {
2116                    crops.push(crop);
2117                    bboxes.push(elem.bbox.clone());
2118                }
2119                Err(err) => {
2120                    tracing::warn!("Formula region crop failed: {}", err);
2121                }
2122            }
2123        }
2124
2125        if crops.is_empty() {
2126            tracing::debug!(
2127                "Formula recognition skipped: all formula crops failed for {} regions",
2128                formula_elements.len()
2129            );
2130            return Ok(Vec::new());
2131        }
2132
2133        let t_formula = Instant::now();
2134        let batch_size = formula_adapter.recommended_batch_size().max(1);
2135        let crop_count = bboxes.len();
2136        let mut formula_results = Vec::with_capacity(crop_count);
2137        let mut score_results = Vec::with_capacity(crop_count);
2138        let mut remaining_crops = crops.into_iter();
2139        loop {
2140            let chunk_vec: Vec<_> = remaining_crops.by_ref().take(batch_size).collect();
2141            if chunk_vec.is_empty() {
2142                break;
2143            }
2144
2145            let output = formula_adapter.execute(ImageTaskInput::new(chunk_vec), None)?;
2146            formula_results.extend(output.formulas);
2147            score_results.extend(output.scores);
2148        }
2149        tracing::debug!(
2150            "structure stage: formula recognition {:.1} ms, crops={}, batches={}, batch_size={}",
2151            t_formula.elapsed().as_secs_f64() * 1000.0,
2152            crop_count,
2153            crop_count.div_ceil(batch_size),
2154            batch_size
2155        );
2156
2157        let mut formulas = Vec::new();
2158        for ((bbox, formula), score) in bboxes.into_iter().zip(formula_results).zip(score_results) {
2159            let width = bbox.x_max() - bbox.x_min();
2160            let height = bbox.y_max() - bbox.y_min();
2161            if width <= 0.0 || height <= 0.0 {
2162                tracing::warn!(
2163                    "Skipping formula with non-positive bbox dimensions: w={:.2}, h={:.2}",
2164                    width,
2165                    height
2166                );
2167                continue;
2168            }
2169
2170            formulas.push(FormulaResult {
2171                bbox,
2172                latex: formula,
2173                confidence: score.unwrap_or(0.0),
2174            });
2175        }
2176
2177        Ok(formulas)
2178    }
2179
2180    fn detect_seal_text(
2181        &self,
2182        page_image: &image::RgbImage,
2183        layout_elements: &mut Vec<crate::domain::structure::LayoutElement>,
2184    ) -> Result<(), OCRError> {
2185        use oar_ocr_core::core::traits::task::ImageTaskInput;
2186        use oar_ocr_core::domain::structure::{LayoutElement, LayoutElementType};
2187        use oar_ocr_core::processors::Point;
2188        use oar_ocr_core::utils::BBoxCrop;
2189
2190        let Some(ref seal_adapter) = self.pipeline.seal_text_detection_adapter else {
2191            return Ok(());
2192        };
2193
2194        let seal_regions: Vec<_> = layout_elements
2195            .iter()
2196            .filter(|e| e.element_type == LayoutElementType::Seal)
2197            .map(|e| e.bbox.clone())
2198            .collect();
2199
2200        if seal_regions.is_empty() {
2201            tracing::debug!("Seal detection skipped: no seal regions from layout detection");
2202            return Ok(());
2203        }
2204
2205        let mut seal_crops = Vec::new();
2206        let mut crop_offsets = Vec::new();
2207
2208        for region_bbox in &seal_regions {
2209            match BBoxCrop::crop_bounding_box(page_image, region_bbox) {
2210                Ok(crop) => {
2211                    seal_crops.push(crop);
2212                    crop_offsets.push((region_bbox.x_min(), region_bbox.y_min()));
2213                }
2214                Err(err) => {
2215                    tracing::warn!("Seal region crop failed: {}", err);
2216                }
2217            }
2218        }
2219
2220        if seal_crops.is_empty() {
2221            return Ok(());
2222        }
2223
2224        let input = ImageTaskInput::new(seal_crops);
2225        let seal_result = seal_adapter.execute(input, None)?;
2226
2227        for ((dx, dy), detections) in crop_offsets.iter().zip(seal_result.detections) {
2228            for detection in detections {
2229                let translated_bbox = crate::processors::BoundingBox::new(
2230                    detection
2231                        .bbox
2232                        .points
2233                        .iter()
2234                        .map(|p| Point::new(p.x + dx, p.y + dy))
2235                        .collect(),
2236                );
2237
2238                layout_elements.push(
2239                    LayoutElement::new(translated_bbox, LayoutElementType::Seal, detection.score)
2240                        .with_label("seal".to_string()),
2241                );
2242            }
2243        }
2244
2245        Ok(())
2246    }
2247
2248    fn sort_layout_elements_enhanced(
2249        layout_elements: &mut Vec<crate::domain::structure::LayoutElement>,
2250        page_width: f32,
2251        page_height: f32,
2252    ) {
2253        use oar_ocr_core::processors::layout_sorting::{SortableElement, sort_layout_enhanced};
2254
2255        if layout_elements.is_empty() {
2256            return;
2257        }
2258
2259        let sortable_elements: Vec<_> = layout_elements
2260            .iter()
2261            .map(|e| SortableElement {
2262                bbox: e.bbox.clone(),
2263                element_type: e.element_type,
2264                num_lines: e.num_lines,
2265            })
2266            .collect();
2267
2268        let sorted_indices = sort_layout_enhanced(&sortable_elements, page_width, page_height);
2269        if sorted_indices.len() != layout_elements.len() {
2270            return;
2271        }
2272
2273        let sorted_elements: Vec<_> = sorted_indices
2274            .into_iter()
2275            .map(|idx| layout_elements[idx].clone())
2276            .collect();
2277        *layout_elements = sorted_elements;
2278    }
2279
2280    fn assign_region_block_membership(
2281        region_blocks: &mut [crate::domain::structure::RegionBlock],
2282        layout_elements: &[crate::domain::structure::LayoutElement],
2283    ) {
2284        use std::cmp::Ordering;
2285
2286        if region_blocks.is_empty() {
2287            return;
2288        }
2289
2290        region_blocks.sort_by(|a, b| {
2291            a.bbox
2292                .y_min()
2293                .partial_cmp(&b.bbox.y_min())
2294                .unwrap_or(Ordering::Equal)
2295                .then_with(|| {
2296                    a.bbox
2297                        .x_min()
2298                        .partial_cmp(&b.bbox.x_min())
2299                        .unwrap_or(Ordering::Equal)
2300                })
2301        });
2302
2303        for (i, region) in region_blocks.iter_mut().enumerate() {
2304            region.order_index = Some((i + 1) as u32);
2305            region.element_indices.clear();
2306        }
2307
2308        if layout_elements.is_empty() {
2309            return;
2310        }
2311
2312        for (elem_idx, elem) in layout_elements.iter().enumerate() {
2313            let elem_area = elem.bbox.area();
2314            if elem_area <= 0.0 {
2315                continue;
2316            }
2317
2318            let mut best_region: Option<usize> = None;
2319            let mut best_ioa = 0.0f32;
2320
2321            for (region_idx, region) in region_blocks.iter().enumerate() {
2322                let intersection = elem.bbox.intersection_area(&region.bbox);
2323                if intersection <= 0.0 {
2324                    continue;
2325                }
2326                let ioa = intersection / elem_area;
2327                if ioa > best_ioa {
2328                    best_ioa = ioa;
2329                    best_region = Some(region_idx);
2330                }
2331            }
2332
2333            if let Some(region_idx) = best_region
2334                && best_ioa >= REGION_MEMBERSHIP_IOA_THRESHOLD
2335            {
2336                region_blocks[region_idx].element_indices.push(elem_idx);
2337            }
2338        }
2339    }
2340
2341    fn run_overall_ocr(
2342        &self,
2343        page_image: &image::RgbImage,
2344        layout_elements: &[crate::domain::structure::LayoutElement],
2345        region_blocks: Option<&[crate::domain::structure::RegionBlock]>,
2346    ) -> Result<Vec<crate::oarocr::TextRegion>, OCRError> {
2347        use crate::oarocr::TextRegion;
2348        use oar_ocr_core::core::traits::task::ImageTaskInput;
2349        use std::sync::Arc;
2350
2351        let Some(ref text_detection_adapter) = self.pipeline.text_detection_adapter else {
2352            return Ok(Vec::new());
2353        };
2354        let Some(ref text_recognition_adapter) = self.pipeline.text_recognition_adapter else {
2355            return Ok(Vec::new());
2356        };
2357
2358        let mut text_regions = Vec::new();
2359
2360        // Mask formula regions before text detection only when formula
2361        // recognition is enabled. With formula recognition disabled, PaddleX
2362        // keeps formula-like regions in overall OCR output.
2363        let mut ocr_image = page_image.clone();
2364        if self.pipeline.formula_recognition_adapter.is_some() {
2365            let mask_bboxes: Vec<crate::processors::BoundingBox> = layout_elements
2366                .iter()
2367                .filter(|e| e.element_type.is_formula())
2368                .map(|e| e.bbox.clone())
2369                .collect();
2370
2371            if !mask_bboxes.is_empty() {
2372                crate::utils::mask_regions(&mut ocr_image, &mask_bboxes, [255, 255, 255]);
2373            }
2374        }
2375
2376        // Text detection (on masked image).
2377        let input = ImageTaskInput::new(vec![ocr_image.clone()]);
2378        let t_text_det = Instant::now();
2379        let det_result = text_detection_adapter.execute(input, None)?;
2380        let text_det_dur = t_text_det.elapsed();
2381
2382        let mut detection_boxes = if let Some(detections) = det_result.detections.first() {
2383            detections
2384                .iter()
2385                .map(|d| d.bbox.clone())
2386                .collect::<Vec<_>>()
2387        } else {
2388            Vec::new()
2389        };
2390
2391        // Debug: raw text detection boxes from overall OCR (before any splitting).
2392        let raw_detection_boxes = detection_boxes.clone();
2393        if tracing::enabled!(tracing::Level::DEBUG) && !raw_detection_boxes.is_empty() {
2394            let raw_rects: Vec<[f32; 4]> = raw_detection_boxes
2395                .iter()
2396                .map(|b| [b.x_min(), b.y_min(), b.x_max(), b.y_max()])
2397                .collect();
2398            tracing::debug!("overall OCR text det boxes (raw): {:?}", raw_rects);
2399        }
2400
2401        // Cross-layout re-recognition: split text det boxes that span multiple layout/region boxes.
2402        if !detection_boxes.is_empty() {
2403            let mut split_boxes = Vec::new();
2404            let mut split_count = 0usize;
2405
2406            let container_boxes: Vec<crate::processors::BoundingBox> =
2407                if let Some(regions) = region_blocks {
2408                    regions.iter().map(|r| r.bbox.clone()).collect()
2409                } else {
2410                    layout_elements
2411                        .iter()
2412                        .filter(|e| {
2413                            matches!(
2414                            e.element_type,
2415                            crate::domain::structure::LayoutElementType::DocTitle
2416                                | crate::domain::structure::LayoutElementType::ParagraphTitle
2417                                | crate::domain::structure::LayoutElementType::Text
2418                                | crate::domain::structure::LayoutElementType::Content
2419                                | crate::domain::structure::LayoutElementType::Abstract
2420                                | crate::domain::structure::LayoutElementType::Header
2421                                | crate::domain::structure::LayoutElementType::Footer
2422                                | crate::domain::structure::LayoutElementType::Footnote
2423                                | crate::domain::structure::LayoutElementType::Number
2424                                | crate::domain::structure::LayoutElementType::Reference
2425                                | crate::domain::structure::LayoutElementType::ReferenceContent
2426                                | crate::domain::structure::LayoutElementType::Algorithm
2427                                | crate::domain::structure::LayoutElementType::AsideText
2428                                | crate::domain::structure::LayoutElementType::List
2429                                | crate::domain::structure::LayoutElementType::FigureTitle
2430                                | crate::domain::structure::LayoutElementType::TableTitle
2431                                | crate::domain::structure::LayoutElementType::ChartTitle
2432                                | crate::domain::structure::LayoutElementType::FigureTableChartTitle
2433                        )
2434                        })
2435                        .map(|e| e.bbox.clone())
2436                        .collect()
2437                };
2438
2439            if !container_boxes.is_empty() {
2440                for bbox in detection_boxes.into_iter() {
2441                    let mut intersections: Vec<crate::processors::BoundingBox> = Vec::new();
2442                    let self_area = bbox.area();
2443                    if self_area <= 0.0 {
2444                        split_boxes.push(bbox);
2445                        continue;
2446                    }
2447
2448                    for container in &container_boxes {
2449                        let inter_x_min = bbox.x_min().max(container.x_min());
2450                        let inter_y_min = bbox.y_min().max(container.y_min());
2451                        let inter_x_max = bbox.x_max().min(container.x_max());
2452                        let inter_y_max = bbox.y_max().min(container.y_max());
2453
2454                        if inter_x_max - inter_x_min <= 2.0 || inter_y_max - inter_y_min <= 2.0 {
2455                            continue;
2456                        }
2457
2458                        let inter_bbox = crate::processors::BoundingBox::from_coords(
2459                            inter_x_min,
2460                            inter_y_min,
2461                            inter_x_max,
2462                            inter_y_max,
2463                        );
2464                        let inter_area = inter_bbox.area();
2465                        if inter_area <= 0.0 {
2466                            continue;
2467                        }
2468
2469                        let ioa = inter_area / self_area;
2470                        if ioa >= TEXT_BOX_SPLIT_IOA_THRESHOLD {
2471                            intersections.push(inter_bbox);
2472                        }
2473                    }
2474
2475                    if intersections.len() >= 2 {
2476                        split_count += intersections.len();
2477                        split_boxes.extend(intersections);
2478                    } else {
2479                        split_boxes.push(bbox);
2480                    }
2481                }
2482
2483                if split_count > 0 {
2484                    tracing::debug!(
2485                        "Cross-layout re-recognition: split {} text boxes into {} sub-boxes",
2486                        split_count,
2487                        split_boxes.len()
2488                    );
2489                }
2490
2491                detection_boxes = split_boxes;
2492            }
2493        }
2494
2495        // PaddleX sorts OCR detection boxes in reading order before cropping/recognition.
2496        if !detection_boxes.is_empty() {
2497            detection_boxes = oar_ocr_core::processors::sort_quad_boxes(&detection_boxes);
2498        }
2499
2500        // Debug: boxes actually used for recognition cropping (after cross-layout splitting).
2501        if tracing::enabled!(tracing::Level::DEBUG) && !detection_boxes.is_empty() {
2502            let pre_rec_rects: Vec<[f32; 4]> = detection_boxes
2503                .iter()
2504                .map(|b| [b.x_min(), b.y_min(), b.x_max(), b.y_max()])
2505                .collect();
2506            tracing::debug!(
2507                "overall OCR boxes pre-recognition (after splitting): {:?}",
2508                pre_rec_rects
2509            );
2510        }
2511
2512        if !detection_boxes.is_empty() {
2513            use crate::oarocr::processors::{EdgeProcessor, TextCroppingProcessor};
2514
2515            let processor = TextCroppingProcessor::new(true);
2516            let cropped =
2517                processor.process((Arc::new(page_image.clone()), detection_boxes.clone()))?;
2518
2519            let mut cropped_images: Vec<image::RgbImage> = Vec::new();
2520            let mut valid_indices: Vec<usize> = Vec::new();
2521
2522            for (idx, crop_result) in cropped.into_iter().enumerate() {
2523                if let Some(img) = crop_result {
2524                    cropped_images.push((*img).clone());
2525                    valid_indices.push(idx);
2526                }
2527            }
2528
2529            if !cropped_images.is_empty() {
2530                // PaddleX applies textline orientation in detection order first.
2531                if let Some(ref tlo_adapter) = self.pipeline.text_line_orientation_adapter {
2532                    let tlo_input = ImageTaskInput::new(cropped_images.clone());
2533                    match tlo_adapter.execute(tlo_input, None) {
2534                        Ok(tlo_result) => {
2535                            for (i, classifications) in
2536                                tlo_result.classifications.iter().enumerate()
2537                            {
2538                                if i >= cropped_images.len() {
2539                                    break;
2540                                }
2541                                if let Some(top_cls) = classifications.first()
2542                                    && top_cls.class_id == 1
2543                                {
2544                                    cropped_images[i] =
2545                                        image::imageops::rotate180(&cropped_images[i]);
2546                                }
2547                            }
2548                        }
2549                        Err(err) => {
2550                            tracing::warn!(
2551                                "Text-line orientation failed; proceeding without rotation: {}",
2552                                err
2553                            );
2554                        }
2555                    }
2556                }
2557
2558                let mut items: Vec<(usize, f32, image::RgbImage)> = valid_indices
2559                    .into_iter()
2560                    .zip(cropped_images)
2561                    .map(|(det_idx, img)| {
2562                        let wh_ratio = img.width() as f32 / img.height().max(1) as f32;
2563                        (det_idx, wh_ratio, img)
2564                    })
2565                    .collect();
2566
2567                items.sort_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal));
2568
2569                let batch_size = self
2570                    .pipeline
2571                    .region_batch_size
2572                    .unwrap_or_else(|| text_recognition_adapter.recommended_batch_size())
2573                    .max(1);
2574                let mut recognized_by_det_idx: Vec<Option<(String, f32)>> =
2575                    vec![None; detection_boxes.len()];
2576                let mut rec_batches = 0usize;
2577                let t_text_rec = Instant::now();
2578
2579                let mut remaining_items = items.into_iter();
2580                loop {
2581                    let batch_items: Vec<_> = remaining_items.by_ref().take(batch_size).collect();
2582                    if batch_items.is_empty() {
2583                        break;
2584                    }
2585
2586                    let mut det_indices: Vec<usize> = Vec::with_capacity(batch_items.len());
2587                    let mut rec_imgs: Vec<image::RgbImage> = Vec::with_capacity(batch_items.len());
2588                    for (det_idx, _ratio, img) in batch_items {
2589                        det_indices.push(det_idx);
2590                        rec_imgs.push(img);
2591                    }
2592
2593                    let rec_input = ImageTaskInput::new(rec_imgs);
2594                    rec_batches += 1;
2595                    match text_recognition_adapter.execute(rec_input, None) {
2596                        Ok(rec_result) => {
2597                            for ((det_idx, text), score) in det_indices
2598                                .into_iter()
2599                                .zip(rec_result.texts)
2600                                .zip(rec_result.scores)
2601                            {
2602                                if text.is_empty() {
2603                                    continue;
2604                                }
2605                                if let Some(slot) = recognized_by_det_idx.get_mut(det_idx) {
2606                                    *slot = Some((text, score));
2607                                }
2608                            }
2609                        }
2610                        // Mirror the batch path (`precompute_overall_ocr_across_pages`):
2611                        // surface the failure instead of silently dropping the text
2612                        // for these crops.
2613                        Err(err) => {
2614                            tracing::warn!(
2615                                "Text recognition batch failed for {} crops and will be skipped: {}",
2616                                det_indices.len(),
2617                                err
2618                            );
2619                        }
2620                    }
2621                }
2622                tracing::debug!(
2623                    "structure stage: text recognition {:.1} ms, crops={}, batches={}, batch_size={}",
2624                    t_text_rec.elapsed().as_secs_f64() * 1000.0,
2625                    detection_boxes.len(),
2626                    rec_batches,
2627                    batch_size
2628                );
2629
2630                // Emit OCR regions in original detection order, matching PaddleX.
2631                for (det_idx, rec) in recognized_by_det_idx.into_iter().enumerate() {
2632                    let Some((text, score)) = rec else {
2633                        continue;
2634                    };
2635                    let bbox = detection_boxes[det_idx].clone();
2636                    text_regions.push(TextRegion {
2637                        bounding_box: bbox.clone(),
2638                        dt_poly: Some(bbox.clone()),
2639                        rec_poly: Some(bbox),
2640                        text: Some(Arc::from(text)),
2641                        confidence: Some(score),
2642                        orientation_angle: None,
2643                        word_boxes: None,
2644                        label: None,
2645                    });
2646                }
2647            }
2648        }
2649
2650        let batch_size = self
2651            .pipeline
2652            .region_batch_size
2653            .unwrap_or_else(|| text_recognition_adapter.recommended_batch_size())
2654            .max(1);
2655        Self::refine_overall_ocr_with_layout(
2656            &mut text_regions,
2657            layout_elements,
2658            region_blocks,
2659            page_image,
2660            text_recognition_adapter,
2661            batch_size,
2662        )?;
2663        tracing::debug!(
2664            "structure stage: text detection {:.1} ms, boxes={}, recognized_regions={}",
2665            text_det_dur.as_secs_f64() * 1000.0,
2666            detection_boxes.len(),
2667            text_regions.len()
2668        );
2669
2670        Ok(text_regions)
2671    }
2672
2673    /// Analyzes the structure of a document image from a path.
2674    ///
2675    /// # Arguments
2676    ///
2677    /// * `image_path` - Path to the input image
2678    ///
2679    /// # Returns
2680    ///
2681    /// A `StructureResult` containing detected layout elements, tables, formulas, and text.
2682    pub fn predict(&self, image_path: impl Into<PathBuf>) -> Result<StructureResult, OCRError> {
2683        let image_path = image_path.into();
2684
2685        // Load the image
2686        let image = image::open(&image_path).map_err(|e| OCRError::InvalidInput {
2687            message: format!(
2688                "failed to load image from '{}': {}",
2689                image_path.display(),
2690                e
2691            ),
2692        })?;
2693
2694        let mut result = self.predict_image(image.to_rgb8())?;
2695        result.input_path = std::sync::Arc::from(image_path.to_string_lossy().as_ref());
2696        Ok(result)
2697    }
2698
2699    /// Preprocesses a page image before layout detection. Batch callers fill the
2700    /// layout fields later so model inference can run across pages.
2701    fn preprocess_page(&self, image: image::RgbImage) -> Result<PreparedPage, OCRError> {
2702        use crate::oarocr::preprocess::DocumentPreprocessor;
2703        use std::sync::Arc;
2704
2705        let preprocessor = DocumentPreprocessor::new(
2706            self.pipeline.document_orientation_adapter.as_ref(),
2707            self.pipeline.rectification_adapter.as_ref(),
2708        );
2709        let preprocess = preprocessor.preprocess(Arc::new(image))?;
2710        let current_image = preprocess.image;
2711        let orientation_angle = preprocess.orientation_angle;
2712        let rectified_img = preprocess.rectified_img;
2713        let rotation = preprocess.rotation;
2714
2715        Ok(PreparedPage {
2716            current_image,
2717            orientation_angle,
2718            rectified_img,
2719            rotation,
2720            layout_elements: Vec::new(),
2721            detected_region_blocks: None,
2722            precomputed_text_regions: None,
2723        })
2724    }
2725
2726    /// Preprocesses a page image and runs layout detection, returning intermediate
2727    /// results ready for formula recognition and downstream processing.
2728    fn prepare_page(&self, image: image::RgbImage) -> Result<PreparedPage, OCRError> {
2729        let mut prepared = self.preprocess_page(image)?;
2730        let (layout_elements, detected_region_blocks) =
2731            self.detect_layout_and_regions(&prepared.current_image)?;
2732        prepared.layout_elements = layout_elements;
2733        prepared.detected_region_blocks = detected_region_blocks;
2734        Ok(prepared)
2735    }
2736
2737    /// Completes page analysis given a `PreparedPage` and pre-computed formula results.
2738    /// Runs seal detection, OCR, table analysis, stitching, and coordinate transforms.
2739    fn complete_page(
2740        &self,
2741        prepared: PreparedPage,
2742        mut formulas: Vec<crate::domain::structure::FormulaResult>,
2743    ) -> Result<StructureResult, OCRError> {
2744        use std::sync::Arc;
2745
2746        let PreparedPage {
2747            current_image,
2748            orientation_angle,
2749            rectified_img,
2750            rotation,
2751            mut layout_elements,
2752            mut detected_region_blocks,
2753            precomputed_text_regions,
2754        } = prepared;
2755
2756        let mut tables = Vec::new();
2757
2758        self.detect_seal_text(&current_image, &mut layout_elements)?;
2759
2760        // Sort layout elements after all detection/augmentation steps (formulas/seals)
2761        // so reading order includes any injected blocks.
2762        if !layout_elements.is_empty() {
2763            let (width, height) = if let Some(img) = &rectified_img {
2764                (img.width() as f32, img.height() as f32)
2765            } else {
2766                (current_image.width() as f32, current_image.height() as f32)
2767            };
2768            Self::sort_layout_elements_enhanced(&mut layout_elements, width, height);
2769        }
2770
2771        if let Some(ref mut regions) = detected_region_blocks {
2772            Self::assign_region_block_membership(regions, &layout_elements);
2773        }
2774
2775        let t_ocr = Instant::now();
2776        let mut text_regions = if let Some(text_regions) = precomputed_text_regions {
2777            text_regions
2778        } else {
2779            self.run_overall_ocr(
2780                &current_image,
2781                &layout_elements,
2782                detected_region_blocks.as_deref(),
2783            )?
2784        };
2785        let ocr_dur = t_ocr.elapsed();
2786
2787        {
2788            let t_tables = Instant::now();
2789            let analyzer = crate::oarocr::table_analyzer::TableAnalyzer::new(
2790                crate::oarocr::table_analyzer::TableAnalyzerConfig {
2791                    table_classification_adapter: self
2792                        .pipeline
2793                        .table_classification_adapter
2794                        .as_ref(),
2795                    table_orientation_adapter: self.pipeline.table_orientation_adapter.as_ref(),
2796                    table_structure_recognition_adapter: self
2797                        .pipeline
2798                        .table_structure_recognition_adapter
2799                        .as_ref(),
2800                    wired_table_structure_adapter: self
2801                        .pipeline
2802                        .wired_table_structure_adapter
2803                        .as_ref(),
2804                    wireless_table_structure_adapter: self
2805                        .pipeline
2806                        .wireless_table_structure_adapter
2807                        .as_ref(),
2808                    table_cell_detection_adapter: self
2809                        .pipeline
2810                        .table_cell_detection_adapter
2811                        .as_ref(),
2812                    wired_table_cell_adapter: self.pipeline.wired_table_cell_adapter.as_ref(),
2813                    wireless_table_cell_adapter: self.pipeline.wireless_table_cell_adapter.as_ref(),
2814                    use_e2e_wired_table_rec: self.pipeline.use_e2e_wired_table_rec,
2815                    use_e2e_wireless_table_rec: self.pipeline.use_e2e_wireless_table_rec,
2816                    use_wired_table_cells_trans_to_html: self
2817                        .pipeline
2818                        .use_wired_table_cells_trans_to_html,
2819                    use_wireless_table_cells_trans_to_html: self
2820                        .pipeline
2821                        .use_wireless_table_cells_trans_to_html,
2822                },
2823            );
2824            tables.extend(analyzer.analyze_tables(&current_image, &layout_elements)?);
2825            tracing::debug!(
2826                "structure stage: table analysis {:.1} ms, tables={}",
2827                t_tables.elapsed().as_secs_f64() * 1000.0,
2828                tables.len()
2829            );
2830        }
2831        tracing::debug!(
2832            "structure stage: overall OCR total {:.1} ms, regions={}",
2833            ocr_dur.as_secs_f64() * 1000.0,
2834            text_regions.len()
2835        );
2836
2837        // 5b. Optional OCR box splitting by table cell boundaries.
2838        //
2839        // Split OCR boxes that span multiple table cells horizontally and re-recognize
2840        // the smaller segments. This mirrors `split_ocr_bboxes_by_table_cells`:
2841        // - For each OCR box that overlaps >= k table cells, split at cell boundaries
2842        // - Re-run recognition on each split crop
2843        // - Replace the original OCR box with the split boxes + texts
2844        let has_detection_backed_table_cells = tables.iter().any(|table| !table.is_e2e);
2845        if has_detection_backed_table_cells
2846            && !text_regions.is_empty()
2847            && let Some(ref text_rec_adapter) = self.pipeline.text_recognition_adapter
2848        {
2849            Self::split_ocr_bboxes_by_table_cells(
2850                &tables,
2851                &mut text_regions,
2852                &current_image,
2853                text_rec_adapter,
2854            )?;
2855        }
2856
2857        // Transform bounding boxes back to original coordinate system if rotation was applied.
2858        // If rectification was applied, keep coordinates in rectified space (UVDoc can't be inverted).
2859        if let Some(rot) = rotation {
2860            let rotated_width = rot.rotated_width;
2861            let rotated_height = rot.rotated_height;
2862            let angle = rot.angle;
2863
2864            // Transform layout elements
2865            for element in &mut layout_elements {
2866                element.bbox =
2867                    element
2868                        .bbox
2869                        .rotate_back_to_original(angle, rotated_width, rotated_height);
2870            }
2871
2872            // Transform table bounding boxes and cells
2873            for table in &mut tables {
2874                table.bbox =
2875                    table
2876                        .bbox
2877                        .rotate_back_to_original(angle, rotated_width, rotated_height);
2878
2879                // Transform cell bounding boxes
2880                for cell in &mut table.cells {
2881                    cell.bbox =
2882                        cell.bbox
2883                            .rotate_back_to_original(angle, rotated_width, rotated_height);
2884                }
2885            }
2886
2887            // Transform formula bounding boxes
2888            for formula in &mut formulas {
2889                formula.bbox =
2890                    formula
2891                        .bbox
2892                        .rotate_back_to_original(angle, rotated_width, rotated_height);
2893            }
2894
2895            // Transform text region polygons, bounding boxes, and word boxes
2896            for region in &mut text_regions {
2897                region.dt_poly = region
2898                    .dt_poly
2899                    .take()
2900                    .map(|poly| poly.rotate_back_to_original(angle, rotated_width, rotated_height));
2901                region.rec_poly = region
2902                    .rec_poly
2903                    .take()
2904                    .map(|poly| poly.rotate_back_to_original(angle, rotated_width, rotated_height));
2905                region.bounding_box = region.bounding_box.rotate_back_to_original(
2906                    angle,
2907                    rotated_width,
2908                    rotated_height,
2909                );
2910
2911                if let Some(ref word_boxes) = region.word_boxes {
2912                    let transformed_word_boxes: Vec<_> = word_boxes
2913                        .iter()
2914                        .map(|wb| wb.rotate_back_to_original(angle, rotated_width, rotated_height))
2915                        .collect();
2916                    region.word_boxes = Some(transformed_word_boxes);
2917                }
2918            }
2919
2920            // Transform region block bounding boxes
2921            if let Some(ref mut regions) = detected_region_blocks {
2922                for region in regions.iter_mut() {
2923                    region.bbox =
2924                        region
2925                            .bbox
2926                            .rotate_back_to_original(angle, rotated_width, rotated_height);
2927                }
2928            }
2929        }
2930
2931        // PaddleX: convert_formula_res_to_ocr_format — inject formula results into
2932        // the overall OCR pool so they participate in normal block matching and table
2933        // cell matching. The raw LaTeX text is used here (no $...$ wrapping);
2934        // wrapping is handled by to_markdown() for formula elements, by
2935        // stitch_tables() for table cells, and by sort_and_join_texts for inline formulas.
2936        for formula in &formulas {
2937            let w = formula.bbox.x_max() - formula.bbox.x_min();
2938            let h = formula.bbox.y_max() - formula.bbox.y_min();
2939            if w > 1.0 && h > 1.0 {
2940                let mut region = crate::oarocr::TextRegion::new(formula.bbox.clone());
2941                region.text = Some(formula.latex.clone().into());
2942                region.confidence = Some(1.0);
2943                region.label = Some("formula".into()); // Mark as formula for inline wrapping
2944                text_regions.push(region);
2945            }
2946        }
2947
2948        // Construct and return result
2949        // Ensure rectified_img is always set for markdown image extraction
2950        // If no rectification was applied, use current_image
2951        let final_image = rectified_img.unwrap_or_else(|| Arc::new((*current_image).clone()));
2952        let mut result = StructureResult {
2953            input_path: Arc::from("memory"),
2954            index: 0,
2955            layout_elements,
2956            tables,
2957            formulas,
2958            text_regions: if text_regions.is_empty() {
2959                None
2960            } else {
2961                Some(text_regions)
2962            },
2963            orientation_angle,
2964            region_blocks: detected_region_blocks,
2965            rectified_img: Some(final_image),
2966            page_continuation_flags: None,
2967        };
2968
2969        // Stitch text results into layout elements and tables
2970        // Note: When region_blocks is present, stitching preserves the hierarchical order
2971        use crate::oarocr::stitching::{ResultStitcher, StitchConfig};
2972        let stitch_cfg = StitchConfig::default();
2973        ResultStitcher::stitch_with_config(&mut result, &stitch_cfg);
2974
2975        Ok(result)
2976    }
2977
2978    /// Analyzes the structure of a single document image.
2979    pub fn predict_image(&self, image: image::RgbImage) -> Result<StructureResult, OCRError> {
2980        let t_total = Instant::now();
2981        let prepared = self.prepare_page(image)?;
2982        let formulas =
2983            self.recognize_formulas(&prepared.current_image, &prepared.layout_elements)?;
2984        let result = self.complete_page(prepared, formulas)?;
2985        tracing::debug!(
2986            "structure stage: total predict_image {:.1} ms",
2987            t_total.elapsed().as_secs_f64() * 1000.0
2988        );
2989        Ok(result)
2990    }
2991
2992    fn precompute_overall_ocr_across_pages(
2993        &self,
2994        prepared_pages: &mut [Result<PreparedPage, OCRError>],
2995    ) {
2996        use crate::oarocr::TextRegion;
2997        use crate::oarocr::processors::{EdgeProcessor, TextCroppingProcessor};
2998        use oar_ocr_core::core::traits::task::ImageTaskInput;
2999        use std::sync::Arc;
3000
3001        let Some(ref text_detection_adapter) = self.pipeline.text_detection_adapter else {
3002            return;
3003        };
3004        let Some(ref text_recognition_adapter) = self.pipeline.text_recognition_adapter else {
3005            return;
3006        };
3007
3008        // Seal detection augments layout before OCR in the single-page path. Keep
3009        // that path for seal-enabled pipelines until seal detection is batched too.
3010        if self.pipeline.seal_text_detection_adapter.is_some() {
3011            return;
3012        }
3013
3014        let image_batch_size = self
3015            .pipeline
3016            .image_batch_size
3017            .unwrap_or_else(|| text_detection_adapter.recommended_batch_size())
3018            .max(1);
3019
3020        let t_total = Instant::now();
3021
3022        #[derive(Default)]
3023        struct PageOcrState {
3024            detection_boxes: Vec<crate::processors::BoundingBox>,
3025            recognized: Vec<Option<(String, f32)>>,
3026        }
3027
3028        struct RecItem {
3029            page_idx: usize,
3030            det_idx: usize,
3031            wh_ratio: f32,
3032            image: Arc<image::RgbImage>,
3033        }
3034
3035        let mut page_states: Vec<Option<PageOcrState>> =
3036            (0..prepared_pages.len()).map(|_| None).collect();
3037        let mut rec_items: Vec<RecItem> = Vec::new();
3038        let cropper = TextCroppingProcessor::new(true);
3039        let mut batched_detection_boxes: Vec<Option<Vec<crate::processors::BoundingBox>>> =
3040            (0..prepared_pages.len()).map(|_| None).collect();
3041
3042        let t_detection = Instant::now();
3043        let mut det_page_indices = Vec::new();
3044        let mut det_images = Vec::new();
3045        for (page_idx, prepared) in prepared_pages.iter().enumerate() {
3046            let Ok(prepared) = prepared else {
3047                continue;
3048            };
3049
3050            let mut ocr_image = (*prepared.current_image).clone();
3051            if self.pipeline.formula_recognition_adapter.is_some() {
3052                let mask_bboxes: Vec<crate::processors::BoundingBox> = prepared
3053                    .layout_elements
3054                    .iter()
3055                    .filter(|e| e.element_type.is_formula())
3056                    .map(|e| e.bbox.clone())
3057                    .collect();
3058                if !mask_bboxes.is_empty() {
3059                    crate::utils::mask_regions(&mut ocr_image, &mask_bboxes, [255, 255, 255]);
3060                }
3061            }
3062
3063            det_page_indices.push(page_idx);
3064            det_images.push(ocr_image);
3065        }
3066
3067        let mut remaining_pages = det_page_indices.into_iter().zip(det_images);
3068        loop {
3069            let mut batch_page_indices = Vec::with_capacity(image_batch_size);
3070            let mut batch_images = Vec::with_capacity(image_batch_size);
3071            for (page_idx, image) in remaining_pages.by_ref().take(image_batch_size) {
3072                batch_page_indices.push(page_idx);
3073                batch_images.push(image);
3074            }
3075            if batch_page_indices.is_empty() {
3076                break;
3077            }
3078            match text_detection_adapter.execute(ImageTaskInput::new(batch_images), None) {
3079                Ok(det_result) => {
3080                    for (offset, detections) in det_result.detections.into_iter().enumerate() {
3081                        let page_idx = batch_page_indices[offset];
3082                        batched_detection_boxes[page_idx] =
3083                            Some(detections.into_iter().map(|d| d.bbox).collect());
3084                    }
3085                }
3086                Err(err) => {
3087                    tracing::warn!(
3088                        "Batch structure OCR text detection failed; falling back to per-page detection: {}",
3089                        err
3090                    );
3091                }
3092            }
3093        }
3094        let detection_ms = t_detection.elapsed().as_secs_f64() * 1000.0;
3095
3096        let t_crop = Instant::now();
3097        for page_idx in 0..prepared_pages.len() {
3098            let prepared = match &prepared_pages[page_idx] {
3099                Ok(prepared) => prepared,
3100                Err(_) => continue,
3101            };
3102
3103            let mut detection_boxes = if let Some(boxes) = batched_detection_boxes[page_idx].take()
3104            {
3105                boxes
3106            } else {
3107                let mut ocr_image = (*prepared.current_image).clone();
3108                if self.pipeline.formula_recognition_adapter.is_some() {
3109                    let mask_bboxes: Vec<crate::processors::BoundingBox> = prepared
3110                        .layout_elements
3111                        .iter()
3112                        .filter(|e| e.element_type.is_formula())
3113                        .map(|e| e.bbox.clone())
3114                        .collect();
3115                    if !mask_bboxes.is_empty() {
3116                        crate::utils::mask_regions(&mut ocr_image, &mask_bboxes, [255, 255, 255]);
3117                    }
3118                }
3119
3120                let det_result = match text_detection_adapter
3121                    .execute(ImageTaskInput::new(vec![ocr_image]), None)
3122                {
3123                    Ok(result) => result,
3124                    Err(err) => {
3125                        prepared_pages[page_idx] = Err(err);
3126                        continue;
3127                    }
3128                };
3129
3130                det_result
3131                    .detections
3132                    .first()
3133                    .map(|detections| {
3134                        detections
3135                            .iter()
3136                            .map(|d| d.bbox.clone())
3137                            .collect::<Vec<_>>()
3138                    })
3139                    .unwrap_or_default()
3140            };
3141
3142            if !detection_boxes.is_empty() {
3143                let mut split_boxes = Vec::new();
3144                let container_boxes: Vec<crate::processors::BoundingBox> = prepared
3145                    .detected_region_blocks
3146                    .as_ref()
3147                    .map(|regions| regions.iter().map(|r| r.bbox.clone()).collect())
3148                    .unwrap_or_else(|| {
3149                        prepared
3150                            .layout_elements
3151                            .iter()
3152                            .filter(|e| {
3153                                matches!(
3154                                    e.element_type,
3155                                    crate::domain::structure::LayoutElementType::DocTitle
3156                                        | crate::domain::structure::LayoutElementType::ParagraphTitle
3157                                        | crate::domain::structure::LayoutElementType::Text
3158                                        | crate::domain::structure::LayoutElementType::Content
3159                                        | crate::domain::structure::LayoutElementType::Abstract
3160                                        | crate::domain::structure::LayoutElementType::Header
3161                                        | crate::domain::structure::LayoutElementType::Footer
3162                                        | crate::domain::structure::LayoutElementType::Footnote
3163                                        | crate::domain::structure::LayoutElementType::Number
3164                                        | crate::domain::structure::LayoutElementType::Reference
3165                                        | crate::domain::structure::LayoutElementType::ReferenceContent
3166                                        | crate::domain::structure::LayoutElementType::Algorithm
3167                                        | crate::domain::structure::LayoutElementType::AsideText
3168                                        | crate::domain::structure::LayoutElementType::List
3169                                        | crate::domain::structure::LayoutElementType::FigureTitle
3170                                        | crate::domain::structure::LayoutElementType::TableTitle
3171                                        | crate::domain::structure::LayoutElementType::ChartTitle
3172                                        | crate::domain::structure::LayoutElementType::FigureTableChartTitle
3173                                )
3174                            })
3175                            .map(|e| e.bbox.clone())
3176                            .collect()
3177                    });
3178
3179                if !container_boxes.is_empty() {
3180                    for bbox in detection_boxes.into_iter() {
3181                        let mut intersections: Vec<crate::processors::BoundingBox> = Vec::new();
3182                        let self_area = bbox.area();
3183                        if self_area <= 0.0 {
3184                            split_boxes.push(bbox);
3185                            continue;
3186                        }
3187
3188                        for container in &container_boxes {
3189                            let inter_x_min = bbox.x_min().max(container.x_min());
3190                            let inter_y_min = bbox.y_min().max(container.y_min());
3191                            let inter_x_max = bbox.x_max().min(container.x_max());
3192                            let inter_y_max = bbox.y_max().min(container.y_max());
3193
3194                            if inter_x_max - inter_x_min <= 2.0 || inter_y_max - inter_y_min <= 2.0
3195                            {
3196                                continue;
3197                            }
3198
3199                            let inter_bbox = crate::processors::BoundingBox::from_coords(
3200                                inter_x_min,
3201                                inter_y_min,
3202                                inter_x_max,
3203                                inter_y_max,
3204                            );
3205                            let inter_area = inter_bbox.area();
3206                            if inter_area <= 0.0 {
3207                                continue;
3208                            }
3209
3210                            if inter_area / self_area >= TEXT_BOX_SPLIT_IOA_THRESHOLD {
3211                                intersections.push(inter_bbox);
3212                            }
3213                        }
3214
3215                        if intersections.len() >= 2 {
3216                            split_boxes.extend(intersections);
3217                        } else {
3218                            split_boxes.push(bbox);
3219                        }
3220                    }
3221                    detection_boxes = split_boxes;
3222                }
3223            }
3224
3225            if !detection_boxes.is_empty() {
3226                detection_boxes = oar_ocr_core::processors::sort_quad_boxes(&detection_boxes);
3227            }
3228
3229            let state = PageOcrState {
3230                recognized: vec![None; detection_boxes.len()],
3231                detection_boxes,
3232            };
3233
3234            if !state.detection_boxes.is_empty() {
3235                match cropper.process((
3236                    Arc::clone(&prepared.current_image),
3237                    state.detection_boxes.clone(),
3238                )) {
3239                    Ok(cropped) => {
3240                        for (det_idx, crop_result) in cropped.into_iter().enumerate() {
3241                            let Some(img) = crop_result else {
3242                                continue;
3243                            };
3244                            let wh_ratio = img.width() as f32 / img.height().max(1) as f32;
3245                            rec_items.push(RecItem {
3246                                page_idx,
3247                                det_idx,
3248                                wh_ratio,
3249                                image: img,
3250                            });
3251                        }
3252                    }
3253                    Err(err) => {
3254                        prepared_pages[page_idx] = Err(err);
3255                        continue;
3256                    }
3257                }
3258            }
3259
3260            page_states[page_idx] = Some(state);
3261        }
3262        let crop_ms = t_crop.elapsed().as_secs_f64() * 1000.0;
3263
3264        let mut tlo_ms = 0.0;
3265        let mut recognition_ms = 0.0;
3266        if !rec_items.is_empty() {
3267            if let Some(ref tlo_adapter) = self.pipeline.text_line_orientation_adapter {
3268                let t_tlo = Instant::now();
3269                let input = ImageTaskInput::from_arc_images(
3270                    rec_items
3271                        .iter()
3272                        .map(|item| Arc::clone(&item.image))
3273                        .collect(),
3274                );
3275                match tlo_adapter.execute(input, None) {
3276                    Ok(tlo_result) => {
3277                        for (item, classifications) in
3278                            rec_items.iter_mut().zip(tlo_result.classifications)
3279                        {
3280                            if let Some(top_cls) = classifications.first()
3281                                && top_cls.class_id == 1
3282                            {
3283                                item.image = Arc::new(image::imageops::rotate180(&*item.image));
3284                            }
3285                        }
3286                    }
3287                    Err(err) => {
3288                        tracing::warn!(
3289                            "Text-line orientation failed; proceeding without rotation: {}",
3290                            err
3291                        );
3292                    }
3293                }
3294                tlo_ms = t_tlo.elapsed().as_secs_f64() * 1000.0;
3295            }
3296
3297            rec_items.sort_by(|a, b| {
3298                a.wh_ratio
3299                    .partial_cmp(&b.wh_ratio)
3300                    .unwrap_or(std::cmp::Ordering::Equal)
3301            });
3302
3303            let batch_size = self
3304                .pipeline
3305                .region_batch_size
3306                .unwrap_or_else(|| text_recognition_adapter.recommended_batch_size())
3307                .max(1);
3308
3309            let t_recognition = Instant::now();
3310            let mut start = 0usize;
3311            while start < rec_items.len() {
3312                let end = (start + batch_size).min(rec_items.len());
3313                let chunk = &rec_items[start..end];
3314                let rec_input = ImageTaskInput::from_arc_images(
3315                    chunk.iter().map(|item| Arc::clone(&item.image)).collect(),
3316                );
3317                match text_recognition_adapter.execute(rec_input, None) {
3318                    Ok(rec_result) => {
3319                        for (i, item) in chunk.iter().enumerate() {
3320                            let text = rec_result.texts.get(i).cloned().unwrap_or_default();
3321                            if text.is_empty() {
3322                                continue;
3323                            }
3324                            let score = rec_result.scores.get(i).copied().unwrap_or(0.0);
3325                            if let Some(Some(state)) = page_states.get_mut(item.page_idx)
3326                                && let Some(slot) = state.recognized.get_mut(item.det_idx)
3327                            {
3328                                *slot = Some((text, score));
3329                            }
3330                        }
3331                    }
3332                    Err(err) => {
3333                        tracing::warn!(
3334                            "Text recognition batch failed for {} crops and will be skipped: {}",
3335                            end - start,
3336                            err
3337                        );
3338                    }
3339                }
3340                start = end;
3341            }
3342            recognition_ms = t_recognition.elapsed().as_secs_f64() * 1000.0;
3343        }
3344
3345        let batch_size = self
3346            .pipeline
3347            .region_batch_size
3348            .unwrap_or_else(|| text_recognition_adapter.recommended_batch_size())
3349            .max(1);
3350
3351        let t_refine = Instant::now();
3352        let mut precomputed_pages = 0usize;
3353        let mut text_region_count = 0usize;
3354        for page_idx in 0..prepared_pages.len() {
3355            let Some(state) = page_states[page_idx].take() else {
3356                continue;
3357            };
3358            let Ok(prepared) = &mut prepared_pages[page_idx] else {
3359                continue;
3360            };
3361
3362            let mut text_regions = Vec::new();
3363            for (det_idx, rec) in state.recognized.into_iter().enumerate() {
3364                let Some((text, score)) = rec else {
3365                    continue;
3366                };
3367                let bbox = state.detection_boxes[det_idx].clone();
3368                text_regions.push(TextRegion {
3369                    bounding_box: bbox.clone(),
3370                    dt_poly: Some(bbox.clone()),
3371                    rec_poly: Some(bbox),
3372                    text: Some(Arc::from(text)),
3373                    confidence: Some(score),
3374                    orientation_angle: None,
3375                    word_boxes: None,
3376                    label: None,
3377                });
3378            }
3379
3380            if let Err(err) = Self::refine_overall_ocr_with_layout(
3381                &mut text_regions,
3382                &prepared.layout_elements,
3383                prepared.detected_region_blocks.as_deref(),
3384                &prepared.current_image,
3385                text_recognition_adapter,
3386                batch_size,
3387            ) {
3388                prepared_pages[page_idx] = Err(err);
3389                continue;
3390            }
3391
3392            text_region_count += text_regions.len();
3393            prepared.precomputed_text_regions = Some(text_regions);
3394            precomputed_pages += 1;
3395        }
3396        let refine_ms = t_refine.elapsed().as_secs_f64() * 1000.0;
3397
3398        tracing::debug!(
3399            "structure batch OCR: pages={}, regions={}, detection={:.1} ms, crop/split={:.1} ms, tlo={:.1} ms, recognition={:.1} ms, refine={:.1} ms, total={:.1} ms",
3400            precomputed_pages,
3401            text_region_count,
3402            detection_ms,
3403            crop_ms,
3404            tlo_ms,
3405            recognition_ms,
3406            refine_ms,
3407            t_total.elapsed().as_secs_f64() * 1000.0
3408        );
3409    }
3410
3411    /// Analyzes multiple document page images with configured batching.
3412    ///
3413    /// Image-level stages are chunked according to `image_batch_size` when
3414    /// configured, otherwise the layout adapter's recommended batch size is used.
3415    /// OCR recognition crops are aggregated across the full input set and split
3416    /// according to `region_batch_size` when configured.
3417    ///
3418    /// Per-page errors are returned individually so that a failure on one page does
3419    /// not abort the remaining pages.
3420    pub fn predict_images(
3421        &self,
3422        images: Vec<image::RgbImage>,
3423    ) -> Vec<Result<StructureResult, OCRError>> {
3424        use oar_ocr_core::core::traits::task::ImageTaskInput;
3425        use oar_ocr_core::domain::structure::FormulaResult;
3426        use oar_ocr_core::utils::BBoxCrop;
3427
3428        let image_batch_size = self
3429            .pipeline
3430            .image_batch_size
3431            .unwrap_or_else(|| {
3432                self.pipeline
3433                    .layout_detection_adapter
3434                    .recommended_batch_size()
3435            })
3436            .max(1);
3437
3438        if images.is_empty() {
3439            return Vec::new();
3440        }
3441
3442        let t_total = Instant::now();
3443
3444        // Phase 1: Preprocess every page, then run layout/region detection in
3445        // batches. The original single-page path is still used as a fallback if
3446        // a batched layout call fails.
3447        // Pages that fail preparation are recorded as Err and skipped in later phases.
3448        let t_preprocess = Instant::now();
3449        let mut prepared_pages: Vec<Result<PreparedPage, OCRError>> = images
3450            .into_iter()
3451            .map(|image| self.preprocess_page(image))
3452            .collect();
3453        let preprocess_ms = t_preprocess.elapsed().as_secs_f64() * 1000.0;
3454
3455        let batch_pages: Vec<(usize, std::sync::Arc<image::RgbImage>)> = prepared_pages
3456            .iter()
3457            .enumerate()
3458            .filter_map(|(page_idx, prepared)| {
3459                prepared
3460                    .as_ref()
3461                    .ok()
3462                    .map(|page| (page_idx, std::sync::Arc::clone(&page.current_image)))
3463            })
3464            .collect();
3465
3466        let t_layout = Instant::now();
3467        if !batch_pages.is_empty() {
3468            for page_chunk in batch_pages.chunks(image_batch_size) {
3469                let layout_input = ImageTaskInput::from_arc_images(
3470                    page_chunk
3471                        .iter()
3472                        .map(|(_, img)| std::sync::Arc::clone(img))
3473                        .collect(),
3474                );
3475                match self
3476                    .pipeline
3477                    .layout_detection_adapter
3478                    .execute(layout_input, None)
3479                {
3480                    Ok(layout_result) => {
3481                        for (batch_idx, (page_idx, _)) in page_chunk.iter().enumerate() {
3482                            if let Ok(prepared) = &mut prepared_pages[*page_idx] {
3483                                let mut layout_elements = layout_result
3484                                    .elements
3485                                    .get(batch_idx)
3486                                    .map(|elements| Self::layout_elements_from_detection(elements))
3487                                    .unwrap_or_default();
3488                                Self::finish_layout_elements(&mut layout_elements);
3489                                prepared.layout_elements = layout_elements;
3490                            }
3491                        }
3492
3493                        if let Some(ref region_adapter) = self.pipeline.region_detection_adapter {
3494                            let region_input = ImageTaskInput::from_arc_images(
3495                                page_chunk
3496                                    .iter()
3497                                    .map(|(_, img)| std::sync::Arc::clone(img))
3498                                    .collect(),
3499                            );
3500                            match region_adapter.execute(region_input, None) {
3501                                Ok(region_result) => {
3502                                    for (batch_idx, (page_idx, _)) in page_chunk.iter().enumerate()
3503                                    {
3504                                        let Some(region_elements) =
3505                                            region_result.elements.get(batch_idx)
3506                                        else {
3507                                            continue;
3508                                        };
3509                                        if region_elements.is_empty() {
3510                                            continue;
3511                                        }
3512                                        if let Ok(prepared) = &mut prepared_pages[*page_idx] {
3513                                            prepared.detected_region_blocks = Some(
3514                                                region_elements
3515                                                    .iter()
3516                                                    .map(|e| {
3517                                                        crate::domain::structure::RegionBlock {
3518                                                            bbox: e.bbox.clone(),
3519                                                            confidence: e.score,
3520                                                            order_index: None,
3521                                                            element_indices: Vec::new(),
3522                                                        }
3523                                                    })
3524                                                    .collect(),
3525                                            );
3526                                        }
3527                                    }
3528                                }
3529                                Err(err) => {
3530                                    tracing::warn!("Batch region detection failed: {}", err);
3531                                }
3532                            }
3533                        }
3534                    }
3535                    Err(err) => {
3536                        tracing::warn!(
3537                            "Batch layout detection failed; falling back to per-page layout: {}",
3538                            err
3539                        );
3540                        for (page_idx, _) in page_chunk {
3541                            if let Ok(prepared) = &mut prepared_pages[*page_idx] {
3542                                match self.detect_layout_and_regions(&prepared.current_image) {
3543                                    Ok((layout_elements, region_blocks)) => {
3544                                        prepared.layout_elements = layout_elements;
3545                                        prepared.detected_region_blocks = region_blocks;
3546                                    }
3547                                    Err(err) => {
3548                                        prepared_pages[*page_idx] = Err(err);
3549                                    }
3550                                }
3551                            }
3552                        }
3553                    }
3554                }
3555            }
3556        }
3557        let layout_ms = t_layout.elapsed().as_secs_f64() * 1000.0;
3558
3559        // Phase 2: Batch formula recognition across all successfully prepared pages.
3560        let t_formula = Instant::now();
3561        let num_pages = prepared_pages.len();
3562        let mut per_page_formulas: Vec<Vec<FormulaResult>> =
3563            (0..num_pages).map(|_| Vec::new()).collect();
3564
3565        if let Some(ref formula_adapter) = self.pipeline.formula_recognition_adapter {
3566            let mut all_crops: Vec<image::RgbImage> = Vec::new();
3567            let mut crop_meta: Vec<(usize, oar_ocr_core::processors::BoundingBox)> = Vec::new();
3568
3569            for (page_idx, prepared) in prepared_pages.iter().enumerate() {
3570                let prepared = match prepared {
3571                    Ok(p) => p,
3572                    Err(_) => continue,
3573                };
3574                for elem in prepared
3575                    .layout_elements
3576                    .iter()
3577                    .filter(|e| e.element_type.is_formula())
3578                {
3579                    match BBoxCrop::crop_bounding_box(&prepared.current_image, &elem.bbox) {
3580                        Ok(crop) => {
3581                            all_crops.push(crop);
3582                            crop_meta.push((page_idx, elem.bbox.clone()));
3583                        }
3584                        Err(err) => {
3585                            tracing::warn!("Formula region crop failed (batch): {}", err);
3586                        }
3587                    }
3588                }
3589            }
3590
3591            if !all_crops.is_empty() {
3592                let batch_size = formula_adapter.recommended_batch_size().max(1);
3593                let mut remaining_crops = all_crops.into_iter();
3594                let mut meta_offset = 0;
3595
3596                loop {
3597                    let chunk_vec: Vec<_> = remaining_crops.by_ref().take(batch_size).collect();
3598                    let chunk_len = chunk_vec.len();
3599                    if chunk_len == 0 {
3600                        break;
3601                    }
3602
3603                    let chunk_meta = &crop_meta[meta_offset..meta_offset + chunk_len];
3604                    match formula_adapter.execute(ImageTaskInput::new(chunk_vec), None) {
3605                        Ok(formula_output) => {
3606                            for ((page_idx, bbox), (formula_text, score)) in
3607                                chunk_meta.iter().cloned().zip(
3608                                    formula_output
3609                                        .formulas
3610                                        .into_iter()
3611                                        .zip(formula_output.scores),
3612                                )
3613                            {
3614                                let width = bbox.x_max() - bbox.x_min();
3615                                let height = bbox.y_max() - bbox.y_min();
3616                                if width > 0.0 && height > 0.0 {
3617                                    per_page_formulas[page_idx].push(FormulaResult {
3618                                        bbox,
3619                                        latex: formula_text,
3620                                        confidence: score.unwrap_or(0.0),
3621                                    });
3622                                }
3623                            }
3624                        }
3625                        Err(err) => {
3626                            tracing::warn!("Batch formula recognition failed: {}", err);
3627                        }
3628                    }
3629                    meta_offset += chunk_len;
3630                }
3631            }
3632        }
3633        let formula_ms = t_formula.elapsed().as_secs_f64() * 1000.0;
3634
3635        let t_ocr = Instant::now();
3636        self.precompute_overall_ocr_across_pages(&mut prepared_pages);
3637        let ocr_ms = t_ocr.elapsed().as_secs_f64() * 1000.0;
3638
3639        // Phase 3: Complete each page with its pre-computed formula results.
3640        let t_complete = Instant::now();
3641        let results: Vec<_> = prepared_pages
3642            .into_iter()
3643            .zip(per_page_formulas)
3644            .enumerate()
3645            .map(|(index, (prepared, formulas))| {
3646                let mut result = self.complete_page(prepared?, formulas)?;
3647                result.index = index;
3648                Ok(result)
3649            })
3650            .collect();
3651        tracing::debug!(
3652            "structure batch: pages={}, preprocess={:.1} ms, layout/region={:.1} ms, formula={:.1} ms, ocr={:.1} ms, complete={:.1} ms, total={:.1} ms",
3653            num_pages,
3654            preprocess_ms,
3655            layout_ms,
3656            formula_ms,
3657            ocr_ms,
3658            t_complete.elapsed().as_secs_f64() * 1000.0,
3659            t_total.elapsed().as_secs_f64() * 1000.0
3660        );
3661        results
3662    }
3663}
3664
3665#[cfg(test)]
3666mod tests {
3667    use super::*;
3668
3669    #[test]
3670    fn pp_structurev3_preset_sets_the_benchmark_stack() {
3671        let builder = OARStructureBuilder::pp_structurev3();
3672        assert_eq!(
3673            builder.layout_detection_model.as_path(),
3674            Some(std::path::Path::new("pp-doclayoutv3.onnx"))
3675        );
3676        assert_eq!(builder.layout_model_name.as_deref(), Some("PP-DocLayoutV3"));
3677        assert_eq!(
3678            builder
3679                .text_detection_model
3680                .as_ref()
3681                .and_then(|m| m.as_path()),
3682            Some(std::path::Path::new("pp-ocrv6_tiny_det.onnx"))
3683        );
3684        assert_eq!(
3685            builder
3686                .text_recognition_model
3687                .as_ref()
3688                .and_then(|m| m.as_path()),
3689            Some(std::path::Path::new("pp-ocrv6_tiny_rec.onnx"))
3690        );
3691        assert_eq!(
3692            builder.character_dict_path.as_deref(),
3693            Some(std::path::Path::new("ppocrv6_tiny_dict.txt"))
3694        );
3695        assert_eq!(
3696            builder
3697                .table_classification_model
3698                .as_ref()
3699                .and_then(|m| m.as_path()),
3700            Some(std::path::Path::new("pp-lcnet_x1_0_table_cls.onnx"))
3701        );
3702        assert_eq!(
3703            builder
3704                .wired_table_structure_model
3705                .as_ref()
3706                .and_then(|m| m.as_path()),
3707            Some(std::path::Path::new("slanext_wired.onnx"))
3708        );
3709        assert_eq!(
3710            builder
3711                .wireless_table_structure_model
3712                .as_ref()
3713                .and_then(|m| m.as_path()),
3714            Some(std::path::Path::new("slanet_plus.onnx"))
3715        );
3716        assert_eq!(
3717            builder
3718                .wired_table_cell_model
3719                .as_ref()
3720                .and_then(|m| m.as_path()),
3721            Some(std::path::Path::new("rt-detr-l_wired_table_cell_det.onnx"))
3722        );
3723        assert_eq!(
3724            builder.table_structure_dict_path.as_deref(),
3725            Some(std::path::Path::new("table_structure_dict_ch.txt"))
3726        );
3727    }
3728
3729    #[test]
3730    fn layout_presets_include_doclayout_v2_and_v3() {
3731        for (name, expected) in [
3732            ("PP-DocLayoutV2", "pp-doclayoutv2"),
3733            ("pp_doclayout_v2", "pp-doclayoutv2"),
3734            ("PP-DocLayoutV3", "pp-doclayoutv3"),
3735            ("PP-DocLayout-V3", "pp-doclayoutv3"),
3736        ] {
3737            let config = layout_model_config_for(name);
3738            assert_eq!(config.model_name, expected, "{name}");
3739            assert_eq!(config.num_classes, 25, "{name}");
3740        }
3741        assert_eq!(
3742            layout_model_config_for("unknown").model_name,
3743            layout_model_config_for("PP-DocLayout_plus-L").model_name
3744        );
3745    }
3746
3747    #[test]
3748    fn layout_presets_select_their_detection_defaults() {
3749        let json = |config: LayoutDetectionConfig| serde_json::to_value(config).unwrap();
3750        for (preset, expected) in [
3751            (
3752                "PP-DocLayoutV2",
3753                LayoutDetectionConfig::with_pp_doclayoutv2_defaults(),
3754            ),
3755            (
3756                "PP-DocLayoutV3",
3757                LayoutDetectionConfig::with_pp_doclayoutv3_defaults(),
3758            ),
3759            (
3760                "PP-DocLayout_plus-L",
3761                LayoutDetectionConfig::with_pp_structurev3_defaults(),
3762            ),
3763        ] {
3764            let name = layout_model_config_for(preset).model_name;
3765            assert_eq!(
3766                json(default_layout_detection_config(&name)),
3767                json(expected),
3768                "{preset}"
3769            );
3770        }
3771    }
3772
3773    #[test]
3774    fn arena_shrinkage_defaults_on_for_cuda_only() {
3775        use oar_ocr_core::core::config::OrtExecutionProvider;
3776        let cuda = || OrtExecutionProvider::CUDA {
3777            device_id: None,
3778            gpu_mem_limit: None,
3779            arena_extend_strategy: None,
3780            cudnn_conv_algo_search: None,
3781            cudnn_conv_use_max_workspace: None,
3782        };
3783
3784        let mut config = OrtSessionConfig::new().with_execution_providers(vec![cuda()]);
3785        default_arena_shrinkage(&mut config);
3786        assert_eq!(config.arena_shrinkage, Some(true));
3787
3788        let mut explicit = OrtSessionConfig::new()
3789            .with_execution_providers(vec![cuda()])
3790            .with_arena_shrinkage(false);
3791        default_arena_shrinkage(&mut explicit);
3792        assert_eq!(explicit.arena_shrinkage, Some(false));
3793
3794        let mut cpu =
3795            OrtSessionConfig::new().with_execution_providers(vec![OrtExecutionProvider::CPU]);
3796        default_arena_shrinkage(&mut cpu);
3797        assert_eq!(cpu.arena_shrinkage, None);
3798    }
3799
3800    #[test]
3801    fn generic_table_setters_replace_same_type_preset_models() {
3802        let builder = OARStructureBuilder::pp_structurev3()
3803            .with_table_structure_recognition("custom_structure.onnx", "wired")
3804            .with_table_cell_detection("custom_cells.onnx", "wired");
3805        assert!(builder.wired_table_structure_model.is_none());
3806        assert!(builder.wired_table_cell_model.is_none());
3807        assert!(builder.wireless_table_structure_model.is_some());
3808    }
3809
3810    #[test]
3811    fn test_structure_builder_new() {
3812        let builder = OARStructureBuilder::new("layout.onnx");
3813        assert_eq!(
3814            builder.layout_detection_model.as_path(),
3815            Some(std::path::Path::new("layout.onnx"))
3816        );
3817        assert!(builder.table_classification_model.is_none());
3818        assert!(builder.formula_recognition_model.is_none());
3819    }
3820
3821    #[test]
3822    fn test_structure_builder_with_table_components() {
3823        let builder = OARStructureBuilder::new("layout.onnx")
3824            .with_table_classification("table_cls.onnx")
3825            .with_table_cell_detection("table_cell.onnx", TableType::Wired)
3826            .with_table_structure_recognition("table_struct.onnx", TableType::Wired)
3827            .table_structure_dict_path("table_structure_dict.txt");
3828
3829        assert!(builder.table_classification_model.is_some());
3830        assert!(builder.table_cell_detection_model.is_some());
3831        assert!(builder.table_structure_recognition_model.is_some());
3832        assert_eq!(
3833            builder.table_cell_detection_type,
3834            Some(Ok(TableType::Wired))
3835        );
3836        assert_eq!(
3837            builder.table_structure_recognition_type,
3838            Some(Ok(TableType::Wired))
3839        );
3840        assert_eq!(
3841            builder.table_structure_dict_path,
3842            Some(PathBuf::from("table_structure_dict.txt"))
3843        );
3844    }
3845
3846    #[test]
3847    fn test_structure_builder_with_formula() {
3848        let builder = OARStructureBuilder::new("layout.onnx").with_formula_recognition(
3849            "formula.onnx",
3850            "tokenizer.json",
3851            FormulaModelKind::PPFormulaNet,
3852        );
3853
3854        assert!(builder.formula_recognition_model.is_some());
3855        assert!(builder.formula_tokenizer_path.is_some());
3856        assert_eq!(
3857            builder.formula_recognition_type,
3858            Some(Ok(FormulaModelKind::PPFormulaNet))
3859        );
3860    }
3861
3862    #[test]
3863    fn test_structure_builder_accepts_strict_string_kinds() {
3864        let builder = OARStructureBuilder::new("layout.onnx")
3865            .with_table_cell_detection("table_cell.onnx", "wireless")
3866            .with_table_structure_recognition("table_struct.onnx", "wired")
3867            .with_formula_recognition("formula.onnx", "tokenizer.json", "pp_formulanet");
3868
3869        assert_eq!(
3870            builder.table_cell_detection_type,
3871            Some(Ok(TableType::Wireless))
3872        );
3873        assert_eq!(
3874            builder.table_structure_recognition_type,
3875            Some(Ok(TableType::Wired))
3876        );
3877        assert_eq!(
3878            builder.formula_recognition_type,
3879            Some(Ok(FormulaModelKind::PPFormulaNet))
3880        );
3881    }
3882
3883    #[test]
3884    fn test_structure_builder_preserves_invalid_string_kinds() {
3885        let builder = OARStructureBuilder::new("layout.onnx")
3886            .with_table_cell_detection("table_cell.onnx", "border-ish")
3887            .with_formula_recognition("formula.onnx", "tokenizer.json", "mystery");
3888
3889        assert!(matches!(
3890            builder.table_cell_detection_type,
3891            Some(Err(ref message)) if message.contains("border-ish")
3892        ));
3893        assert!(matches!(
3894            builder.formula_recognition_type,
3895            Some(Err(ref message)) if message.contains("mystery")
3896        ));
3897    }
3898
3899    #[test]
3900    fn test_structure_builder_reports_invalid_string_before_loading_models() {
3901        let result = OARStructureBuilder::new("missing-layout.onnx")
3902            .with_table_cell_detection("missing-table.onnx", "border-ish")
3903            .build();
3904        let error = match result {
3905            Ok(_) => panic!("invalid selector must fail"),
3906            Err(error) => error,
3907        };
3908        let message = error.to_string();
3909        assert!(message.contains("table_cell_detection"));
3910        assert!(message.contains("border-ish"));
3911    }
3912
3913    #[test]
3914    fn test_structure_builder_with_ocr() {
3915        let builder =
3916            OARStructureBuilder::new("layout.onnx").with_ocr("det.onnx", "rec.onnx", "dict.txt");
3917
3918        assert!(builder.text_detection_model.is_some());
3919        assert!(builder.text_recognition_model.is_some());
3920        assert!(builder.character_dict_path.is_some());
3921    }
3922
3923    #[test]
3924    fn test_structure_builder_with_configuration() {
3925        let layout_config = LayoutDetectionConfig {
3926            score_threshold: 0.5,
3927            max_elements: 100,
3928            ..Default::default()
3929        };
3930
3931        let builder = OARStructureBuilder::new("layout.onnx")
3932            .layout_detection_config(layout_config.clone())
3933            .image_batch_size(4)
3934            .region_batch_size(64);
3935
3936        assert!(builder.layout_detection_config.is_some());
3937        assert_eq!(builder.image_batch_size, Some(4));
3938        assert_eq!(builder.region_batch_size, Some(64));
3939    }
3940
3941    #[test]
3942    fn test_structure_batch_size_validation() {
3943        assert!(OARStructureBuilder::validate_batch_size("image_batch_size", 1).is_ok());
3944        assert!(
3945            OARStructureBuilder::validate_batch_size(
3946                "region_batch_size",
3947                OARStructureBuilder::MAX_BATCH_SIZE,
3948            )
3949            .is_ok()
3950        );
3951
3952        let err = OARStructureBuilder::validate_batch_size("image_batch_size", 0).unwrap_err();
3953        let msg = err.to_string();
3954        assert!(msg.contains("image_batch_size"));
3955        assert!(msg.contains(&format!("1..={}", OARStructureBuilder::MAX_BATCH_SIZE)));
3956    }
3957}