output_formatter_v2.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332
  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. 6. 金额数字标准化(全角→半角转换)
  11. 模块结构:
  12. - json_formatters.py: JSON 格式化工具
  13. - markdown_generator.py: Markdown 生成器
  14. - html_generator.py: HTML 生成器
  15. - visualization_utils.py: 可视化工具
  16. """
  17. import json
  18. import sys
  19. from pathlib import Path
  20. from typing import Dict, Any, List, Optional
  21. from loguru import logger
  22. # 导入子模块
  23. from .json_formatters import JSONFormatters
  24. from .markdown_generator import MarkdownGenerator
  25. from .html_generator import HTMLGenerator
  26. from .visualization_utils import VisualizationUtils
  27. # 导入数字标准化工具
  28. from .normalize_financial_numbers import normalize_markdown_table, normalize_json_table
  29. class OutputFormatterV2:
  30. """
  31. 统一输出格式化器
  32. 严格遵循 MinerU mineru_vllm_results_cell_bbox 格式:
  33. - middle.json: MinerU 标准格式,用于生成 Markdown
  34. - page_xxx.json: 每页独立的 JSON,包含 table_cells
  35. - Markdown: 带 bbox 注释
  36. - 表格: HTML 格式,带 data-bbox 属性
  37. 命名规则:
  38. - PDF输入: 文件名_page_001.*(按页编号)
  39. - 图片输入: 文件名.*(不加页码后缀)
  40. """
  41. # 颜色映射(导出供其他模块使用)
  42. COLOR_MAP = VisualizationUtils.COLOR_MAP
  43. OCR_BOX_COLOR = VisualizationUtils.OCR_BOX_COLOR
  44. CELL_BOX_COLOR = VisualizationUtils.CELL_BOX_COLOR
  45. def __init__(self, output_dir: str):
  46. """
  47. 初始化格式化器
  48. Args:
  49. output_dir: 输出目录
  50. """
  51. self.output_dir = Path(output_dir)
  52. self.output_dir.mkdir(parents=True, exist_ok=True)
  53. @staticmethod
  54. def is_pdf_input(results: Dict[str, Any]) -> bool:
  55. """
  56. 判断输入是否为 PDF
  57. Args:
  58. results: 处理结果
  59. Returns:
  60. True 如果输入是 PDF,否则 False
  61. """
  62. doc_path = results.get('document_path', '')
  63. if doc_path:
  64. return Path(doc_path).suffix.lower() == '.pdf'
  65. # 如果没有 document_path,检查 metadata
  66. input_type = results.get('metadata', {}).get('input_type', '')
  67. return input_type == 'pdf'
  68. @staticmethod
  69. def get_page_name(doc_name: str, page_idx: int, is_pdf: bool, total_pages: int = 1) -> str:
  70. """
  71. 获取页面名称
  72. Args:
  73. doc_name: 文档名称
  74. page_idx: 页码索引(从0开始)
  75. is_pdf: 是否为 PDF 输入
  76. total_pages: 总页数
  77. Returns:
  78. 页面名称(不含扩展名)
  79. """
  80. if is_pdf or total_pages > 1:
  81. # PDF 或多页输入:添加页码后缀
  82. return f"{doc_name}_page_{page_idx + 1:03d}"
  83. else:
  84. # 单个图片:不添加页码后缀
  85. return doc_name
  86. def save_results(
  87. self,
  88. results: Dict[str, Any],
  89. output_config: Dict[str, Any]
  90. ) -> Dict[str, Any]:
  91. """
  92. 保存处理结果
  93. 命名规则:
  94. - PDF输入: 文件名_page_001.*(按页编号)
  95. - 图片输入: 文件名.*(不加页码后缀)
  96. Args:
  97. results: 处理结果
  98. output_config: 输出配置,支持以下选项:
  99. - create_subdir: 是否在输出目录下创建文档名子目录(默认 False)
  100. - ... 其他选项见 save_mineru_format 函数
  101. Returns:
  102. 输出文件路径字典
  103. """
  104. output_paths: Dict[str, Any] = {
  105. 'images': [],
  106. 'json_pages': [],
  107. }
  108. # 创建文档输出目录
  109. doc_name = Path(results['document_path']).stem
  110. # 是否创建子目录(默认不创建,直接使用指定的输出目录)
  111. create_subdir = output_config.get('create_subdir', False)
  112. if create_subdir:
  113. doc_output_dir = self.output_dir / doc_name
  114. else:
  115. doc_output_dir = self.output_dir
  116. doc_output_dir.mkdir(parents=True, exist_ok=True)
  117. # 判断输入类型
  118. is_pdf = self.is_pdf_input(results)
  119. total_pages = len(results.get('pages', []))
  120. # 创建 images 子目录
  121. images_dir = doc_output_dir / 'images'
  122. images_dir.mkdir(exist_ok=True)
  123. # 1. 首先保存图片元素(设置 image_path)
  124. image_paths = VisualizationUtils.save_image_elements(
  125. results, images_dir, doc_name, is_pdf=is_pdf
  126. )
  127. if image_paths:
  128. output_paths['images'] = image_paths
  129. # 2. 转换为 MinerU middle.json 格式
  130. middle_json = JSONFormatters.convert_to_middle_json(results)
  131. # 3. 保存 middle.json
  132. if output_config.get('save_json', True):
  133. json_path = doc_output_dir / f"{doc_name}_middle.json"
  134. json_content = json.dumps(middle_json, ensure_ascii=False, indent=2)
  135. # 金额数字标准化
  136. normalize_numbers = output_config.get('normalize_numbers', True)
  137. if normalize_numbers:
  138. original_content = json_content
  139. json_content = normalize_json_table(json_content)
  140. # 检查是否有变化
  141. if json_content != original_content:
  142. # 保存原始文件
  143. original_path = doc_output_dir / f"{doc_name}_middle_original.json"
  144. with open(original_path, 'w', encoding='utf-8') as f:
  145. f.write(original_content)
  146. logger.info(f"📄 Original middle JSON saved: {original_path}")
  147. output_paths['middle_json_original'] = str(original_path)
  148. with open(json_path, 'w', encoding='utf-8') as f:
  149. f.write(json_content)
  150. output_paths['middle_json'] = str(json_path)
  151. logger.info(f"📄 Middle JSON saved: {json_path}")
  152. # 4. 保存每页独立的 mineru_vllm_results_cell_bbox 格式 JSON
  153. if output_config.get('save_page_json', True):
  154. normalize_numbers = output_config.get('normalize_numbers', True)
  155. page_json_paths = JSONFormatters.save_page_jsons(
  156. results, doc_output_dir, doc_name, is_pdf=is_pdf,
  157. normalize_numbers=normalize_numbers
  158. )
  159. output_paths['json_pages'] = page_json_paths
  160. # 5. 保存 Markdown(完整版)
  161. if output_config.get('save_markdown', True):
  162. normalize_numbers = output_config.get('normalize_numbers', True)
  163. md_path, original_md_path = MarkdownGenerator.save_markdown(
  164. results, middle_json, doc_output_dir, doc_name,
  165. normalize_numbers=normalize_numbers
  166. )
  167. output_paths['markdown'] = str(md_path)
  168. if original_md_path:
  169. output_paths['markdown_original'] = str(original_md_path)
  170. # 5.5 保存每页独立的 Markdown
  171. if output_config.get('save_page_markdown', True):
  172. normalize_numbers = output_config.get('normalize_numbers', True)
  173. page_md_paths = MarkdownGenerator.save_page_markdowns(
  174. results, doc_output_dir, doc_name, is_pdf=is_pdf,
  175. normalize_numbers=normalize_numbers
  176. )
  177. output_paths['markdown_pages'] = page_md_paths
  178. # 6. 保存表格 HTML
  179. if output_config.get('save_html', True):
  180. html_dir = HTMLGenerator.save_table_htmls(
  181. results, doc_output_dir, doc_name, is_pdf=is_pdf
  182. )
  183. output_paths['table_htmls'] = str(html_dir)
  184. # 7. Debug 模式:保存可视化图片
  185. if output_config.get('save_layout_image', False):
  186. layout_paths = VisualizationUtils.save_layout_images(
  187. results, doc_output_dir, doc_name,
  188. draw_type_label=output_config.get('draw_type_label', True),
  189. draw_bbox_number=output_config.get('draw_bbox_number', True),
  190. is_pdf=is_pdf
  191. )
  192. output_paths['layout_images'] = layout_paths
  193. if output_config.get('save_ocr_image', False):
  194. ocr_paths = VisualizationUtils.save_ocr_images(
  195. results, doc_output_dir, doc_name, is_pdf=is_pdf
  196. )
  197. output_paths['ocr_images'] = ocr_paths
  198. logger.info(f"✅ All results saved to: {doc_output_dir}")
  199. return output_paths
  200. # ==================== 便捷函数 ====================
  201. def save_mineru_format(
  202. results: Dict[str, Any],
  203. output_dir: str,
  204. output_config: Optional[Dict[str, Any]] = None
  205. ) -> Dict[str, Any]:
  206. """
  207. 便捷函数:保存为 MinerU 格式
  208. Args:
  209. results: pipeline 处理结果
  210. output_dir: 输出目录
  211. output_config: 输出配置,支持以下选项:
  212. - create_subdir: 在输出目录下创建文档名子目录(默认 False)
  213. - save_json: 保存 middle.json
  214. - save_page_json: 保存每页 JSON
  215. - save_markdown: 保存完整 Markdown
  216. - save_page_markdown: 保存每页 Markdown
  217. - save_html: 保存表格 HTML
  218. - save_layout_image: 保存布局可视化图
  219. - save_ocr_image: 保存 OCR 可视化图
  220. - normalize_numbers: 标准化金额数字(全角→半角)
  221. Returns:
  222. 输出文件路径字典
  223. """
  224. if output_config is None:
  225. output_config = {
  226. 'create_subdir': False, # 默认不创建子目录,直接使用指定目录
  227. 'save_json': True,
  228. 'save_page_json': True,
  229. 'save_markdown': True,
  230. 'save_page_markdown': True,
  231. 'save_html': True,
  232. 'save_layout_image': False,
  233. 'save_ocr_image': False,
  234. 'normalize_numbers': True, # 默认启用数字标准化
  235. }
  236. formatter = OutputFormatterV2(output_dir)
  237. return formatter.save_results(results, output_config)
  238. if __name__ == "__main__":
  239. # 测试代码
  240. sample_results = {
  241. "document_path": "/path/to/sample.pdf",
  242. "scene": "bank_statement",
  243. "pages": [
  244. {
  245. "page_idx": 0,
  246. "image_shape": [1654, 2338, 3],
  247. "elements": [
  248. {
  249. "type": "title",
  250. "bbox": [100, 50, 800, 100],
  251. "content": {"text": "银行流水"},
  252. "confidence": 0.98
  253. },
  254. {
  255. "type": "table",
  256. "bbox": [100, 200, 800, 600],
  257. "content": {
  258. "html": "<table><tr><td>日期</td><td>金额</td></tr></table>",
  259. "cells": [
  260. {"text": "日期", "bbox": [100, 200, 200, 250], "row": 1, "col": 1},
  261. {"text": "金额", "bbox": [200, 200, 300, 250], "row": 1, "col": 2}
  262. ]
  263. }
  264. }
  265. ]
  266. }
  267. ]
  268. }
  269. output_files = save_mineru_format(
  270. sample_results,
  271. "./test_output_v2",
  272. {
  273. "save_json": True,
  274. "save_page_json": True,
  275. "save_markdown": True,
  276. "save_page_markdown": True,
  277. "save_html": True,
  278. "save_layout_image": False,
  279. "save_ocr_image": False
  280. }
  281. )
  282. print("Generated files:", output_files)