mineru_output_formatter.py 21 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632
  1. """
  2. MinerU 标准输出格式化器
  3. 严格遵循 MinerU 的输出格式:
  4. 1. JSON: 与 mineru_vllm_results_cell_bbox 格式完全一致
  5. 2. Markdown: 带 bbox 注释的 Markdown
  6. 3. 表格单元格坐标: table_cells 数组
  7. 封装 MinerU 的 vlm_middle_json_mkcontent 和 pipeline_middle_json_mkcontent 实现
  8. """
  9. import json
  10. import os
  11. import sys
  12. from pathlib import Path
  13. from typing import Dict, Any, List, Optional, Tuple
  14. from loguru import logger
  15. import numpy as np
  16. from PIL import Image, ImageDraw
  17. import cv2
  18. import hashlib
  19. # 添加 MinerU 路径
  20. mineru_path = Path(__file__).parents[3]
  21. if str(mineru_path) not in sys.path:
  22. sys.path.insert(0, str(mineru_path))
  23. # 导入 MinerU 组件
  24. try:
  25. from mineru.utils.enum_class import MakeMode, BlockType, ContentType
  26. MINERU_AVAILABLE = True
  27. except ImportError as e:
  28. logger.warning(f"MinerU components not available: {e}")
  29. MINERU_AVAILABLE = False
  30. # 定义占位符
  31. class MakeMode:
  32. MM_MD = 'mm_md'
  33. NLP_MD = 'nlp_md'
  34. CONTENT_LIST = 'content_list'
  35. class BlockType:
  36. TEXT = 'text'
  37. TABLE = 'table'
  38. IMAGE = 'image'
  39. TITLE = 'title'
  40. LIST = 'list'
  41. HEADER = 'header'
  42. FOOTER = 'footer'
  43. TABLE_BODY = 'table_body'
  44. TABLE_CAPTION = 'table_caption'
  45. TABLE_FOOTNOTE = 'table_footnote'
  46. class ContentType:
  47. TEXT = 'text'
  48. TABLE = 'table'
  49. IMAGE = 'image'
  50. class MinerUOutputFormatter:
  51. """
  52. MinerU 标准输出格式化器
  53. 输出格式严格遵循 MinerU 的 mineru_vllm_results_cell_bbox 格式:
  54. - JSON: 每页一个 JSON 文件,包含元素列表
  55. - Markdown: 每页一个 MD 文件,带 bbox 注释
  56. - 表格: HTML 格式,带 data-bbox 属性
  57. - 单元格坐标: table_cells 数组
  58. """
  59. def __init__(self, output_dir: str, save_images: bool = True):
  60. """
  61. 初始化格式化器
  62. Args:
  63. output_dir: 输出目录
  64. save_images: 是否保存图片
  65. """
  66. self.output_dir = Path(output_dir)
  67. self.output_dir.mkdir(parents=True, exist_ok=True)
  68. self.save_images = save_images
  69. # 创建 images 子目录
  70. self.images_dir = self.output_dir / "images"
  71. self.images_dir.mkdir(parents=True, exist_ok=True)
  72. def format_and_save(
  73. self,
  74. results: Dict[str, Any],
  75. doc_name: str = None
  76. ) -> Dict[str, List[str]]:
  77. """
  78. 格式化并保存结果
  79. Args:
  80. results: pipeline 处理结果
  81. doc_name: 文档名称(可选)
  82. Returns:
  83. 输出文件路径字典
  84. """
  85. if doc_name is None:
  86. doc_name = Path(results.get('document_path', 'unknown')).stem
  87. output_paths = {
  88. 'json': [],
  89. 'markdown': [],
  90. 'images': []
  91. }
  92. for page in results.get('pages', []):
  93. page_idx = page.get('page_idx', 0)
  94. page_name = f"{doc_name}_page_{page_idx + 1:03d}"
  95. # 1. 转换为 MinerU 格式
  96. mineru_elements = self._convert_page_to_mineru_format(
  97. page,
  98. results.get('metadata', {})
  99. )
  100. # 2. 保存 JSON
  101. json_path = self.output_dir / f"{page_name}.json"
  102. with open(json_path, 'w', encoding='utf-8') as f:
  103. json.dump(mineru_elements, f, ensure_ascii=False, indent=2)
  104. output_paths['json'].append(str(json_path))
  105. # 3. 生成并保存 Markdown
  106. md_content = self._generate_markdown(mineru_elements)
  107. md_path = self.output_dir / f"{page_name}.md"
  108. with open(md_path, 'w', encoding='utf-8') as f:
  109. f.write(md_content)
  110. output_paths['markdown'].append(str(md_path))
  111. # 4. 保存原始图片
  112. if self.save_images and 'original_image' in page:
  113. img_path = self._save_page_image(
  114. page['original_image'],
  115. page_name
  116. )
  117. if img_path:
  118. output_paths['images'].append(img_path)
  119. logger.info(f"✅ MinerU format output saved to: {self.output_dir}")
  120. return output_paths
  121. def _convert_page_to_mineru_format(
  122. self,
  123. page: Dict[str, Any],
  124. metadata: Dict[str, Any]
  125. ) -> List[Dict[str, Any]]:
  126. """
  127. 将页面转换为 MinerU 标准格式
  128. Args:
  129. page: 页面数据
  130. metadata: 文档元数据
  131. Returns:
  132. MinerU 格式的元素列表
  133. """
  134. mineru_elements = []
  135. page_idx = page.get('page_idx', 0)
  136. for element in page.get('elements', []):
  137. mineru_elem = self._convert_element_to_mineru(element, page_idx)
  138. if mineru_elem:
  139. mineru_elements.append(mineru_elem)
  140. return mineru_elements
  141. def _convert_element_to_mineru(
  142. self,
  143. element: Dict[str, Any],
  144. page_idx: int
  145. ) -> Optional[Dict[str, Any]]:
  146. """
  147. 将单个元素转换为 MinerU 格式
  148. Args:
  149. element: 元素数据
  150. page_idx: 页面索引
  151. Returns:
  152. MinerU 格式的元素
  153. """
  154. elem_type = element.get('type', '')
  155. bbox = element.get('bbox', [0, 0, 0, 0])
  156. content = element.get('content', {})
  157. # 确保 bbox 是整数列表
  158. bbox = [int(x) for x in bbox[:4]] if bbox else [0, 0, 0, 0]
  159. result = {
  160. 'type': self._map_element_type(elem_type),
  161. 'bbox': bbox,
  162. 'page_idx': page_idx
  163. }
  164. # 处理不同类型的元素
  165. if elem_type in ['text', 'title', 'ref_text', 'ocr_text']:
  166. text = content.get('text', '') if isinstance(content, dict) else str(content)
  167. result['text'] = text
  168. elif elem_type in ['header', 'footer']:
  169. text = content.get('text', '') if isinstance(content, dict) else str(content)
  170. result['text'] = text
  171. elif elem_type == 'list':
  172. result['sub_type'] = 'text'
  173. list_items = content.get('list_items', []) if isinstance(content, dict) else []
  174. if not list_items and isinstance(content, dict):
  175. text = content.get('text', '')
  176. if text:
  177. list_items = [text]
  178. result['list_items'] = list_items
  179. elif elem_type in ['table', 'table_body']:
  180. result = self._convert_table_element(element, page_idx)
  181. elif elem_type in ['image', 'image_body', 'figure']:
  182. result = self._convert_image_element(element, page_idx)
  183. elif elem_type in ['interline_equation', 'inline_equation', 'equation']:
  184. latex = content.get('latex', '') if isinstance(content, dict) else ''
  185. result['text'] = latex
  186. result['text_format'] = 'latex'
  187. return result
  188. def _convert_table_element(
  189. self,
  190. element: Dict[str, Any],
  191. page_idx: int
  192. ) -> Dict[str, Any]:
  193. """
  194. 转换表格元素为 MinerU 格式
  195. 关键:
  196. 1. table_body: 带 data-bbox 属性的 HTML
  197. 2. table_cells: 单元格坐标数组
  198. 3. image_rotation_angle: 旋转角度
  199. 4. skew_angle: 倾斜角度
  200. """
  201. bbox = element.get('bbox', [0, 0, 0, 0])
  202. bbox = [int(x) for x in bbox[:4]] if bbox else [0, 0, 0, 0]
  203. content = element.get('content', {})
  204. # 获取表格 HTML
  205. table_html = content.get('html', '') or content.get('original_html', '')
  206. # 获取表格标题和脚注
  207. table_caption = content.get('table_caption', [])
  208. if isinstance(table_caption, str):
  209. table_caption = [table_caption] if table_caption else []
  210. table_footnote = content.get('table_footnote', [])
  211. if isinstance(table_footnote, str):
  212. table_footnote = [table_footnote] if table_footnote else []
  213. # 保存表格图片
  214. img_path = ''
  215. if self.save_images and 'table_image' in content:
  216. img_path = self._save_table_image(
  217. content['table_image'],
  218. f"table_{page_idx}_{bbox[0]}_{bbox[1]}"
  219. )
  220. result = {
  221. 'type': 'table',
  222. 'img_path': img_path,
  223. 'table_caption': table_caption,
  224. 'table_footnote': table_footnote,
  225. 'table_body': table_html,
  226. 'bbox': bbox,
  227. 'page_idx': page_idx
  228. }
  229. # 添加单元格坐标信息
  230. cells = content.get('cells', [])
  231. if cells:
  232. result['table_cells'] = self._format_table_cells(cells)
  233. # 添加旋转和倾斜信息
  234. if 'table_angle' in content:
  235. result['image_rotation_angle'] = float(content['table_angle'])
  236. if 'skew_angle' in content:
  237. result['skew_angle'] = float(content['skew_angle'])
  238. return result
  239. def _format_table_cells(self, cells: List[Dict]) -> List[Dict[str, Any]]:
  240. """
  241. 格式化表格单元格为 MinerU 格式
  242. MinerU 格式:
  243. {
  244. "type": "table_cell",
  245. "text": "单元格内容",
  246. "matched_text": "OCR匹配文本",
  247. "bbox": [y1, x1, y2, x2], # 注意坐标顺序
  248. "row": 2,
  249. "col": 1,
  250. "score": 100.0,
  251. "paddle_bbox_indices": [11, 12]
  252. }
  253. """
  254. formatted_cells = []
  255. for cell in cells:
  256. formatted_cell = {
  257. 'type': 'table_cell',
  258. 'text': cell.get('text', ''),
  259. 'matched_text': cell.get('matched_text', cell.get('text', '')),
  260. 'bbox': cell.get('bbox', [0, 0, 0, 0]),
  261. 'row': cell.get('row', 0),
  262. 'col': cell.get('col', 0),
  263. 'score': cell.get('score', 100.0),
  264. 'paddle_bbox_indices': cell.get('paddle_bbox_indices',
  265. cell.get('paddle_indices', []))
  266. }
  267. # 确保 bbox 是整数列表
  268. if formatted_cell['bbox']:
  269. formatted_cell['bbox'] = [int(x) for x in formatted_cell['bbox'][:4]]
  270. formatted_cells.append(formatted_cell)
  271. return formatted_cells
  272. def _convert_image_element(
  273. self,
  274. element: Dict[str, Any],
  275. page_idx: int
  276. ) -> Dict[str, Any]:
  277. """转换图片元素为 MinerU 格式"""
  278. bbox = element.get('bbox', [0, 0, 0, 0])
  279. bbox = [int(x) for x in bbox[:4]] if bbox else [0, 0, 0, 0]
  280. content = element.get('content', {})
  281. # 保存图片
  282. img_path = ''
  283. if self.save_images and 'image_data' in content:
  284. img_path = self._save_content_image(
  285. content['image_data'],
  286. f"image_{page_idx}_{bbox[0]}_{bbox[1]}"
  287. )
  288. return {
  289. 'type': 'image',
  290. 'img_path': img_path,
  291. 'image_caption': content.get('caption', []),
  292. 'image_footnote': content.get('footnote', []),
  293. 'bbox': bbox,
  294. 'page_idx': page_idx
  295. }
  296. def _map_element_type(self, elem_type: str) -> str:
  297. """映射元素类型到 MinerU 标准类型"""
  298. type_mapping = {
  299. 'text': 'text',
  300. 'title': 'title',
  301. 'ref_text': 'ref_text',
  302. 'ocr_text': 'text',
  303. 'header': 'header',
  304. 'footer': 'footer',
  305. 'table': 'table',
  306. 'table_body': 'table',
  307. 'image': 'image',
  308. 'image_body': 'image',
  309. 'figure': 'image',
  310. 'list': 'list',
  311. 'interline_equation': 'equation',
  312. 'inline_equation': 'equation',
  313. 'equation': 'equation',
  314. }
  315. return type_mapping.get(elem_type, elem_type)
  316. def _generate_markdown(self, elements: List[Dict[str, Any]]) -> str:
  317. """
  318. 生成 MinerU 格式的 Markdown
  319. 格式:
  320. <!-- bbox: [x1, y1, x2, y2] -->
  321. 内容
  322. """
  323. md_lines = []
  324. for elem in elements:
  325. elem_type = elem.get('type', '')
  326. bbox = elem.get('bbox', [])
  327. # 添加 bbox 注释
  328. bbox_comment = f"<!-- bbox: {bbox} -->"
  329. if elem_type == 'header':
  330. md_lines.append(bbox_comment)
  331. md_lines.append(f"<!-- 页眉: {elem.get('text', '')} -->")
  332. md_lines.append("")
  333. elif elem_type == 'footer':
  334. md_lines.append(bbox_comment)
  335. md_lines.append(f"<!-- 页脚: {elem.get('text', '')} -->")
  336. md_lines.append("")
  337. elif elem_type == 'title':
  338. md_lines.append(bbox_comment)
  339. level = elem.get('text_level', 1)
  340. md_lines.append(f"{'#' * min(level, 6)} {elem.get('text', '')}")
  341. md_lines.append("")
  342. elif elem_type == 'text':
  343. md_lines.append(bbox_comment)
  344. md_lines.append(elem.get('text', ''))
  345. md_lines.append("")
  346. elif elem_type == 'list':
  347. md_lines.append(bbox_comment)
  348. for item in elem.get('list_items', []):
  349. md_lines.append(item)
  350. md_lines.append("")
  351. elif elem_type == 'table':
  352. # 表格标题
  353. for caption in elem.get('table_caption', []):
  354. md_lines.append(f"**{caption}**")
  355. md_lines.append("")
  356. md_lines.append(bbox_comment)
  357. # 表格 HTML
  358. table_body = elem.get('table_body', '')
  359. if table_body:
  360. md_lines.append(table_body)
  361. else:
  362. # 如果没有 HTML,显示图片
  363. img_path = elem.get('img_path', '')
  364. if img_path:
  365. md_lines.append(f"![Image]({img_path})")
  366. md_lines.append("")
  367. elif elem_type == 'image':
  368. md_lines.append(bbox_comment)
  369. img_path = elem.get('img_path', '')
  370. if img_path:
  371. md_lines.append(f"![Image]({img_path})")
  372. md_lines.append("")
  373. elif elem_type == 'equation':
  374. md_lines.append(bbox_comment)
  375. latex = elem.get('text', '')
  376. if latex:
  377. md_lines.append(f"$$\n{latex}\n$$")
  378. md_lines.append("")
  379. return '\n'.join(md_lines)
  380. def _save_page_image(
  381. self,
  382. image: np.ndarray,
  383. name: str
  384. ) -> Optional[str]:
  385. """保存页面图片"""
  386. try:
  387. # 生成唯一文件名
  388. img_hash = hashlib.md5(image.tobytes()[:1000]).hexdigest()[:16]
  389. filename = f"{name}_{img_hash}.png"
  390. filepath = self.images_dir / filename
  391. if isinstance(image, np.ndarray):
  392. if len(image.shape) == 3 and image.shape[2] == 3:
  393. # BGR to RGB
  394. image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
  395. Image.fromarray(image).save(filepath)
  396. return f"images/{filename}"
  397. except Exception as e:
  398. logger.warning(f"Failed to save page image: {e}")
  399. return None
  400. def _save_table_image(
  401. self,
  402. image: np.ndarray,
  403. name: str
  404. ) -> str:
  405. """保存表格图片"""
  406. try:
  407. img_hash = hashlib.md5(image.tobytes()[:1000]).hexdigest()[:32]
  408. filename = f"{img_hash}.jpg"
  409. filepath = self.images_dir / filename
  410. if isinstance(image, np.ndarray):
  411. if len(image.shape) == 3 and image.shape[2] == 3:
  412. image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
  413. Image.fromarray(image).save(filepath, quality=95)
  414. return f"images/{filename}"
  415. except Exception as e:
  416. logger.warning(f"Failed to save table image: {e}")
  417. return ""
  418. def _save_content_image(
  419. self,
  420. image: np.ndarray,
  421. name: str
  422. ) -> str:
  423. """保存内容图片"""
  424. return self._save_table_image(image, name)
  425. class MinerUFormatConverter:
  426. """
  427. MinerU 格式转换器
  428. 用于将 pipeline 的内部格式转换为 MinerU 标准输出格式
  429. """
  430. @staticmethod
  431. def convert_pipeline_result_to_mineru(
  432. results: Dict[str, Any]
  433. ) -> List[Dict[str, Any]]:
  434. """
  435. 将 pipeline 结果转换为 MinerU content_list 格式
  436. Args:
  437. results: pipeline 处理结果
  438. Returns:
  439. MinerU content_list 格式的结果
  440. """
  441. content_list = []
  442. for page in results.get('pages', []):
  443. page_idx = page.get('page_idx', 0)
  444. page_size = page.get('image_shape', [0, 0])[:2]
  445. for element in page.get('elements', []):
  446. content = MinerUFormatConverter._element_to_content(
  447. element, page_idx, page_size
  448. )
  449. if content:
  450. content_list.append(content)
  451. return content_list
  452. @staticmethod
  453. def _element_to_content(
  454. element: Dict[str, Any],
  455. page_idx: int,
  456. page_size: Tuple[int, int]
  457. ) -> Optional[Dict[str, Any]]:
  458. """将元素转换为 MinerU content 格式"""
  459. elem_type = element.get('type', '')
  460. bbox = element.get('bbox', [0, 0, 0, 0])
  461. content = element.get('content', {})
  462. # 转换 bbox 为相对坐标(千分比)
  463. page_height, page_width = page_size if len(page_size) >= 2 else (1, 1)
  464. if page_width > 0 and page_height > 0 and bbox:
  465. relative_bbox = [
  466. int(bbox[0] * 1000 / page_width),
  467. int(bbox[1] * 1000 / page_height),
  468. int(bbox[2] * 1000 / page_width),
  469. int(bbox[3] * 1000 / page_height)
  470. ]
  471. else:
  472. relative_bbox = bbox
  473. result = {
  474. 'bbox': relative_bbox,
  475. 'page_idx': page_idx
  476. }
  477. if elem_type in ['text', 'title', 'ref_text', 'header', 'footer']:
  478. text = content.get('text', '') if isinstance(content, dict) else str(content)
  479. result['type'] = elem_type
  480. result['text'] = text
  481. elif elem_type in ['table', 'table_body']:
  482. result['type'] = 'table'
  483. result['img_path'] = content.get('img_path', '')
  484. result['table_caption'] = content.get('table_caption', [])
  485. result['table_footnote'] = content.get('table_footnote', [])
  486. result['table_body'] = content.get('html', '')
  487. # 添加单元格信息
  488. cells = content.get('cells', [])
  489. if cells:
  490. result['table_cells'] = cells
  491. # 添加旋转信息
  492. if 'table_angle' in content:
  493. result['image_rotation_angle'] = content['table_angle']
  494. if 'skew_angle' in content:
  495. result['skew_angle'] = content['skew_angle']
  496. elif elem_type in ['image', 'figure']:
  497. result['type'] = 'image'
  498. result['img_path'] = content.get('img_path', '')
  499. result['image_caption'] = content.get('caption', [])
  500. result['image_footnote'] = content.get('footnote', [])
  501. elif elem_type == 'list':
  502. result['type'] = 'list'
  503. result['sub_type'] = 'text'
  504. result['list_items'] = content.get('list_items', [])
  505. else:
  506. return None
  507. return result
  508. def save_mineru_format(
  509. results: Dict[str, Any],
  510. output_dir: str,
  511. doc_name: str = None,
  512. save_images: bool = True
  513. ) -> Dict[str, List[str]]:
  514. """
  515. 便捷函数:保存为 MinerU 格式
  516. Args:
  517. results: pipeline 处理结果
  518. output_dir: 输出目录
  519. doc_name: 文档名称
  520. save_images: 是否保存图片
  521. Returns:
  522. 输出文件路径
  523. """
  524. formatter = MinerUOutputFormatter(output_dir, save_images)
  525. return formatter.format_and_save(results, doc_name)