output_formatter_v2.py 7.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218
  1. """
  2. 统一输出格式化器 v2
  3. 严格遵循 MinerU mineru_vllm_results_cell_bbox 格式
  4. 支持:
  5. 1. MinerU 标准 middle.json 格式(用于 union_make 生成 Markdown)
  6. 2. mineru_vllm_results_cell_bbox 格式(每页独立 JSON)
  7. 3. Markdown 输出(复用 MinerU union_make)
  8. 4. Debug 模式:layout 图片、OCR 图片
  9. 5. 表格 HTML 输出(带坐标信息)
  10. 模块结构:
  11. - json_formatters.py: JSON 格式化工具
  12. - markdown_generator.py: Markdown 生成器
  13. - html_generator.py: HTML 生成器
  14. - visualization_utils.py: 可视化工具
  15. """
  16. import json
  17. from pathlib import Path
  18. from typing import Dict, Any, List, Optional
  19. from loguru import logger
  20. # 导入子模块
  21. from .json_formatters import JSONFormatters
  22. from .markdown_generator import MarkdownGenerator
  23. from .html_generator import HTMLGenerator
  24. from .visualization_utils import VisualizationUtils
  25. class OutputFormatterV2:
  26. """
  27. 统一输出格式化器
  28. 严格遵循 MinerU mineru_vllm_results_cell_bbox 格式:
  29. - middle.json: MinerU 标准格式,用于生成 Markdown
  30. - page_xxx.json: 每页独立的 JSON,包含 table_cells
  31. - Markdown: 带 bbox 注释
  32. - 表格: HTML 格式,带 data-bbox 属性
  33. """
  34. # 颜色映射(导出供其他模块使用)
  35. COLOR_MAP = VisualizationUtils.COLOR_MAP
  36. OCR_BOX_COLOR = VisualizationUtils.OCR_BOX_COLOR
  37. CELL_BOX_COLOR = VisualizationUtils.CELL_BOX_COLOR
  38. def __init__(self, output_dir: str):
  39. """
  40. 初始化格式化器
  41. Args:
  42. output_dir: 输出目录
  43. """
  44. self.output_dir = Path(output_dir)
  45. self.output_dir.mkdir(parents=True, exist_ok=True)
  46. def save_results(
  47. self,
  48. results: Dict[str, Any],
  49. output_config: Dict[str, Any]
  50. ) -> Dict[str, Any]:
  51. """
  52. 保存处理结果
  53. Args:
  54. results: 处理结果
  55. output_config: 输出配置
  56. Returns:
  57. 输出文件路径字典
  58. """
  59. output_paths: Dict[str, Any] = {
  60. 'images': [],
  61. 'json_pages': [],
  62. }
  63. # 创建文档输出目录
  64. doc_name = Path(results['document_path']).stem
  65. doc_output_dir = self.output_dir / doc_name
  66. doc_output_dir.mkdir(parents=True, exist_ok=True)
  67. # 创建 images 子目录
  68. images_dir = doc_output_dir / 'images'
  69. images_dir.mkdir(exist_ok=True)
  70. # 1. 首先保存图片元素(设置 image_path)
  71. image_paths = VisualizationUtils.save_image_elements(results, images_dir, doc_name)
  72. if image_paths:
  73. output_paths['images'] = image_paths
  74. # 2. 转换为 MinerU middle.json 格式
  75. middle_json = JSONFormatters.convert_to_middle_json(results)
  76. # 3. 保存 middle.json
  77. if output_config.get('save_json', True):
  78. json_path = doc_output_dir / f"{doc_name}_middle.json"
  79. with open(json_path, 'w', encoding='utf-8') as f:
  80. json.dump(middle_json, f, ensure_ascii=False, indent=2)
  81. output_paths['middle_json'] = str(json_path)
  82. logger.info(f"📄 Middle JSON saved: {json_path}")
  83. # 4. 保存每页独立的 mineru_vllm_results_cell_bbox 格式 JSON
  84. if output_config.get('save_page_json', True):
  85. page_json_paths = JSONFormatters.save_page_jsons(results, doc_output_dir, doc_name)
  86. output_paths['json_pages'] = page_json_paths
  87. # 5. 保存 Markdown(完整版)
  88. if output_config.get('save_markdown', True):
  89. md_path = MarkdownGenerator.save_markdown(results, middle_json, doc_output_dir, doc_name)
  90. output_paths['markdown'] = str(md_path)
  91. # 5.5 保存每页独立的 Markdown
  92. if output_config.get('save_page_markdown', True):
  93. page_md_paths = MarkdownGenerator.save_page_markdowns(results, doc_output_dir, doc_name)
  94. output_paths['markdown_pages'] = page_md_paths
  95. # 6. 保存表格 HTML
  96. if output_config.get('save_html', True):
  97. html_dir = HTMLGenerator.save_table_htmls(results, doc_output_dir, doc_name)
  98. output_paths['table_htmls'] = str(html_dir)
  99. # 7. Debug 模式:保存可视化图片
  100. if output_config.get('save_layout_image', False):
  101. layout_paths = VisualizationUtils.save_layout_images(
  102. results, doc_output_dir, doc_name,
  103. draw_type_label=output_config.get('draw_type_label', True),
  104. draw_bbox_number=output_config.get('draw_bbox_number', True)
  105. )
  106. output_paths['layout_images'] = layout_paths
  107. if output_config.get('save_ocr_image', False):
  108. ocr_paths = VisualizationUtils.save_ocr_images(results, doc_output_dir, doc_name)
  109. output_paths['ocr_images'] = ocr_paths
  110. logger.info(f"✅ All results saved to: {doc_output_dir}")
  111. return output_paths
  112. # ==================== 便捷函数 ====================
  113. def save_mineru_format(
  114. results: Dict[str, Any],
  115. output_dir: str,
  116. output_config: Optional[Dict[str, Any]] = None
  117. ) -> Dict[str, Any]:
  118. """
  119. 便捷函数:保存为 MinerU 格式
  120. Args:
  121. results: pipeline 处理结果
  122. output_dir: 输出目录
  123. output_config: 输出配置
  124. Returns:
  125. 输出文件路径字典
  126. """
  127. if output_config is None:
  128. output_config = {
  129. 'save_json': True,
  130. 'save_page_json': True,
  131. 'save_markdown': True,
  132. 'save_page_markdown': True,
  133. 'save_html': True,
  134. 'save_layout_image': False,
  135. 'save_ocr_image': False,
  136. }
  137. formatter = OutputFormatterV2(output_dir)
  138. return formatter.save_results(results, output_config)
  139. if __name__ == "__main__":
  140. # 测试代码
  141. sample_results = {
  142. "document_path": "/path/to/sample.pdf",
  143. "scene": "bank_statement",
  144. "pages": [
  145. {
  146. "page_idx": 0,
  147. "image_shape": [1654, 2338, 3],
  148. "elements": [
  149. {
  150. "type": "title",
  151. "bbox": [100, 50, 800, 100],
  152. "content": {"text": "银行流水"},
  153. "confidence": 0.98
  154. },
  155. {
  156. "type": "table",
  157. "bbox": [100, 200, 800, 600],
  158. "content": {
  159. "html": "<table><tr><td>日期</td><td>金额</td></tr></table>",
  160. "cells": [
  161. {"text": "日期", "bbox": [100, 200, 200, 250], "row": 1, "col": 1},
  162. {"text": "金额", "bbox": [200, 200, 300, 250], "row": 1, "col": 2}
  163. ]
  164. }
  165. }
  166. ]
  167. }
  168. ]
  169. }
  170. output_files = save_mineru_format(
  171. sample_results,
  172. "./test_output_v2",
  173. {
  174. "save_json": True,
  175. "save_page_json": True,
  176. "save_markdown": True,
  177. "save_page_markdown": True,
  178. "save_html": True,
  179. "save_layout_image": False,
  180. "save_ocr_image": False
  181. }
  182. )
  183. print("Generated files:", output_files)