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