Browse Source

feat: enhance EnhancedDocPipeline with detailed element categorization, improved processing methods for text, tables, images, equations, and code, and added support for handling discarded elements and OCR box transformations

zhch158_admin 2 days ago
parent
commit
5e7ec5c568
1 changed files with 418 additions and 56 deletions
  1. 418 56
      zhch/universal_doc_parser/core/pipeline_manager_v2.py

+ 418 - 56
zhch/universal_doc_parser/core/pipeline_manager_v2.py

@@ -62,6 +62,39 @@ except ImportError as e:
 class EnhancedDocPipeline:
     """增强版文档处理流水线"""
     
+    # 元素类别分组(参考 MinerU CategoryId 定义)
+    # CategoryId: 0=Title, 1=Text, 2=Abandon, 3=ImageBody, 4=ImageCaption, 
+    #             5=TableBody, 6=TableCaption, 7=TableFootnote, 8=InterlineEquation_Layout,
+    #             9=InterlineEquationNumber_Layout, 13=InlineEquation, 14=InterlineEquation_YOLO,
+    #             15=OcrText, 16=LowScoreText, 101=ImageFootnote
+    
+    # 文本类元素:需要OCR识别或PDF文本提取
+    TEXT_CATEGORIES = [
+        'text', 'title', 'ocr_text', 'low_score_text',  # 基础文本
+        'header', 'footer', 'page_number',               # 页面元素
+        'ref_text', 'aside_text', 'page_footnote',       # 辅助文本
+    ]
+    
+    # 表格相关元素
+    TABLE_BODY_CATEGORIES = ['table', 'table_body']      # 表格主体:需要OCR+VLM处理
+    TABLE_TEXT_CATEGORIES = ['table_caption', 'table_footnote']  # 表格附属文本
+    
+    # 图片相关元素
+    IMAGE_BODY_CATEGORIES = ['image', 'image_body', 'figure']    # 图片主体
+    IMAGE_TEXT_CATEGORIES = ['image_caption', 'image_footnote']  # 图片附属文本
+    
+    # 公式类元素
+    EQUATION_CATEGORIES = [
+        'interline_equation', 'inline_equation', 'equation',
+        'interline_equation_yolo', 'interline_equation_number'
+    ]
+    
+    # 丢弃类元素(水印、装饰等)
+    DISCARD_CATEGORIES = ['abandon', 'discarded']
+    
+    # 代码类元素
+    CODE_CATEGORIES = ['code', 'code_body', 'code_caption', 'algorithm']
+    
     def __init__(self, config_path: str):
         """
         初始化流水线
@@ -311,71 +344,154 @@ class EnhancedDocPipeline:
         
         page_result['layout_raw'] = layout_results
         
-        # 3. 分类元素
-        text_elements = []
-        table_elements = []
-        other_elements = []
+        # 3. 分类元素(参考 MinerU 的分类策略)
+        text_elements = []           # 纯文本元素
+        table_body_elements = []     # 表格主体
+        table_text_elements = []     # 表格标题/脚注
+        image_body_elements = []     # 图片主体
+        image_text_elements = []     # 图片标题/脚注
+        equation_elements = []       # 公式
+        code_elements = []           # 代码
+        discard_elements = []        # 丢弃元素
         
         for item in layout_results:
-            category = item.get('category', '')
-            if category in ['text', 'title', 'ocr_text', 'ref_text', 'header', 'footer', 
-                            'image_caption', 'image_footnote', 'table_caption', 'table_footnote']:
+            category = item.get('category', '').lower()
+            
+            if category in self.TEXT_CATEGORIES:
                 text_elements.append(item)
-            elif category in ['table', 'table_body']:
-                table_elements.append(item)
+            elif category in self.TABLE_BODY_CATEGORIES:
+                table_body_elements.append(item)
+            elif category in self.TABLE_TEXT_CATEGORIES:
+                table_text_elements.append(item)
+            elif category in self.IMAGE_BODY_CATEGORIES:
+                image_body_elements.append(item)
+            elif category in self.IMAGE_TEXT_CATEGORIES:
+                image_text_elements.append(item)
+            elif category in self.EQUATION_CATEGORIES:
+                equation_elements.append(item)
+            elif category in self.CODE_CATEGORIES:
+                code_elements.append(item)
+            elif category in self.DISCARD_CATEGORIES:
+                discard_elements.append(item)
             else:
-                other_elements.append(item)
+                # 未知类别归入文本处理
+                logger.debug(f"Unknown category '{category}', treating as text")
+                text_elements.append(item)
         
-        # 4. 并行处理文本和表格(在旋转后的图片上进行识别)
+        logger.info(f"📊 Page {page_idx} elements: "
+                   f"text={len(text_elements)}, "
+                   f"table={len(table_body_elements)}, "
+                   f"table_text={len(table_text_elements)}, "
+                   f"image={len(image_body_elements)}, "
+                   f"image_text={len(image_text_elements)}, "
+                   f"equation={len(equation_elements)}, "
+                   f"code={len(code_elements)}, "
+                   f"discard={len(discard_elements)}")
+        
+        # 4. 并行处理各类元素(在旋转后的图片上进行识别)
         processed_elements = []
+        discarded_elements_result = []  # 单独存储丢弃元素
+        
+        # 辅助函数:处理元素并转换坐标
+        def process_and_transform(item, processor_func, *args):
+            element = processor_func(detection_image, item, *args)
+            if rotate_angle != 0:
+                element = self._transform_coords_to_original(
+                    element, rotate_angle, detection_image.shape, original_image.shape
+                )
+            return element
         
-        # 4.1 处理文本区域
+        # 4.1 处理文本区域
         for text_item in text_elements:
             try:
-                element = self._process_text_element(
-                    detection_image, text_item, pdf_type, pdf_doc, page_idx, scale
+                element = process_and_transform(
+                    text_item, self._process_text_element, pdf_type, pdf_doc, page_idx, scale
                 )
-                # **关键**: 将坐标转换回原始图片坐标系
-                if rotate_angle != 0:
-                    element = self._transform_coords_to_original(
-                        element, rotate_angle, detection_image.shape, original_image.shape
-                    )
                 processed_elements.append(element)
             except Exception as e:
                 logger.warning(f"⚠️ Text processing failed: {e}")
                 processed_elements.append(self._create_error_element(text_item, str(e)))
         
-        # 4.2 处理表格区域(OCR检测 + VLM结构识别 + 坐标匹配
-        for table_item in table_elements:
+        # 4.2 处理表格标题/脚注(作为文本处理
+        for table_text_item in table_text_elements:
             try:
-                element = self._process_table_element(
-                    detection_image, table_item, scale
+                element = process_and_transform(
+                    table_text_item, self._process_text_element, pdf_type, pdf_doc, page_idx, scale
+                )
+                processed_elements.append(element)
+            except Exception as e:
+                logger.warning(f"⚠️ Table text processing failed: {e}")
+                processed_elements.append(self._create_error_element(table_text_item, str(e)))
+        
+        # 4.3 处理图片标题/脚注(作为文本处理)
+        for image_text_item in image_text_elements:
+            try:
+                element = process_and_transform(
+                    image_text_item, self._process_text_element, pdf_type, pdf_doc, page_idx, scale
+                )
+                processed_elements.append(element)
+            except Exception as e:
+                logger.warning(f"⚠️ Image text processing failed: {e}")
+                processed_elements.append(self._create_error_element(image_text_item, str(e)))
+        
+        # 4.4 处理表格主体(OCR检测 + VLM结构识别 + 坐标匹配)
+        for table_item in table_body_elements:
+            try:
+                element = process_and_transform(
+                    table_item, self._process_table_element, scale
                 )
-                # **关键**: 将坐标转换回原始图片坐标系
-                if rotate_angle != 0:
-                    element = self._transform_coords_to_original(
-                        element, rotate_angle, detection_image.shape, original_image.shape
-                    )
                 processed_elements.append(element)
             except Exception as e:
                 logger.warning(f"⚠️ Table processing failed: {e}")
                 processed_elements.append(self._create_error_element(table_item, str(e)))
         
-        # 4.3 处理其他元素
-        for other_item in other_elements:
+        # 4.5 处理公式元素
+        for equation_item in equation_elements:
             try:
-                element = self._process_other_element(detection_image, other_item)
-                # **关键**: 将坐标转换回原始图片坐标系
+                element = process_and_transform(
+                    equation_item, self._process_equation_element
+                )
+                processed_elements.append(element)
+            except Exception as e:
+                logger.warning(f"⚠️ Equation processing failed: {e}")
+                processed_elements.append(self._create_error_element(equation_item, str(e)))
+        
+        # 4.6 处理图片主体
+        for image_item in image_body_elements:
+            try:
+                element = process_and_transform(
+                    image_item, self._process_image_element
+                )
+                processed_elements.append(element)
+            except Exception as e:
+                logger.warning(f"⚠️ Image processing failed: {e}")
+                processed_elements.append(self._create_error_element(image_item, str(e)))
+        
+        # 4.7 处理代码元素(作为文本处理)
+        for code_item in code_elements:
+            try:
+                element = process_and_transform(
+                    code_item, self._process_code_element, pdf_type, pdf_doc, page_idx, scale
+                )
+                processed_elements.append(element)
+            except Exception as e:
+                logger.warning(f"⚠️ Code processing failed: {e}")
+                processed_elements.append(self._create_error_element(code_item, str(e)))
+        
+        # 4.8 处理丢弃元素(仅记录,不做OCR)
+        for discard_item in discard_elements:
+            try:
+                element = self._process_discard_element(discard_item)
                 if rotate_angle != 0:
                     element = self._transform_coords_to_original(
                         element, rotate_angle, detection_image.shape, original_image.shape
                     )
-                processed_elements.append(element)
+                discarded_elements_result.append(element)
             except Exception as e:
-                logger.warning(f"⚠️ Element processing failed: {e}")
-                processed_elements.append(self._create_error_element(other_item, str(e)))
+                logger.debug(f"Discard element processing failed: {e}")
         
         page_result['elements'] = processed_elements
+        page_result['discarded_blocks'] = discarded_elements_result  # MinerU 兼容字段
         return page_result
     
     def _process_text_element(
@@ -541,11 +657,20 @@ class EnhancedDocPipeline:
                 orig_table_size=orig_table_size,
                 table_bbox=bbox
             )
+            # ocr_boxes 也需要逆向旋转并加上表格偏移量
+            ocr_boxes = self._inverse_rotate_ocr_boxes(
+                ocr_boxes=ocr_boxes,
+                rotation_angle=table_angle,
+                orig_table_size=orig_table_size,
+                table_bbox=bbox
+            )
             logger.info(f"📐 Coordinates transformed back to original image")
         else:
             # 没有旋转,只需要加上表格偏移量转换为整页坐标
             cells = self._add_table_offset_to_cells(cells, bbox)
             enhanced_html = self._add_table_offset_to_html(enhanced_html, bbox)
+            # ocr_boxes 也需要加上表格偏移量
+            ocr_boxes = self._add_table_offset_to_ocr_boxes(ocr_boxes, bbox)
         
         return {
             'type': 'table',
@@ -561,36 +686,26 @@ class EnhancedDocPipeline:
             },
         }
     
-    def _process_other_element(
+    def _process_equation_element(
         self,
         image: np.ndarray,
         layout_item: Dict[str, Any]
     ) -> Dict[str, Any]:
-        """处理其他类型元素(公式、图片等)"""
+        """处理公式元素"""
         bbox = layout_item.get('bbox', [0, 0, 0, 0])
         category = layout_item.get('category', '')
         cropped_region = self._crop_region(image, bbox)
         
-        content = {}
+        content = {'latex': '', 'confidence': 0.0}
         
-        # 公式识别
-        if category in ['interline_equation', 'inline_equation', 'equation']:
-            try:
-                formula_result = self.vl_recognizer.recognize_formula(cropped_region)
-                content = {
-                    'latex': formula_result.get('latex', ''),
-                    'confidence': formula_result.get('confidence', 0.0)
-                }
-            except Exception as e:
-                logger.warning(f"Formula recognition failed: {e}")
-                content = {'latex': '', 'confidence': 0.0}
-        
-        # 图片保持原样
-        elif category in ['image', 'image_body', 'figure']:
+        try:
+            formula_result = self.vl_recognizer.recognize_formula(cropped_region)
             content = {
-                'type': 'image',
-                'description': ''
+                'latex': formula_result.get('latex', ''),
+                'confidence': formula_result.get('confidence', 0.0)
             }
+        except Exception as e:
+            logger.warning(f"Formula recognition failed: {e}")
         
         return {
             'type': category,
@@ -599,6 +714,111 @@ class EnhancedDocPipeline:
             'content': content
         }
     
+    def _process_image_element(
+        self,
+        image: np.ndarray,
+        layout_item: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """
+        处理图片元素
+        
+        裁剪图片区域,返回图片信息(实际保存由 OutputFormatter 处理)
+        """
+        bbox = layout_item.get('bbox', [0, 0, 0, 0])
+        category = layout_item.get('category', '')
+        
+        # 裁剪图片区域
+        cropped_image = self._crop_region(image, bbox)
+        
+        return {
+            'type': category,
+            'bbox': bbox,
+            'confidence': layout_item.get('confidence', 0.0),
+            'content': {
+                'type': 'image',
+                'description': '',
+                'image_data': cropped_image,  # 保存裁剪的图片数据,由 OutputFormatter 保存
+                'image_path': ''  # 由 OutputFormatter 填充实际路径
+            }
+        }
+    
+    def _process_code_element(
+        self,
+        image: np.ndarray,
+        layout_item: Dict[str, Any],
+        pdf_type: str,
+        pdf_doc: Optional[Any],
+        page_idx: int,
+        scale: float
+    ) -> Dict[str, Any]:
+        """
+        处理代码元素
+        
+        代码块作为特殊文本处理,保留格式
+        """
+        bbox = layout_item.get('bbox', [0, 0, 0, 0])
+        category = layout_item.get('category', '')
+        cropped_image = self._crop_region(image, bbox)
+        
+        code_content = ""
+        ocr_details = []
+        
+        # 优先从PDF提取(保留格式更好)
+        if pdf_type == 'txt' and pdf_doc is not None:
+            code_content, success = self._extract_text_from_pdf(
+                pdf_doc, page_idx, bbox, scale
+            )
+            if not success:
+                code_content = ""
+        
+        # PDF提取失败或扫描件,使用OCR
+        if not code_content:
+            try:
+                ocr_results = self.ocr_recognizer.recognize_text(cropped_image)
+                if ocr_results:
+                    # 代码保留换行
+                    lines = []
+                    for item in ocr_results:
+                        lines.append(item.get('text', ''))
+                        ocr_details.append({
+                            'text': item.get('text', ''),
+                            'bbox': item.get('bbox', []),
+                            'confidence': item.get('confidence', 0.0)
+                        })
+                    code_content = '\n'.join(lines)
+            except Exception as e:
+                logger.warning(f"Code OCR failed: {e}")
+        
+        return {
+            'type': category,
+            'bbox': bbox,
+            'confidence': layout_item.get('confidence', 0.0),
+            'content': {
+                'code': code_content,
+                'language': '',  # 可以后续添加语言检测
+                'ocr_details': ocr_details
+            }
+        }
+    
+    def _process_discard_element(
+        self,
+        layout_item: Dict[str, Any]
+    ) -> Dict[str, Any]:
+        """
+        处理丢弃元素(水印、装饰等)
+        
+        仅记录位置,不做OCR识别
+        """
+        bbox = layout_item.get('bbox', [0, 0, 0, 0])
+        category = layout_item.get('category', 'abandon')
+        
+        return {
+            'type': 'discarded',
+            'bbox': bbox,
+            'confidence': layout_item.get('confidence', 0.0),
+            'original_category': category
+        }
+    
     def _extract_text_from_pdf(
         self,
         pdf_doc: Any,
@@ -913,6 +1133,129 @@ class EnhancedDocPipeline:
         
         return converted_html
     
+    def _add_table_offset_to_ocr_boxes(
+        self,
+        ocr_boxes: List[Dict],
+        table_bbox: List[float]
+    ) -> List[Dict]:
+        """
+        为 ocr_boxes 添加表格偏移量,将相对坐标转换为页面绝对坐标
+        
+        Args:
+            ocr_boxes: OCR 框列表
+            table_bbox: 表格在页面中的位置 [x1, y1, x2, y2]
+            
+        Returns:
+            转换后的 ocr_boxes
+        """
+        if not ocr_boxes or not table_bbox:
+            return ocr_boxes
+        
+        offset_x = table_bbox[0]
+        offset_y = table_bbox[1]
+        
+        converted_boxes = []
+        for box in ocr_boxes:
+            new_box = box.copy()
+            
+            # 转换 bbox [x1, y1, x2, y2]
+            if 'bbox' in new_box and new_box['bbox']:
+                bbox = new_box['bbox']
+                if len(bbox) >= 4:
+                    new_box['bbox'] = [
+                        bbox[0] + offset_x,
+                        bbox[1] + offset_y,
+                        bbox[2] + offset_x,
+                        bbox[3] + offset_y
+                    ]
+            
+            # 转换 poly [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
+            if 'poly' in new_box and new_box['poly']:
+                poly = new_box['poly']
+                new_poly = []
+                for point in poly:
+                    if isinstance(point, (list, tuple)) and len(point) >= 2:
+                        new_poly.append([point[0] + offset_x, point[1] + offset_y])
+                    else:
+                        new_poly.append(point)
+                new_box['poly'] = new_poly
+            
+            converted_boxes.append(new_box)
+        
+        return converted_boxes
+    
+    def _inverse_rotate_ocr_boxes(
+        self,
+        ocr_boxes: List[Dict],
+        rotation_angle: float,
+        orig_table_size: Tuple[int, int],
+        table_bbox: List[float]
+    ) -> List[Dict]:
+        """
+        对 ocr_boxes 进行逆向旋转并添加表格偏移量
+        
+        Args:
+            ocr_boxes: OCR 框列表
+            rotation_angle: 表格旋转角度
+            orig_table_size: 原始表格尺寸 (width, height)
+            table_bbox: 表格在页面中的位置
+            
+        Returns:
+            转换后的 ocr_boxes
+        """
+        if not ocr_boxes:
+            return ocr_boxes
+        
+        if not MERGER_AVAILABLE:
+            # 如果 merger 不可用,只添加偏移量
+            return self._add_table_offset_to_ocr_boxes(ocr_boxes, table_bbox)
+        
+        offset_x = table_bbox[0]
+        offset_y = table_bbox[1]
+        
+        converted_boxes = []
+        for box in ocr_boxes:
+            new_box = box.copy()
+            
+            # 逆向旋转 bbox
+            if 'bbox' in new_box and new_box['bbox']:
+                bbox = new_box['bbox']
+                if len(bbox) >= 4:
+                    try:
+                        rotated_bbox = BBoxExtractor.inverse_rotate_box_coordinates(
+                            bbox, rotation_angle, orig_table_size
+                        )
+                        # 添加表格偏移量
+                        new_box['bbox'] = [
+                            rotated_bbox[0] + offset_x,
+                            rotated_bbox[1] + offset_y,
+                            rotated_bbox[2] + offset_x,
+                            rotated_bbox[3] + offset_y
+                        ]
+                    except Exception as e:
+                        logger.debug(f"Failed to inverse rotate ocr_box bbox: {e}")
+            
+            # 逆向旋转 poly
+            if 'poly' in new_box and new_box['poly']:
+                poly = new_box['poly']
+                try:
+                    # 将 poly 转换为 list 格式 [[x1,y1], [x2,y2], ...]
+                    poly_list = [[float(p[0]), float(p[1])] for p in poly]
+                    rotated_poly = BBoxExtractor.inverse_rotate_coordinates(
+                        poly_list, rotation_angle, orig_table_size
+                    )
+                    # 添加表格偏移量
+                    new_poly = []
+                    for point in rotated_poly:
+                        new_poly.append([point[0] + offset_x, point[1] + offset_y])
+                    new_box['poly'] = new_poly
+                except Exception as e:
+                    logger.debug(f"Failed to inverse rotate ocr_box poly: {e}")
+            
+            converted_boxes.append(new_box)
+        
+        return converted_boxes
+    
     def _merge_cross_page_tables(self, results: Dict[str, Any]) -> Dict[str, Any]:
         """合并跨页表格"""
         # TODO: 实现跨页表格合并逻辑
@@ -988,12 +1331,31 @@ class EnhancedDocPipeline:
             if ocr_details:
                 for detail in ocr_details:
                     if 'bbox' in detail and detail['bbox']:
-                        detail['bbox'] = BBoxExtractor.inverse_rotate_box_coordinates(
-                            detail.get('bbox', []), rotate_angle, orig_image_size
+                        if self._is_poly_format(detail['bbox']):
+                            detail['bbox'] = BBoxExtractor.inverse_rotate_coordinates(
+                                detail['bbox'], rotate_angle, orig_image_size
+                            )
+                        else:
+                            detail['bbox'] = BBoxExtractor.inverse_rotate_box_coordinates(
+                                detail.get('bbox', []), rotate_angle, orig_image_size
                         )
         
         return element
     
+    def _is_poly_format(self, bbox: Any) -> bool:
+        """
+        检测 bbox 是否为四点多边形格式
+        
+        四点格式: [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
+        矩形格式: [x_min, y_min, x_max, y_max]
+        """
+        if not bbox or not isinstance(bbox, list):
+            return False
+        if len(bbox) != 4:
+            return False
+        # 如果第一个元素是列表,则为四点格式
+        return isinstance(bbox[0], (list, tuple))
+
     def _transform_html_data_bbox(
         self,
         html: str,