Procházet zdrojové kódy

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 před 6 hodinami
rodič
revize
e857ac89a5

+ 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
 import numpy as np
 from loguru import logger
+import traceback
 
 # 添加 ocr_platform 根目录到 Python 路径(用于导入 ocr_utils)
 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:
     from .model_factory import ModelFactory
     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 .pdf_utils import PDFUtils
     from .element_processors import ElementProcessors
 except ImportError:
     from model_factory import ModelFactory
     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 pdf_utils import PDFUtils
     from element_processors import ElementProcessors
 
 # 从 ocr_tools.ocr_merger 导入 merger 组件
@@ -126,6 +128,13 @@ class EnhancedDocPipeline:
                 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识别器(表格、公式)
             if self.config.get('vl_recognition', {}) != {}:
                 self.vl_recognizer = ModelFactory.create_vl_recognizer(
@@ -138,6 +147,11 @@ class EnhancedDocPipeline:
             self.ocr_recognizer = ModelFactory.create_ocr_recognizer(
                 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. 表格分类器(可选)
             self.table_classifier = None
@@ -357,40 +371,65 @@ class EnhancedDocPipeline:
         
         # 2. Layout检测
         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:
             logger.error(f"❌ Layout detection failed: {e}")
             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(关键改进)
         all_text_spans = []
@@ -485,12 +524,12 @@ class EnhancedDocPipeline:
         # 8. 坐标转换回原始图片坐标系
         if rotate_angle != 0:
             sorted_elements = [
-                CoordinateUtils.transform_coords_to_original(
+                TableCoordinateUtils.transform_coords_to_original(
                     element, rotate_angle, detection_image.shape, original_image.shape
                 ) for element in sorted_elements
             ]
             sorted_discarded = [
-                CoordinateUtils.transform_coords_to_original(
+                TableCoordinateUtils.transform_coords_to_original(
                     element, rotate_angle, detection_image.shape, original_image.shape
                 ) for element in sorted_discarded
             ]
@@ -872,8 +911,8 @@ class EnhancedDocPipeline:
                 
                 processed_elements.append(element)
             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']: