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