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