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