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