Răsfoiți Sursa

feat(pipeline_manager): enhance layout detection with OCR integration

- Updated the EnhancedDocPipeline to support OCR-based layout evaluation, allowing for pre-OCR text box detection to improve layout accuracy.
- Integrated smart router functionality to conditionally set the OCR recognizer based on layout detection strategy.
- Improved error handling in table processing with detailed logging of exceptions and tracebacks.
- Refactored import paths for CoordinateUtils and PDFUtils to maintain consistency across modules.
zhch158_admin 11 ore în urmă
părinte
comite
e857ac89a5
1 a modificat fișierele cu 77 adăugiri și 38 ștergeri
  1. 77 38
      ocr_tools/universal_doc_parser/core/pipeline_manager_v2.py

+ 77 - 38
ocr_tools/universal_doc_parser/core/pipeline_manager_v2.py

@@ -21,6 +21,7 @@ from typing import Dict, List, Any, Optional
 from pathlib import Path
 from pathlib import Path
 import numpy as np
 import numpy as np
 from loguru import logger
 from loguru import logger
+import traceback
 
 
 # 添加 ocr_platform 根目录到 Python 路径(用于导入 ocr_utils)
 # 添加 ocr_platform 根目录到 Python 路径(用于导入 ocr_utils)
 ocr_platform_root = Path(__file__).parents[3]  # core -> universal_doc_parser -> ocr_tools -> ocr_platform -> repository.git
 ocr_platform_root = Path(__file__).parents[3]  # core -> universal_doc_parser -> ocr_tools -> ocr_platform -> repository.git
@@ -36,16 +37,17 @@ if str(module_root) not in sys.path:
 try:
 try:
     from .model_factory import ModelFactory
     from .model_factory import ModelFactory
     from .config_manager import ConfigManager
     from .config_manager import ConfigManager
-    from .coordinate_utils import CoordinateUtils
+    from ocr_utils.coordinate_utils import CoordinateUtils
+    from ocr_utils import PDFUtils
+    from .table_coordinate_utils import TableCoordinateUtils
     from .layout_utils import LayoutUtils, SpanMatcher
     from .layout_utils import LayoutUtils, SpanMatcher
-    from .pdf_utils import PDFUtils
     from .element_processors import ElementProcessors
     from .element_processors import ElementProcessors
 except ImportError:
 except ImportError:
     from model_factory import ModelFactory
     from model_factory import ModelFactory
     from config_manager import ConfigManager
     from config_manager import ConfigManager
-    from coordinate_utils import CoordinateUtils
+    from ocr_utils import CoordinateUtils, PDFUtils
+    from table_coordinate_utils import TableCoordinateUtils
     from layout_utils import LayoutUtils, SpanMatcher
     from layout_utils import LayoutUtils, SpanMatcher
-    from pdf_utils import PDFUtils
     from element_processors import ElementProcessors
     from element_processors import ElementProcessors
 
 
 # 从 ocr_tools.ocr_merger 导入 merger 组件
 # 从 ocr_tools.ocr_merger 导入 merger 组件
@@ -126,6 +128,13 @@ class EnhancedDocPipeline:
                 self.config['layout_detection']
                 self.config['layout_detection']
             )
             )
             
             
+            # 如果是智能路由器且使用ocr_eval策略,需要设置OCR识别器
+            if hasattr(self.layout_detector, 'set_ocr_recognizer'):
+                # OCR识别器会在后面初始化,这里先标记
+                self._smart_router_needs_ocr = True
+            else:
+                self._smart_router_needs_ocr = False
+            
             # 3. VL识别器(表格、公式)
             # 3. VL识别器(表格、公式)
             if self.config.get('vl_recognition', {}) != {}:
             if self.config.get('vl_recognition', {}) != {}:
                 self.vl_recognizer = ModelFactory.create_vl_recognizer(
                 self.vl_recognizer = ModelFactory.create_vl_recognizer(
@@ -138,6 +147,11 @@ class EnhancedDocPipeline:
             self.ocr_recognizer = ModelFactory.create_ocr_recognizer(
             self.ocr_recognizer = ModelFactory.create_ocr_recognizer(
                 self.config['ocr_recognition']
                 self.config['ocr_recognition']
             )
             )
+            
+            # 如果使用智能路由器,设置OCR识别器
+            if self._smart_router_needs_ocr and hasattr(self.layout_detector, 'set_ocr_recognizer'):
+                self.layout_detector.set_ocr_recognizer(self.ocr_recognizer)
+                logger.info("✅ OCR recognizer set for smart router")
 
 
             # 5. 表格分类器(可选)
             # 5. 表格分类器(可选)
             self.table_classifier = None
             self.table_classifier = None
@@ -357,40 +371,65 @@ class EnhancedDocPipeline:
         
         
         # 2. Layout检测
         # 2. Layout检测
         try:
         try:
-            layout_results = self.layout_detector.detect(detection_image)
-            logger.info(f"📋 Page {page_idx}: detected {len(layout_results)} elements (before dedup)")
+            # 如果使用智能路由器且策略是ocr_eval,需要先获取OCR spans(只检测文本框,不识别文字)
+            ocr_spans_for_layout = None
+            if hasattr(self.layout_detector, 'strategy') and self.layout_detector.strategy == 'ocr_eval':
+                # 先运行OCR检测获取文本框(用于layout评估,不需要识别文字)
+                if self.ocr_recognizer:
+                    try:
+                        # 使用 detect_text_boxes 只检测文本框,不识别文字(性能优化)
+                        if hasattr(self.ocr_recognizer, 'detect_text_boxes'):
+                            ocr_spans_for_layout = self.ocr_recognizer.detect_text_boxes(detection_image)
+                        else:
+                            # 如果没有 detect_text_boxes 方法,使用完整识别
+                            ocr_spans_for_layout = self.ocr_recognizer.recognize_text(detection_image)
+                        
+                        ocr_spans_for_layout = SpanMatcher.remove_duplicate_spans(ocr_spans_for_layout)
+                        ocr_spans_for_layout = self._sort_spans_by_position(ocr_spans_for_layout)
+                        logger.info(f"📝 Pre-OCR text box detection for layout evaluation: {len(ocr_spans_for_layout)} boxes")
+                    except Exception as e:
+                        logger.warning(f"⚠️ Pre-OCR text box detection for layout evaluation failed: {e}")
+            
+            # 如果是 SmartLayoutRouter 且启用调试模式,设置调试信息
+            if hasattr(self.layout_detector, 'set_ocr_recognizer'):
+                # 这是 SmartLayoutRouter
+                if self.debug_mode and output_dir:
+                    self.layout_detector.debug_mode = self.debug_mode  # type: ignore
+                    self.layout_detector.output_dir = output_dir  # type: ignore
+            
+            # 调用layout检测(传递OCR spans如果可用)
+            if ocr_spans_for_layout is not None and hasattr(self.layout_detector, 'detect'):
+                # SmartLayoutRouter 的 detect 方法支持 ocr_spans 和 page_name 参数
+                try:
+                    layout_results = self.layout_detector.detect(
+                        detection_image, 
+                        ocr_spans=ocr_spans_for_layout,
+                        page_name=page_name
+                    )
+                except TypeError:
+                    # 如果方法不支持这些参数,使用默认调用
+                    try:
+                        layout_results = self.layout_detector.detect(detection_image, ocr_spans=ocr_spans_for_layout)
+                    except TypeError:
+                        layout_results = self.layout_detector.detect(detection_image)
+            else:
+                # 普通 detector 或没有 ocr_spans
+                if hasattr(self.layout_detector, 'detect') and hasattr(self.layout_detector, 'set_ocr_recognizer'):
+                    # SmartLayoutRouter,尝试传递 page_name
+                    try:
+                        layout_results = self.layout_detector.detect(detection_image, page_name=page_name)
+                    except TypeError:
+                        layout_results = self.layout_detector.detect(detection_image)
+                else:
+                    layout_results = self.layout_detector.detect(detection_image)
+            
+            logger.info(f"📋 Page {page_idx}: detected {len(layout_results)} elements (post-processed)")
         except Exception as e:
         except Exception as e:
             logger.error(f"❌ Layout detection failed: {e}")
             logger.error(f"❌ Layout detection failed: {e}")
             layout_results = []
             layout_results = []
         
         
-        # 2.5 处理重叠框
-        if layout_results:
-            original_count = len(layout_results)
-            layout_results = LayoutUtils.remove_overlapping_boxes(layout_results)
-            removed_count = original_count - len(layout_results)
-            if removed_count > 0:
-                logger.info(f"🔄 Page {page_idx}: removed {removed_count} overlapping boxes")
-        
-        # 2.6 将大面积文本块转换为表格(后处理)
-        if layout_results:
-            convert_large_text = self.config.get('layout', {}).get('convert_large_text_to_table', False)
-            if convert_large_text:
-                # 获取图像尺寸 (height, width)
-                img_shape = detection_image.shape
-                if len(img_shape) >= 2:
-                    image_shape = (int(img_shape[0]), int(img_shape[1]))
-                else:
-                    image_shape = (0, 0)
-                
-                layout_results = LayoutUtils.convert_large_text_to_table(
-                    layout_results,
-                    image_shape,
-                    min_area_ratio=self.config.get('layout', {}).get('min_text_area_ratio', 0.25),
-                    min_width_ratio=self.config.get('layout', {}).get('min_text_width_ratio', 0.4),
-                    min_height_ratio=self.config.get('layout', {}).get('min_text_height_ratio', 0.3)
-                )
-        
-        page_result['layout_raw'] = layout_results
+        # 2.5 后处理布局结果
+        # 注意:所有 detector 的 detect() 方法已经自动执行了后处理
         
         
         # 3. 整页 OCR 获取所有 text spans(关键改进)
         # 3. 整页 OCR 获取所有 text spans(关键改进)
         all_text_spans = []
         all_text_spans = []
@@ -485,12 +524,12 @@ class EnhancedDocPipeline:
         # 8. 坐标转换回原始图片坐标系
         # 8. 坐标转换回原始图片坐标系
         if rotate_angle != 0:
         if rotate_angle != 0:
             sorted_elements = [
             sorted_elements = [
-                CoordinateUtils.transform_coords_to_original(
+                TableCoordinateUtils.transform_coords_to_original(
                     element, rotate_angle, detection_image.shape, original_image.shape
                     element, rotate_angle, detection_image.shape, original_image.shape
                 ) for element in sorted_elements
                 ) for element in sorted_elements
             ]
             ]
             sorted_discarded = [
             sorted_discarded = [
-                CoordinateUtils.transform_coords_to_original(
+                TableCoordinateUtils.transform_coords_to_original(
                     element, rotate_angle, detection_image.shape, original_image.shape
                     element, rotate_angle, detection_image.shape, original_image.shape
                 ) for element in sorted_discarded
                 ) for element in sorted_discarded
             ]
             ]
@@ -872,8 +911,8 @@ class EnhancedDocPipeline:
                 
                 
                 processed_elements.append(element)
                 processed_elements.append(element)
             except Exception as e:
             except Exception as e:
-                logger.warning(f"⚠️ Table processing failed: {e}")
-                processed_elements.append(ElementProcessors.create_error_element(item, str(e)))
+                logger.error(f"⚠️ Table processing failed: {e}, traceback: {traceback.format_exc()}")
+                raise Exception(f"Table processing failed: {e}, traceback: {traceback.format_exc()}")
         
         
         # 处理公式元素
         # 处理公式元素
         for item in classified_elements['equation']:
         for item in classified_elements['equation']: