json_formatters.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396
  1. """
  2. JSON 格式化工具模块
  3. 提供 JSON 输出格式化功能:
  4. - MinerU middle.json 格式转换
  5. - mineru_vllm_results_cell_bbox 格式转换
  6. - 表格单元格格式化
  7. - 金额数字标准化(全角→半角)
  8. """
  9. import json
  10. import sys
  11. from pathlib import Path
  12. from typing import Dict, Any, List, Optional
  13. from loguru import logger
  14. # 导入数字标准化工具
  15. from .normalize_financial_numbers import normalize_json_table
  16. class JSONFormatters:
  17. """JSON 格式化工具类"""
  18. @staticmethod
  19. def convert_to_middle_json(results: Dict[str, Any]) -> Dict[str, Any]:
  20. """
  21. 转换为 MinerU 标准 middle.json 格式
  22. 用于 vlm_union_make 生成 Markdown
  23. Args:
  24. results: 处理结果
  25. Returns:
  26. MinerU middle.json 格式的字典
  27. """
  28. middle_json = {
  29. "pdf_info": [],
  30. "_backend": "vlm",
  31. "_scene": results.get('scene', 'unknown'),
  32. "_version_name": "2.5.0"
  33. }
  34. for page in results.get('pages', []):
  35. page_info = {
  36. 'page_idx': page['page_idx'],
  37. 'page_size': list(page.get('image_shape', [0, 0])[:2][::-1]),
  38. 'angle': page.get('angle', 0),
  39. 'para_blocks': [],
  40. 'discarded_blocks': []
  41. }
  42. # 处理普通元素
  43. for element in page.get('elements', []):
  44. block = JSONFormatters._element_to_middle_block(element)
  45. if block:
  46. elem_type = element.get('type', '')
  47. if elem_type in ['header', 'footer', 'page_number', 'aside_text', 'abandon', 'discarded']:
  48. page_info['discarded_blocks'].append(block)
  49. else:
  50. page_info['para_blocks'].append(block)
  51. # 处理丢弃元素(从 discarded_blocks 字段)
  52. for element in page.get('discarded_blocks', []):
  53. block = JSONFormatters._element_to_middle_block(element)
  54. if block:
  55. page_info['discarded_blocks'].append(block)
  56. middle_json['pdf_info'].append(page_info)
  57. return middle_json
  58. @staticmethod
  59. def _element_to_middle_block(element: Dict[str, Any]) -> Optional[Dict[str, Any]]:
  60. """
  61. 将元素转换为 MinerU middle.json block 格式
  62. MinerU 期望的嵌套结构:
  63. - image 类型: { type: "image", blocks: [{ type: "image_body", lines: [...] }] }
  64. - table 类型: { type: "table", blocks: [{ type: "table_body", lines: [...] }] }
  65. """
  66. elem_type = element.get('type', '')
  67. bbox = element.get('bbox', [0, 0, 0, 0])
  68. content = element.get('content', {})
  69. block = {
  70. 'type': elem_type,
  71. 'bbox': bbox,
  72. 'angle': element.get('angle', 0),
  73. 'reading_order': element.get('reading_order', 0),
  74. 'lines': []
  75. }
  76. # 文本类型
  77. if elem_type in ['text', 'title', 'ref_text', 'header', 'footer', 'ocr_text']:
  78. text = content.get('text', '') if isinstance(content, dict) else str(content)
  79. if text:
  80. block['lines'] = [{
  81. 'bbox': bbox,
  82. 'spans': [{
  83. 'bbox': bbox,
  84. 'type': 'text',
  85. 'content': text
  86. }]
  87. }]
  88. # 表格类型 - 嵌套结构
  89. elif elem_type in ['table', 'table_body']:
  90. table_html = content.get('html', '')
  91. cells = content.get('cells', [])
  92. block['type'] = 'table'
  93. block['blocks'] = [{
  94. 'type': 'table_body',
  95. 'bbox': bbox,
  96. 'angle': 0,
  97. 'lines': [{
  98. 'bbox': bbox,
  99. 'spans': [{
  100. 'bbox': bbox,
  101. 'type': 'table',
  102. 'html': table_html,
  103. 'cells': cells
  104. }]
  105. }]
  106. }]
  107. # 图片类型 - 嵌套结构
  108. elif elem_type in ['image', 'image_body', 'figure']:
  109. block['type'] = 'image'
  110. block['blocks'] = [{
  111. 'type': 'image_body',
  112. 'bbox': bbox,
  113. 'angle': element.get('angle', 0),
  114. 'lines': [{
  115. 'bbox': bbox,
  116. 'spans': [{
  117. 'bbox': bbox,
  118. 'type': 'image',
  119. 'image_path': content.get('image_path', ''),
  120. 'description': content.get('description', '')
  121. }]
  122. }]
  123. }]
  124. # 公式类型
  125. elif elem_type in ['interline_equation', 'inline_equation', 'equation']:
  126. latex = content.get('latex', '')
  127. block['lines'] = [{
  128. 'bbox': bbox,
  129. 'spans': [{
  130. 'bbox': bbox,
  131. 'type': 'interline_equation' if 'interline' in elem_type else 'inline_equation',
  132. 'content': latex
  133. }]
  134. }]
  135. # 表格/图片附属文本
  136. elif elem_type in ['table_caption', 'table_footnote', 'image_caption', 'image_footnote']:
  137. text = content.get('text', '') if isinstance(content, dict) else str(content)
  138. if text:
  139. block['lines'] = [{
  140. 'bbox': bbox,
  141. 'spans': [{
  142. 'bbox': bbox,
  143. 'type': 'text',
  144. 'content': text
  145. }]
  146. }]
  147. # 丢弃类型
  148. elif elem_type in ['abandon', 'discarded']:
  149. block['type'] = 'abandon'
  150. text = content.get('text', '') if isinstance(content, dict) else str(content)
  151. if text:
  152. block['lines'] = [{
  153. 'bbox': bbox,
  154. 'spans': [{
  155. 'bbox': bbox,
  156. 'type': 'text',
  157. 'content': text
  158. }]
  159. }]
  160. return block
  161. @staticmethod
  162. def save_page_jsons(
  163. results: Dict[str, Any],
  164. output_dir: Path,
  165. doc_name: str,
  166. is_pdf: bool = True,
  167. normalize_numbers: bool = True
  168. ) -> List[str]:
  169. """
  170. 保存每页独立的 JSON(mineru_vllm_results_cell_bbox 格式)
  171. 命名规则:
  172. - PDF输入: 文件名_page_001.json
  173. - 图片输入(单页): 文件名.json
  174. Args:
  175. results: 处理结果
  176. output_dir: 输出目录
  177. doc_name: 文档名称
  178. is_pdf: 是否为 PDF 输入
  179. normalize_numbers: 是否标准化金额数字(全角→半角)
  180. Returns:
  181. 保存的文件路径列表
  182. """
  183. saved_paths = []
  184. total_pages = len(results.get('pages', []))
  185. for page in results.get('pages', []):
  186. page_idx = page.get('page_idx', 0)
  187. # 根据输入类型决定命名
  188. if is_pdf or total_pages > 1:
  189. page_name = f"{doc_name}_page_{page_idx + 1:03d}"
  190. else:
  191. page_name = doc_name
  192. # 转换为 mineru_vllm_results_cell_bbox 格式
  193. page_elements = []
  194. for element in page.get('elements', []):
  195. converted = JSONFormatters._element_to_cell_bbox_format(element, page_idx)
  196. if converted:
  197. page_elements.append(converted)
  198. # 添加丢弃元素
  199. for element in page.get('discarded_blocks', []):
  200. converted = JSONFormatters._element_to_cell_bbox_format(element, page_idx)
  201. if converted:
  202. page_elements.append(converted)
  203. # 转换为 JSON 字符串
  204. json_content = json.dumps(page_elements, ensure_ascii=False, indent=2)
  205. # 金额数字标准化
  206. if normalize_numbers:
  207. original_content = json_content
  208. json_content = normalize_json_table(json_content)
  209. if json_content != original_content:
  210. original_path = output_dir / f"{page_name}_original.json"
  211. with open(original_path, 'w', encoding='utf-8') as f:
  212. f.write(original_content)
  213. logger.debug(f"📄 Original page JSON saved: {original_path}")
  214. # 保存 JSON
  215. json_path = output_dir / f"{page_name}.json"
  216. with open(json_path, 'w', encoding='utf-8') as f:
  217. f.write(json_content)
  218. saved_paths.append(str(json_path))
  219. logger.debug(f"📄 Page JSON saved: {json_path}")
  220. if saved_paths:
  221. logger.info(f"📄 {len(saved_paths)} page JSONs saved")
  222. return saved_paths
  223. @staticmethod
  224. def _element_to_cell_bbox_format(
  225. element: Dict[str, Any],
  226. page_idx: int
  227. ) -> Optional[Dict[str, Any]]:
  228. """
  229. 将元素转换为 mineru_vllm_results_cell_bbox 格式
  230. """
  231. elem_type = element.get('type', '')
  232. bbox = element.get('bbox', [0, 0, 0, 0])
  233. content = element.get('content', {})
  234. # 确保 bbox 是整数列表
  235. bbox = [int(x) for x in bbox[:4]] if bbox else [0, 0, 0, 0]
  236. result = {
  237. 'bbox': bbox,
  238. 'page_idx': page_idx,
  239. 'reading_order': element.get('reading_order', 0)
  240. }
  241. # 文本类型
  242. if elem_type in ['text', 'title', 'ref_text', 'ocr_text']:
  243. text = content.get('text', '') if isinstance(content, dict) else str(content)
  244. result['type'] = 'text' if elem_type != 'title' else 'title'
  245. result['text'] = text
  246. if elem_type == 'title':
  247. result['text_level'] = element.get('level', 1)
  248. # 表格类型
  249. elif elem_type in ['table', 'table_body']:
  250. result['type'] = 'table'
  251. result['img_path'] = content.get('table_image_path', '')
  252. result['table_caption'] = JSONFormatters._ensure_list(content.get('table_caption', []))
  253. result['table_footnote'] = JSONFormatters._ensure_list(content.get('table_footnote', []))
  254. result['table_body'] = content.get('html', '')
  255. # 关键:table_cells 数组
  256. cells = content.get('cells', [])
  257. if cells:
  258. result['table_cells'] = JSONFormatters.format_table_cells(cells)
  259. # 旋转和倾斜信息
  260. if 'table_angle' in content:
  261. result['image_rotation_angle'] = float(content['table_angle'])
  262. if 'skew_angle' in content:
  263. result['skew_angle'] = float(content['skew_angle'])
  264. # 图片类型
  265. elif elem_type in ['image', 'image_body', 'figure']:
  266. result['type'] = 'image'
  267. image_filename = content.get('image_path', '')
  268. result['img_path'] = f"images/{image_filename}" if image_filename else ''
  269. result['image_caption'] = JSONFormatters._ensure_list(content.get('caption', []))
  270. result['image_footnote'] = JSONFormatters._ensure_list(content.get('footnote', []))
  271. # 公式类型
  272. elif elem_type in ['interline_equation', 'inline_equation', 'equation']:
  273. result['type'] = 'equation'
  274. result['text'] = content.get('latex', '') if isinstance(content, dict) else ''
  275. result['text_format'] = 'latex'
  276. # 列表类型
  277. elif elem_type == 'list':
  278. result['type'] = 'list'
  279. result['sub_type'] = 'text'
  280. result['list_items'] = content.get('list_items', []) if isinstance(content, dict) else []
  281. # 页眉页脚
  282. elif elem_type in ['header', 'footer']:
  283. result['type'] = elem_type
  284. result['text'] = content.get('text', '') if isinstance(content, dict) else str(content)
  285. # 表格/图片附属文本
  286. elif elem_type in ['table_caption', 'table_footnote', 'image_caption', 'image_footnote']:
  287. result['type'] = elem_type
  288. result['text'] = content.get('text', '') if isinstance(content, dict) else str(content)
  289. # 丢弃元素
  290. elif elem_type in ['discarded', 'abandon']:
  291. result['type'] = 'discarded'
  292. result['original_category'] = element.get('original_category', 'unknown')
  293. result['text'] = content.get('text', '') if isinstance(content, dict) else ''
  294. else:
  295. return None
  296. return result
  297. @staticmethod
  298. def format_table_cells(cells: List[Dict]) -> List[Dict[str, Any]]:
  299. """
  300. 格式化表格单元格为 mineru_vllm_results_cell_bbox 格式
  301. 输出格式:
  302. {
  303. "type": "table_cell",
  304. "text": "单元格内容",
  305. "matched_text": "OCR匹配文本",
  306. "bbox": [x1, y1, x2, y2],
  307. "row": 1,
  308. "col": 1,
  309. "score": 100.0,
  310. "paddle_bbox_indices": [0, 1]
  311. }
  312. """
  313. formatted_cells = []
  314. for cell in cells:
  315. formatted_cell = {
  316. 'type': 'table_cell',
  317. 'text': cell.get('text', ''),
  318. 'matched_text': cell.get('matched_text', cell.get('text', '')),
  319. 'bbox': [float(x) for x in cell.get('bbox', [0, 0, 0, 0])[:4]],
  320. 'row': cell.get('row', 0),
  321. 'col': cell.get('col', 0),
  322. 'score': float(cell.get('score', 100.0)),
  323. 'paddle_bbox_indices': cell.get('paddle_bbox_indices',
  324. cell.get('paddle_indices', []))
  325. }
  326. formatted_cells.append(formatted_cell)
  327. return formatted_cells
  328. @staticmethod
  329. def _ensure_list(value) -> List:
  330. """确保值是列表"""
  331. if value is None:
  332. return []
  333. if isinstance(value, str):
  334. return [value] if value else []
  335. if isinstance(value, list):
  336. return value
  337. return [str(value)]