Skip to main content

oar_ocr/oarocr/
structure.rs

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