Browse Source

feat: enhance document processing by adding page range support and integrating SpanMatcher for improved OCR span matching, allowing for more accurate text extraction and processing across specified pages

zhch158_admin 5 days ago
parent
commit
0449712ccc
1 changed files with 85 additions and 20 deletions
  1. 85 20
      zhch/universal_doc_parser/core/pipeline_manager_v2.py

+ 85 - 20
zhch/universal_doc_parser/core/pipeline_manager_v2.py

@@ -36,14 +36,14 @@ try:
     from .model_factory import ModelFactory
     from .config_manager import ConfigManager
     from .coordinate_utils import CoordinateUtils
-    from .layout_utils import LayoutUtils
+    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 layout_utils import LayoutUtils
+    from layout_utils import LayoutUtils, SpanMatcher
     from pdf_utils import PDFUtils
     from element_processors import ElementProcessors
 
@@ -162,12 +162,19 @@ class EnhancedDocPipeline:
     
     # ==================== 主处理流程 ====================
     
-    def process_document(self, document_path: str) -> Dict[str, Any]:
+    def process_document(
+        self, 
+        document_path: str,
+        page_range: Optional[str] = None
+    ) -> Dict[str, Any]:
         """
         处理文档主流程
         
         Args:
             document_path: 文档路径(PDF、图片或目录)
+            page_range: 页面范围字符串,如 "1-5,7,9-12"
+                       - PDF:按页码(从1开始)
+                       - 图片目录:按文件名排序后的位置(从1开始)
             
         Returns:
             处理结果字典
@@ -186,22 +193,27 @@ class EnhancedDocPipeline:
             # 1. 加载文档并分类
             dpi = self.config.get('input', {}).get('dpi', 200)
             images, pdf_type, pdf_doc = PDFUtils.load_and_classify_document(
-                doc_path, dpi=dpi
+                doc_path, dpi=dpi, page_range=page_range
             )
             results['metadata']['pdf_type'] = pdf_type
             results['metadata']['page_count'] = len(images)
+            results['metadata']['page_range'] = page_range
             
             logger.info(f"📄 Loaded {len(images)} pages, type: {pdf_type}")
             
             # 2. 处理每一页
-            for page_idx, image_dict in enumerate(images):
-                logger.info(f"🔍 Processing page {page_idx + 1}/{len(images)}")
+            for idx, image_dict in enumerate(images):
+                # 使用原始页码索引(支持页面范围过滤)
+                page_idx = image_dict.get('page_idx', idx)
+                page_name = image_dict.get('page_name', f'page_{page_idx + 1:03d}')
+                logger.info(f"🔍 Processing page {idx + 1}/{len(images)} (original index: {page_idx})")
                 
                 page_result = self._process_single_page(
                     image_dict=image_dict,
                     page_idx=page_idx,
                     pdf_type=pdf_type,
-                    pdf_doc=pdf_doc
+                    pdf_doc=pdf_doc,
+                    page_name=page_name
                 )
                 results['pages'].append(page_result)
             
@@ -224,7 +236,8 @@ class EnhancedDocPipeline:
         image_dict: Dict[str, Any],
         page_idx: int,
         pdf_type: str,
-        pdf_doc: Optional[Any] = None
+        pdf_doc: Optional[Any] = None,
+        page_name: Optional[str] = None
     ) -> Dict[str, Any]:
         """
         处理单页文档
@@ -240,8 +253,13 @@ class EnhancedDocPipeline:
         scale = image_dict.get('scale', 1.0)
         original_image = np.array(pil_image)
         
+        # 页面名称(用于输出文件命名)
+        if page_name is None:
+            page_name = f"page_{page_idx + 1:03d}"
+        
         page_result = {
             'page_idx': page_idx,
+            'page_name': page_name,  # 用于输出文件命名
             'elements': [],
             'image_shape': original_image.shape,
             'original_image': original_image,
@@ -283,24 +301,41 @@ class EnhancedDocPipeline:
         
         page_result['layout_raw'] = layout_results
         
-        # 3. 分类元素
+        # 3. 整页 OCR 获取所有 text spans(关键改进)
+        all_ocr_spans = []
+        try:
+            all_ocr_spans = self.ocr_recognizer.recognize_text(detection_image)
+            # 去除重复 spans
+            all_ocr_spans = SpanMatcher.remove_duplicate_spans(all_ocr_spans)
+            logger.info(f"📝 Page {page_idx}: OCR detected {len(all_ocr_spans)} text spans")
+        except Exception as e:
+            logger.warning(f"⚠️ Full-page OCR failed: {e}")
+        
+        # 4. 将 OCR spans 匹配到 layout blocks
+        matched_spans = SpanMatcher.match_spans_to_blocks(
+            all_ocr_spans, layout_results, overlap_threshold=0.5
+        )
+        
+        # 5. 分类元素
         classified_elements = self._classify_elements(layout_results, page_idx)
         
-        # 4. 处理各类元素
+        # 6. 处理各类元素(传入匹配的 spans)
         processed_elements, discarded_elements = self._process_all_elements(
             detection_image=detection_image,
             classified_elements=classified_elements,
             pdf_type=pdf_type,
             pdf_doc=pdf_doc,
             page_idx=page_idx,
-            scale=scale
+            scale=scale,
+            matched_spans=matched_spans,
+            layout_results=layout_results
         )
         
-        # 5. 按阅读顺序排序
+        # 7. 按阅读顺序排序
         sorted_elements = LayoutUtils.sort_elements_by_reading_order(processed_elements)
         sorted_discarded = LayoutUtils.sort_elements_by_reading_order(discarded_elements)
         
-        # 6. 坐标转换回原始图片坐标系
+        # 8. 坐标转换回原始图片坐标系
         if rotate_angle != 0:
             sorted_elements = [
                 CoordinateUtils.transform_coords_to_original(
@@ -388,7 +423,9 @@ class EnhancedDocPipeline:
         pdf_type: str,
         pdf_doc: Optional[Any],
         page_idx: int,
-        scale: float
+        scale: float,
+        matched_spans: Optional[Dict[int, List[Dict[str, Any]]]] = None,
+        layout_results: Optional[List[Dict[str, Any]]] = None
     ) -> tuple:
         """
         处理所有分类后的元素
@@ -400,6 +437,8 @@ class EnhancedDocPipeline:
             pdf_doc: PDF文档对象
             page_idx: 页码
             scale: 缩放比例
+            matched_spans: 匹配的 OCR spans {block_idx: [spans]}
+            layout_results: 原始 layout 检测结果(用于索引匹配)
             
         Returns:
             (processed_elements, discarded_elements)
@@ -407,11 +446,29 @@ class EnhancedDocPipeline:
         processed_elements = []
         discarded_elements = []
         
+        # 构建 layout item 到 block index 的映射
+        item_to_block_idx = {}
+        if layout_results:
+            for idx, item in enumerate(layout_results):
+                # 使用 bbox 作为唯一标识(转为元组方便哈希)
+                bbox_key = tuple(item.get('bbox', [0, 0, 0, 0]))
+                item_to_block_idx[bbox_key] = idx
+        
+        def get_matched_spans_for_item(item: Dict[str, Any]) -> List[Dict[str, Any]]:
+            """获取某个 layout item 匹配的 spans"""
+            if not matched_spans or not layout_results:
+                return []
+            bbox_key = tuple[Any, ...](item.get('bbox', [0, 0, 0, 0]))
+            block_idx = item_to_block_idx.get(bbox_key, -1)
+            return matched_spans.get(block_idx, [])
+        
         # 处理纯文本区域
         for item in classified_elements['text']:
             try:
+                spans = get_matched_spans_for_item(item)
                 element = self.element_processors.process_text_element(
-                    detection_image, item, pdf_type, pdf_doc, page_idx, scale
+                    detection_image, item, pdf_type, pdf_doc, page_idx, scale,
+                    pre_matched_spans=spans
                 )
                 processed_elements.append(element)
             except Exception as e:
@@ -421,8 +478,10 @@ class EnhancedDocPipeline:
         # 处理表格标题/脚注
         for item in classified_elements['table_text']:
             try:
+                spans = get_matched_spans_for_item(item)
                 element = self.element_processors.process_text_element(
-                    detection_image, item, pdf_type, pdf_doc, page_idx, scale
+                    detection_image, item, pdf_type, pdf_doc, page_idx, scale,
+                    pre_matched_spans=spans
                 )
                 processed_elements.append(element)
             except Exception as e:
@@ -432,8 +491,10 @@ class EnhancedDocPipeline:
         # 处理图片标题/脚注
         for item in classified_elements['image_text']:
             try:
+                spans = get_matched_spans_for_item(item)
                 element = self.element_processors.process_text_element(
-                    detection_image, item, pdf_type, pdf_doc, page_idx, scale
+                    detection_image, item, pdf_type, pdf_doc, page_idx, scale,
+                    pre_matched_spans=spans
                 )
                 processed_elements.append(element)
             except Exception as e:
@@ -443,8 +504,9 @@ class EnhancedDocPipeline:
         # 处理表格主体
         for item in classified_elements['table_body']:
             try:
+                spans = get_matched_spans_for_item(item)
                 element = self.element_processors.process_table_element(
-                    detection_image, item, scale
+                    detection_image, item, scale, pre_matched_spans=spans
                 )
                 processed_elements.append(element)
             except Exception as e:
@@ -476,8 +538,10 @@ class EnhancedDocPipeline:
         # 处理代码元素
         for item in classified_elements['code']:
             try:
+                spans = get_matched_spans_for_item(item)
                 element = self.element_processors.process_code_element(
-                    detection_image, item, pdf_type, pdf_doc, page_idx, scale
+                    detection_image, item, pdf_type, pdf_doc, page_idx, scale,
+                    pre_matched_spans=spans
                 )
                 processed_elements.append(element)
             except Exception as e:
@@ -487,8 +551,9 @@ class EnhancedDocPipeline:
         # 处理丢弃元素
         for item in classified_elements['discard']:
             try:
+                spans = get_matched_spans_for_item(item)
                 element = self.element_processors.process_discard_element(
-                    detection_image, item
+                    detection_image, item, pre_matched_spans=spans
                 )
                 discarded_elements.append(element)
             except Exception as e: