Browse Source

feat: add MinerU output formatter for standardized document processing, supporting JSON, Markdown, and table formats with image saving capabilities

zhch158_admin 3 days ago
parent
commit
8ef1d66999
1 changed files with 632 additions and 0 deletions
  1. 632 0
      zhch/universal_doc_parser/utils/mineru_output_formatter.py

+ 632 - 0
zhch/universal_doc_parser/utils/mineru_output_formatter.py

@@ -0,0 +1,632 @@
+"""
+MinerU 标准输出格式化器
+
+严格遵循 MinerU 的输出格式:
+1. JSON: 与 mineru_vllm_results_cell_bbox 格式完全一致
+2. Markdown: 带 bbox 注释的 Markdown
+3. 表格单元格坐标: table_cells 数组
+
+封装 MinerU 的 vlm_middle_json_mkcontent 和 pipeline_middle_json_mkcontent 实现
+"""
+import json
+import os
+import sys
+from pathlib import Path
+from typing import Dict, Any, List, Optional, Tuple
+from loguru import logger
+import numpy as np
+from PIL import Image, ImageDraw
+import cv2
+import hashlib
+
+# 添加 MinerU 路径
+mineru_path = Path(__file__).parents[3]
+if str(mineru_path) not in sys.path:
+    sys.path.insert(0, str(mineru_path))
+
+# 导入 MinerU 组件
+try:
+    from mineru.utils.enum_class import MakeMode, BlockType, ContentType
+    MINERU_AVAILABLE = True
+except ImportError as e:
+    logger.warning(f"MinerU components not available: {e}")
+    MINERU_AVAILABLE = False
+    # 定义占位符
+    class MakeMode:
+        MM_MD = 'mm_md'
+        NLP_MD = 'nlp_md'
+        CONTENT_LIST = 'content_list'
+    
+    class BlockType:
+        TEXT = 'text'
+        TABLE = 'table'
+        IMAGE = 'image'
+        TITLE = 'title'
+        LIST = 'list'
+        HEADER = 'header'
+        FOOTER = 'footer'
+        TABLE_BODY = 'table_body'
+        TABLE_CAPTION = 'table_caption'
+        TABLE_FOOTNOTE = 'table_footnote'
+    
+    class ContentType:
+        TEXT = 'text'
+        TABLE = 'table'
+        IMAGE = 'image'
+
+
+class MinerUOutputFormatter:
+    """
+    MinerU 标准输出格式化器
+    
+    输出格式严格遵循 MinerU 的 mineru_vllm_results_cell_bbox 格式:
+    - JSON: 每页一个 JSON 文件,包含元素列表
+    - Markdown: 每页一个 MD 文件,带 bbox 注释
+    - 表格: HTML 格式,带 data-bbox 属性
+    - 单元格坐标: table_cells 数组
+    """
+    
+    def __init__(self, output_dir: str, save_images: bool = True):
+        """
+        初始化格式化器
+        
+        Args:
+            output_dir: 输出目录
+            save_images: 是否保存图片
+        """
+        self.output_dir = Path(output_dir)
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        self.save_images = save_images
+        
+        # 创建 images 子目录
+        self.images_dir = self.output_dir / "images"
+        self.images_dir.mkdir(parents=True, exist_ok=True)
+    
+    def format_and_save(
+        self,
+        results: Dict[str, Any],
+        doc_name: str = None
+    ) -> Dict[str, List[str]]:
+        """
+        格式化并保存结果
+        
+        Args:
+            results: pipeline 处理结果
+            doc_name: 文档名称(可选)
+            
+        Returns:
+            输出文件路径字典
+        """
+        if doc_name is None:
+            doc_name = Path(results.get('document_path', 'unknown')).stem
+        
+        output_paths = {
+            'json': [],
+            'markdown': [],
+            'images': []
+        }
+        
+        for page in results.get('pages', []):
+            page_idx = page.get('page_idx', 0)
+            page_name = f"{doc_name}_page_{page_idx + 1:03d}"
+            
+            # 1. 转换为 MinerU 格式
+            mineru_elements = self._convert_page_to_mineru_format(
+                page, 
+                results.get('metadata', {})
+            )
+            
+            # 2. 保存 JSON
+            json_path = self.output_dir / f"{page_name}.json"
+            with open(json_path, 'w', encoding='utf-8') as f:
+                json.dump(mineru_elements, f, ensure_ascii=False, indent=2)
+            output_paths['json'].append(str(json_path))
+            
+            # 3. 生成并保存 Markdown
+            md_content = self._generate_markdown(mineru_elements)
+            md_path = self.output_dir / f"{page_name}.md"
+            with open(md_path, 'w', encoding='utf-8') as f:
+                f.write(md_content)
+            output_paths['markdown'].append(str(md_path))
+            
+            # 4. 保存原始图片
+            if self.save_images and 'original_image' in page:
+                img_path = self._save_page_image(
+                    page['original_image'], 
+                    page_name
+                )
+                if img_path:
+                    output_paths['images'].append(img_path)
+        
+        logger.info(f"✅ MinerU format output saved to: {self.output_dir}")
+        return output_paths
+    
+    def _convert_page_to_mineru_format(
+        self,
+        page: Dict[str, Any],
+        metadata: Dict[str, Any]
+    ) -> List[Dict[str, Any]]:
+        """
+        将页面转换为 MinerU 标准格式
+        
+        Args:
+            page: 页面数据
+            metadata: 文档元数据
+            
+        Returns:
+            MinerU 格式的元素列表
+        """
+        mineru_elements = []
+        page_idx = page.get('page_idx', 0)
+        
+        for element in page.get('elements', []):
+            mineru_elem = self._convert_element_to_mineru(element, page_idx)
+            if mineru_elem:
+                mineru_elements.append(mineru_elem)
+        
+        return mineru_elements
+    
+    def _convert_element_to_mineru(
+        self,
+        element: Dict[str, Any],
+        page_idx: int
+    ) -> Optional[Dict[str, Any]]:
+        """
+        将单个元素转换为 MinerU 格式
+        
+        Args:
+            element: 元素数据
+            page_idx: 页面索引
+            
+        Returns:
+            MinerU 格式的元素
+        """
+        elem_type = element.get('type', '')
+        bbox = element.get('bbox', [0, 0, 0, 0])
+        content = element.get('content', {})
+        
+        # 确保 bbox 是整数列表
+        bbox = [int(x) for x in bbox[:4]] if bbox else [0, 0, 0, 0]
+        
+        result = {
+            'type': self._map_element_type(elem_type),
+            'bbox': bbox,
+            'page_idx': page_idx
+        }
+        
+        # 处理不同类型的元素
+        if elem_type in ['text', 'title', 'ref_text', 'ocr_text']:
+            text = content.get('text', '') if isinstance(content, dict) else str(content)
+            result['text'] = text
+            
+        elif elem_type in ['header', 'footer']:
+            text = content.get('text', '') if isinstance(content, dict) else str(content)
+            result['text'] = text
+            
+        elif elem_type == 'list':
+            result['sub_type'] = 'text'
+            list_items = content.get('list_items', []) if isinstance(content, dict) else []
+            if not list_items and isinstance(content, dict):
+                text = content.get('text', '')
+                if text:
+                    list_items = [text]
+            result['list_items'] = list_items
+            
+        elif elem_type in ['table', 'table_body']:
+            result = self._convert_table_element(element, page_idx)
+            
+        elif elem_type in ['image', 'image_body', 'figure']:
+            result = self._convert_image_element(element, page_idx)
+            
+        elif elem_type in ['interline_equation', 'inline_equation', 'equation']:
+            latex = content.get('latex', '') if isinstance(content, dict) else ''
+            result['text'] = latex
+            result['text_format'] = 'latex'
+        
+        return result
+    
+    def _convert_table_element(
+        self,
+        element: Dict[str, Any],
+        page_idx: int
+    ) -> Dict[str, Any]:
+        """
+        转换表格元素为 MinerU 格式
+        
+        关键:
+        1. table_body: 带 data-bbox 属性的 HTML
+        2. table_cells: 单元格坐标数组
+        3. image_rotation_angle: 旋转角度
+        4. skew_angle: 倾斜角度
+        """
+        bbox = element.get('bbox', [0, 0, 0, 0])
+        bbox = [int(x) for x in bbox[:4]] if bbox else [0, 0, 0, 0]
+        content = element.get('content', {})
+        
+        # 获取表格 HTML
+        table_html = content.get('html', '') or content.get('original_html', '')
+        
+        # 获取表格标题和脚注
+        table_caption = content.get('table_caption', [])
+        if isinstance(table_caption, str):
+            table_caption = [table_caption] if table_caption else []
+        
+        table_footnote = content.get('table_footnote', [])
+        if isinstance(table_footnote, str):
+            table_footnote = [table_footnote] if table_footnote else []
+        
+        # 保存表格图片
+        img_path = ''
+        if self.save_images and 'table_image' in content:
+            img_path = self._save_table_image(
+                content['table_image'],
+                f"table_{page_idx}_{bbox[0]}_{bbox[1]}"
+            )
+        
+        result = {
+            'type': 'table',
+            'img_path': img_path,
+            'table_caption': table_caption,
+            'table_footnote': table_footnote,
+            'table_body': table_html,
+            'bbox': bbox,
+            'page_idx': page_idx
+        }
+        
+        # 添加单元格坐标信息
+        cells = content.get('cells', [])
+        if cells:
+            result['table_cells'] = self._format_table_cells(cells)
+        
+        # 添加旋转和倾斜信息
+        if 'table_angle' in content:
+            result['image_rotation_angle'] = float(content['table_angle'])
+        
+        if 'skew_angle' in content:
+            result['skew_angle'] = float(content['skew_angle'])
+        
+        return result
+    
+    def _format_table_cells(self, cells: List[Dict]) -> List[Dict[str, Any]]:
+        """
+        格式化表格单元格为 MinerU 格式
+        
+        MinerU 格式:
+        {
+            "type": "table_cell",
+            "text": "单元格内容",
+            "matched_text": "OCR匹配文本",
+            "bbox": [y1, x1, y2, x2],  # 注意坐标顺序
+            "row": 2,
+            "col": 1,
+            "score": 100.0,
+            "paddle_bbox_indices": [11, 12]
+        }
+        """
+        formatted_cells = []
+        
+        for cell in cells:
+            formatted_cell = {
+                'type': 'table_cell',
+                'text': cell.get('text', ''),
+                'matched_text': cell.get('matched_text', cell.get('text', '')),
+                'bbox': cell.get('bbox', [0, 0, 0, 0]),
+                'row': cell.get('row', 0),
+                'col': cell.get('col', 0),
+                'score': cell.get('score', 100.0),
+                'paddle_bbox_indices': cell.get('paddle_bbox_indices', 
+                                                cell.get('paddle_indices', []))
+            }
+            
+            # 确保 bbox 是整数列表
+            if formatted_cell['bbox']:
+                formatted_cell['bbox'] = [int(x) for x in formatted_cell['bbox'][:4]]
+            
+            formatted_cells.append(formatted_cell)
+        
+        return formatted_cells
+    
+    def _convert_image_element(
+        self,
+        element: Dict[str, Any],
+        page_idx: int
+    ) -> Dict[str, Any]:
+        """转换图片元素为 MinerU 格式"""
+        bbox = element.get('bbox', [0, 0, 0, 0])
+        bbox = [int(x) for x in bbox[:4]] if bbox else [0, 0, 0, 0]
+        content = element.get('content', {})
+        
+        # 保存图片
+        img_path = ''
+        if self.save_images and 'image_data' in content:
+            img_path = self._save_content_image(
+                content['image_data'],
+                f"image_{page_idx}_{bbox[0]}_{bbox[1]}"
+            )
+        
+        return {
+            'type': 'image',
+            'img_path': img_path,
+            'image_caption': content.get('caption', []),
+            'image_footnote': content.get('footnote', []),
+            'bbox': bbox,
+            'page_idx': page_idx
+        }
+    
+    def _map_element_type(self, elem_type: str) -> str:
+        """映射元素类型到 MinerU 标准类型"""
+        type_mapping = {
+            'text': 'text',
+            'title': 'title',
+            'ref_text': 'ref_text',
+            'ocr_text': 'text',
+            'header': 'header',
+            'footer': 'footer',
+            'table': 'table',
+            'table_body': 'table',
+            'image': 'image',
+            'image_body': 'image',
+            'figure': 'image',
+            'list': 'list',
+            'interline_equation': 'equation',
+            'inline_equation': 'equation',
+            'equation': 'equation',
+        }
+        return type_mapping.get(elem_type, elem_type)
+    
+    def _generate_markdown(self, elements: List[Dict[str, Any]]) -> str:
+        """
+        生成 MinerU 格式的 Markdown
+        
+        格式:
+        <!-- bbox: [x1, y1, x2, y2] -->
+        内容
+        """
+        md_lines = []
+        
+        for elem in elements:
+            elem_type = elem.get('type', '')
+            bbox = elem.get('bbox', [])
+            
+            # 添加 bbox 注释
+            bbox_comment = f"<!-- bbox: {bbox} -->"
+            
+            if elem_type == 'header':
+                md_lines.append(bbox_comment)
+                md_lines.append(f"<!-- 页眉: {elem.get('text', '')} -->")
+                md_lines.append("")
+                
+            elif elem_type == 'footer':
+                md_lines.append(bbox_comment)
+                md_lines.append(f"<!-- 页脚: {elem.get('text', '')} -->")
+                md_lines.append("")
+                
+            elif elem_type == 'title':
+                md_lines.append(bbox_comment)
+                level = elem.get('text_level', 1)
+                md_lines.append(f"{'#' * min(level, 6)} {elem.get('text', '')}")
+                md_lines.append("")
+                
+            elif elem_type == 'text':
+                md_lines.append(bbox_comment)
+                md_lines.append(elem.get('text', ''))
+                md_lines.append("")
+                
+            elif elem_type == 'list':
+                md_lines.append(bbox_comment)
+                for item in elem.get('list_items', []):
+                    md_lines.append(item)
+                md_lines.append("")
+                
+            elif elem_type == 'table':
+                # 表格标题
+                for caption in elem.get('table_caption', []):
+                    md_lines.append(f"**{caption}**")
+                md_lines.append("")
+                
+                md_lines.append(bbox_comment)
+                
+                # 表格 HTML
+                table_body = elem.get('table_body', '')
+                if table_body:
+                    md_lines.append(table_body)
+                else:
+                    # 如果没有 HTML,显示图片
+                    img_path = elem.get('img_path', '')
+                    if img_path:
+                        md_lines.append(f"![Image]({img_path})")
+                md_lines.append("")
+                
+            elif elem_type == 'image':
+                md_lines.append(bbox_comment)
+                img_path = elem.get('img_path', '')
+                if img_path:
+                    md_lines.append(f"![Image]({img_path})")
+                md_lines.append("")
+                
+            elif elem_type == 'equation':
+                md_lines.append(bbox_comment)
+                latex = elem.get('text', '')
+                if latex:
+                    md_lines.append(f"$$\n{latex}\n$$")
+                md_lines.append("")
+        
+        return '\n'.join(md_lines)
+    
+    def _save_page_image(
+        self,
+        image: np.ndarray,
+        name: str
+    ) -> Optional[str]:
+        """保存页面图片"""
+        try:
+            # 生成唯一文件名
+            img_hash = hashlib.md5(image.tobytes()[:1000]).hexdigest()[:16]
+            filename = f"{name}_{img_hash}.png"
+            filepath = self.images_dir / filename
+            
+            if isinstance(image, np.ndarray):
+                if len(image.shape) == 3 and image.shape[2] == 3:
+                    # BGR to RGB
+                    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+                Image.fromarray(image).save(filepath)
+            
+            return f"images/{filename}"
+        except Exception as e:
+            logger.warning(f"Failed to save page image: {e}")
+            return None
+    
+    def _save_table_image(
+        self,
+        image: np.ndarray,
+        name: str
+    ) -> str:
+        """保存表格图片"""
+        try:
+            img_hash = hashlib.md5(image.tobytes()[:1000]).hexdigest()[:32]
+            filename = f"{img_hash}.jpg"
+            filepath = self.images_dir / filename
+            
+            if isinstance(image, np.ndarray):
+                if len(image.shape) == 3 and image.shape[2] == 3:
+                    image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+                Image.fromarray(image).save(filepath, quality=95)
+            
+            return f"images/{filename}"
+        except Exception as e:
+            logger.warning(f"Failed to save table image: {e}")
+            return ""
+    
+    def _save_content_image(
+        self,
+        image: np.ndarray,
+        name: str
+    ) -> str:
+        """保存内容图片"""
+        return self._save_table_image(image, name)
+
+
+class MinerUFormatConverter:
+    """
+    MinerU 格式转换器
+    
+    用于将 pipeline 的内部格式转换为 MinerU 标准输出格式
+    """
+    
+    @staticmethod
+    def convert_pipeline_result_to_mineru(
+        results: Dict[str, Any]
+    ) -> List[Dict[str, Any]]:
+        """
+        将 pipeline 结果转换为 MinerU content_list 格式
+        
+        Args:
+            results: pipeline 处理结果
+            
+        Returns:
+            MinerU content_list 格式的结果
+        """
+        content_list = []
+        
+        for page in results.get('pages', []):
+            page_idx = page.get('page_idx', 0)
+            page_size = page.get('image_shape', [0, 0])[:2]
+            
+            for element in page.get('elements', []):
+                content = MinerUFormatConverter._element_to_content(
+                    element, page_idx, page_size
+                )
+                if content:
+                    content_list.append(content)
+        
+        return content_list
+    
+    @staticmethod
+    def _element_to_content(
+        element: Dict[str, Any],
+        page_idx: int,
+        page_size: Tuple[int, int]
+    ) -> Optional[Dict[str, Any]]:
+        """将元素转换为 MinerU content 格式"""
+        elem_type = element.get('type', '')
+        bbox = element.get('bbox', [0, 0, 0, 0])
+        content = element.get('content', {})
+        
+        # 转换 bbox 为相对坐标(千分比)
+        page_height, page_width = page_size if len(page_size) >= 2 else (1, 1)
+        if page_width > 0 and page_height > 0 and bbox:
+            relative_bbox = [
+                int(bbox[0] * 1000 / page_width),
+                int(bbox[1] * 1000 / page_height),
+                int(bbox[2] * 1000 / page_width),
+                int(bbox[3] * 1000 / page_height)
+            ]
+        else:
+            relative_bbox = bbox
+        
+        result = {
+            'bbox': relative_bbox,
+            'page_idx': page_idx
+        }
+        
+        if elem_type in ['text', 'title', 'ref_text', 'header', 'footer']:
+            text = content.get('text', '') if isinstance(content, dict) else str(content)
+            result['type'] = elem_type
+            result['text'] = text
+            
+        elif elem_type in ['table', 'table_body']:
+            result['type'] = 'table'
+            result['img_path'] = content.get('img_path', '')
+            result['table_caption'] = content.get('table_caption', [])
+            result['table_footnote'] = content.get('table_footnote', [])
+            result['table_body'] = content.get('html', '')
+            
+            # 添加单元格信息
+            cells = content.get('cells', [])
+            if cells:
+                result['table_cells'] = cells
+            
+            # 添加旋转信息
+            if 'table_angle' in content:
+                result['image_rotation_angle'] = content['table_angle']
+            if 'skew_angle' in content:
+                result['skew_angle'] = content['skew_angle']
+                
+        elif elem_type in ['image', 'figure']:
+            result['type'] = 'image'
+            result['img_path'] = content.get('img_path', '')
+            result['image_caption'] = content.get('caption', [])
+            result['image_footnote'] = content.get('footnote', [])
+            
+        elif elem_type == 'list':
+            result['type'] = 'list'
+            result['sub_type'] = 'text'
+            result['list_items'] = content.get('list_items', [])
+            
+        else:
+            return None
+        
+        return result
+
+
+def save_mineru_format(
+    results: Dict[str, Any],
+    output_dir: str,
+    doc_name: str = None,
+    save_images: bool = True
+) -> Dict[str, List[str]]:
+    """
+    便捷函数:保存为 MinerU 格式
+    
+    Args:
+        results: pipeline 处理结果
+        output_dir: 输出目录
+        doc_name: 文档名称
+        save_images: 是否保存图片
+        
+    Returns:
+        输出文件路径
+    """
+    formatter = MinerUOutputFormatter(output_dir, save_images)
+    return formatter.format_and_save(results, doc_name)
+