| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632 |
- """
- 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"")
- 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"")
- 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)
|