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