""" 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 格式: 内容 """ md_lines = [] for elem in elements: elem_type = elem.get('type', '') bbox = elem.get('bbox', []) # 添加 bbox 注释 bbox_comment = f"" if elem_type == 'header': md_lines.append(bbox_comment) md_lines.append(f"") md_lines.append("") elif elem_type == 'footer': md_lines.append(bbox_comment) md_lines.append(f"") 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)