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