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