output_formatter_v2.py 38 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058
  1. """
  2. 统一输出格式化器
  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. import json
  12. import os
  13. import sys
  14. import hashlib
  15. from pathlib import Path
  16. from typing import Dict, Any, List, Optional, Union, Tuple
  17. from loguru import logger
  18. import numpy as np
  19. from PIL import Image, ImageDraw, ImageFont
  20. import cv2
  21. # 导入 MinerU 组件
  22. mineru_path = Path(__file__).parents[3]
  23. if str(mineru_path) not in sys.path:
  24. sys.path.insert(0, str(mineru_path))
  25. try:
  26. from mineru.backend.vlm.vlm_middle_json_mkcontent import union_make as vlm_union_make
  27. from mineru.utils.enum_class import MakeMode, BlockType, ContentType
  28. MINERU_AVAILABLE = True
  29. except ImportError as e:
  30. logger.warning(f"MinerU components not available: {e}")
  31. MINERU_AVAILABLE = False
  32. # 占位符定义
  33. class MakeMode:
  34. MM_MD = 'mm_md'
  35. NLP_MD = 'nlp_md'
  36. class OutputFormatterV2:
  37. """
  38. 统一输出格式化器
  39. 严格遵循 MinerU mineru_vllm_results_cell_bbox 格式:
  40. - middle.json: MinerU 标准格式,用于生成 Markdown
  41. - page_xxx.json: 每页独立的 JSON,包含 table_cells
  42. - Markdown: 带 bbox 注释
  43. - 表格: HTML 格式,带 data-bbox 属性
  44. """
  45. # 颜色映射(与 MinerU 保持一致)
  46. COLOR_MAP = {
  47. 'title': (102, 102, 255), # 蓝色
  48. 'text': (153, 0, 76), # 深红
  49. 'image': (153, 255, 51), # 绿色
  50. 'image_body': (153, 255, 51),
  51. 'image_caption': (102, 178, 255),
  52. 'image_footnote': (255, 178, 102),
  53. 'table': (204, 204, 0), # 黄色
  54. 'table_body': (204, 204, 0),
  55. 'table_caption': (255, 255, 102),
  56. 'table_footnote': (229, 255, 204),
  57. 'interline_equation': (0, 255, 0), # 亮绿
  58. 'inline_equation': (0, 200, 0),
  59. 'list': (40, 169, 92),
  60. 'code': (102, 0, 204), # 紫色
  61. 'header': (128, 128, 128), # 灰色
  62. 'footer': (128, 128, 128),
  63. 'ref_text': (180, 180, 180),
  64. 'ocr_text': (153, 0, 76),
  65. 'error': (255, 0, 0), # 红色
  66. }
  67. # OCR 框颜色
  68. OCR_BOX_COLOR = (0, 255, 0) # 绿色
  69. CELL_BOX_COLOR = (255, 165, 0) # 橙色
  70. def __init__(self, output_dir: str):
  71. """
  72. 初始化格式化器
  73. Args:
  74. output_dir: 输出目录
  75. """
  76. self.output_dir = Path(output_dir)
  77. self.output_dir.mkdir(parents=True, exist_ok=True)
  78. def save_results(
  79. self,
  80. results: Dict[str, Any],
  81. output_config: Dict[str, Any]
  82. ) -> Dict[str, Any]:
  83. """
  84. 保存处理结果
  85. Args:
  86. results: 处理结果
  87. output_config: 输出配置
  88. Returns:
  89. 输出文件路径字典
  90. """
  91. output_paths = {
  92. 'images': [],
  93. 'json_pages': [],
  94. }
  95. # 创建文档输出目录
  96. doc_name = Path(results['document_path']).stem
  97. doc_output_dir = self.output_dir / doc_name
  98. doc_output_dir.mkdir(parents=True, exist_ok=True)
  99. # 创建 images 子目录
  100. images_dir = doc_output_dir / 'images'
  101. images_dir.mkdir(exist_ok=True)
  102. # 1. 首先保存图片元素(设置 image_path)
  103. image_paths = self._save_image_elements(results, images_dir, doc_name)
  104. if image_paths:
  105. output_paths['images'] = image_paths
  106. # 2. 转换为 MinerU middle.json 格式(用于 union_make)
  107. middle_json = self._convert_to_middle_json(results)
  108. # 3. 保存 middle.json
  109. if output_config.get('save_json', True):
  110. json_path = doc_output_dir / f"{doc_name}_middle.json"
  111. with open(json_path, 'w', encoding='utf-8') as f:
  112. json.dump(middle_json, f, ensure_ascii=False, indent=2)
  113. output_paths['middle_json'] = str(json_path)
  114. logger.info(f"📄 Middle JSON saved: {json_path}")
  115. # 4. 保存每页独立的 mineru_vllm_results_cell_bbox 格式 JSON
  116. if output_config.get('save_page_json', True):
  117. page_json_paths = self._save_page_jsons(results, doc_output_dir, doc_name)
  118. output_paths['json_pages'] = page_json_paths
  119. # 5. 保存 Markdown
  120. if output_config.get('save_markdown', True):
  121. md_path = self._save_markdown(results, middle_json, doc_output_dir, doc_name)
  122. output_paths['markdown'] = str(md_path)
  123. # 6. 保存表格 HTML
  124. if output_config.get('save_html', True):
  125. html_dir = self._save_table_htmls(results, doc_output_dir, doc_name)
  126. output_paths['table_htmls'] = str(html_dir)
  127. # 7. Debug 模式:保存可视化图片
  128. if output_config.get('save_layout_image', False):
  129. layout_paths = self._save_layout_images(
  130. results, doc_output_dir, doc_name,
  131. draw_type_label=output_config.get('draw_type_label', True),
  132. draw_bbox_number=output_config.get('draw_bbox_number', True)
  133. )
  134. output_paths['layout_images'] = layout_paths
  135. if output_config.get('save_ocr_image', False):
  136. ocr_paths = self._save_ocr_images(results, doc_output_dir, doc_name)
  137. output_paths['ocr_images'] = ocr_paths
  138. logger.info(f"✅ All results saved to: {doc_output_dir}")
  139. return output_paths
  140. # ==================== MinerU middle.json 格式 ====================
  141. def _convert_to_middle_json(self, results: Dict[str, Any]) -> Dict[str, Any]:
  142. """
  143. 转换为 MinerU 标准 middle.json 格式
  144. 用于 vlm_union_make 生成 Markdown
  145. """
  146. middle_json = {
  147. "pdf_info": [],
  148. "_backend": "vlm",
  149. "_scene": results.get('scene', 'unknown'),
  150. "_version_name": "2.5.0"
  151. }
  152. for page in results.get('pages', []):
  153. page_info = {
  154. 'page_idx': page['page_idx'],
  155. 'page_size': list(page.get('image_shape', [0, 0])[:2][::-1]),
  156. 'angle': page.get('angle', 0),
  157. 'para_blocks': [],
  158. 'discarded_blocks': []
  159. }
  160. for element in page.get('elements', []):
  161. block = self._element_to_middle_block(element)
  162. if block:
  163. elem_type = element.get('type', '')
  164. if elem_type in ['header', 'footer', 'page_number', 'aside_text', 'abandon', 'discarded']:
  165. page_info['discarded_blocks'].append(block)
  166. else:
  167. page_info['para_blocks'].append(block)
  168. middle_json['pdf_info'].append(page_info)
  169. return middle_json
  170. def _element_to_middle_block(self, element: Dict[str, Any]) -> Optional[Dict[str, Any]]:
  171. """
  172. 将元素转换为 MinerU middle.json block 格式
  173. MinerU 期望的嵌套结构:
  174. - image 类型: { type: "image", blocks: [{ type: "image_body", lines: [...] }] }
  175. - table 类型: { type: "table", blocks: [{ type: "table_body", lines: [...] }] }
  176. """
  177. elem_type = element.get('type', '')
  178. bbox = element.get('bbox', [0, 0, 0, 0])
  179. content = element.get('content', {})
  180. block = {
  181. 'type': elem_type,
  182. 'bbox': bbox,
  183. 'angle': element.get('angle', 0),
  184. 'lines': []
  185. }
  186. # 文本类型
  187. if elem_type in ['text', 'title', 'ref_text', 'header', 'footer', 'ocr_text']:
  188. text = content.get('text', '') if isinstance(content, dict) else str(content)
  189. if text:
  190. block['lines'] = [{
  191. 'bbox': bbox,
  192. 'spans': [{
  193. 'bbox': bbox,
  194. 'type': 'text',
  195. 'content': text
  196. }]
  197. }]
  198. # 表格类型 - 嵌套结构
  199. elif elem_type in ['table', 'table_body']:
  200. table_html = content.get('html', '')
  201. cells = content.get('cells', [])
  202. block['type'] = 'table'
  203. block['blocks'] = [{
  204. 'type': 'table_body',
  205. 'bbox': bbox,
  206. 'angle': 0,
  207. 'lines': [{
  208. 'bbox': bbox,
  209. 'spans': [{
  210. 'bbox': bbox,
  211. 'type': 'table',
  212. 'html': table_html,
  213. 'cells': cells
  214. }]
  215. }]
  216. }]
  217. # 图片类型 - 嵌套结构
  218. # 注意:MinerU vlm_union_make 期望字段名是 'image_path',不是 'img_path'
  219. elif elem_type in ['image', 'image_body', 'figure']:
  220. block['type'] = 'image'
  221. block['blocks'] = [{
  222. 'type': 'image_body',
  223. 'bbox': bbox,
  224. 'angle': element.get('angle', 0),
  225. 'lines': [{
  226. 'bbox': bbox,
  227. 'spans': [{
  228. 'bbox': bbox,
  229. 'type': 'image',
  230. 'image_path': content.get('image_path', ''), # MinerU 期望 'image_path'
  231. 'description': content.get('description', '')
  232. }]
  233. }]
  234. }]
  235. # 公式类型
  236. elif elem_type in ['interline_equation', 'inline_equation', 'equation']:
  237. latex = content.get('latex', '')
  238. block['lines'] = [{
  239. 'bbox': bbox,
  240. 'spans': [{
  241. 'bbox': bbox,
  242. 'type': 'interline_equation' if 'interline' in elem_type else 'inline_equation',
  243. 'content': latex
  244. }]
  245. }]
  246. # 表格/图片附属文本
  247. elif elem_type in ['table_caption', 'table_footnote', 'image_caption', 'image_footnote']:
  248. text = content.get('text', '') if isinstance(content, dict) else str(content)
  249. if text:
  250. block['lines'] = [{
  251. 'bbox': bbox,
  252. 'spans': [{
  253. 'bbox': bbox,
  254. 'type': 'text',
  255. 'content': text
  256. }]
  257. }]
  258. # 丢弃类型
  259. elif elem_type in ['abandon', 'discarded']:
  260. block['type'] = 'abandon'
  261. return block
  262. # ==================== mineru_vllm_results_cell_bbox 格式 ====================
  263. def _save_page_jsons(
  264. self,
  265. results: Dict[str, Any],
  266. output_dir: Path,
  267. doc_name: str
  268. ) -> List[str]:
  269. """
  270. 保存每页独立的 JSON(mineru_vllm_results_cell_bbox 格式)
  271. 格式示例:
  272. [
  273. {
  274. "type": "table",
  275. "img_path": "images/xxx.jpg",
  276. "table_caption": [],
  277. "table_footnote": [],
  278. "table_body": "<table>...</table>",
  279. "bbox": [x1, y1, x2, y2],
  280. "page_idx": 0,
  281. "table_cells": [
  282. {
  283. "type": "table_cell",
  284. "text": "单元格内容",
  285. "matched_text": "OCR匹配文本",
  286. "bbox": [x1, y1, x2, y2],
  287. "row": 1,
  288. "col": 1,
  289. "score": 100.0,
  290. "paddle_bbox_indices": [0, 1]
  291. }
  292. ],
  293. "image_rotation_angle": 0,
  294. "skew_angle": 0.0
  295. }
  296. ]
  297. """
  298. saved_paths = []
  299. for page in results.get('pages', []):
  300. page_idx = page.get('page_idx', 0)
  301. page_name = f"{doc_name}_page_{page_idx + 1:03d}"
  302. # 转换为 mineru_vllm_results_cell_bbox 格式
  303. page_elements = []
  304. for element in page.get('elements', []):
  305. converted = self._element_to_cell_bbox_format(element, page_idx)
  306. if converted:
  307. page_elements.append(converted)
  308. # 保存 JSON
  309. json_path = output_dir / f"{page_name}.json"
  310. with open(json_path, 'w', encoding='utf-8') as f:
  311. json.dump(page_elements, f, ensure_ascii=False, indent=2)
  312. saved_paths.append(str(json_path))
  313. logger.debug(f"📄 Page JSON saved: {json_path}")
  314. if saved_paths:
  315. logger.info(f"📄 {len(saved_paths)} page JSONs saved")
  316. return saved_paths
  317. def _element_to_cell_bbox_format(
  318. self,
  319. element: Dict[str, Any],
  320. page_idx: int
  321. ) -> Optional[Dict[str, Any]]:
  322. """
  323. 将元素转换为 mineru_vllm_results_cell_bbox 格式
  324. """
  325. elem_type = element.get('type', '')
  326. bbox = element.get('bbox', [0, 0, 0, 0])
  327. content = element.get('content', {})
  328. # 确保 bbox 是整数列表
  329. bbox = [int(x) for x in bbox[:4]] if bbox else [0, 0, 0, 0]
  330. result = {
  331. 'bbox': bbox,
  332. 'page_idx': page_idx
  333. }
  334. # 文本类型
  335. if elem_type in ['text', 'title', 'ref_text', 'ocr_text']:
  336. text = content.get('text', '') if isinstance(content, dict) else str(content)
  337. result['type'] = 'text' if elem_type != 'title' else 'title'
  338. result['text'] = text
  339. if elem_type == 'title':
  340. result['text_level'] = element.get('level', 1)
  341. # 表格类型
  342. elif elem_type in ['table', 'table_body']:
  343. result['type'] = 'table'
  344. result['img_path'] = content.get('table_image_path', '')
  345. result['table_caption'] = self._ensure_list(content.get('table_caption', []))
  346. result['table_footnote'] = self._ensure_list(content.get('table_footnote', []))
  347. result['table_body'] = content.get('html', '')
  348. # 关键:table_cells 数组
  349. cells = content.get('cells', [])
  350. if cells:
  351. result['table_cells'] = self._format_table_cells(cells)
  352. # 旋转和倾斜信息
  353. if 'table_angle' in content:
  354. result['image_rotation_angle'] = float(content['table_angle'])
  355. if 'skew_angle' in content:
  356. result['skew_angle'] = float(content['skew_angle'])
  357. # 图片类型
  358. elif elem_type in ['image', 'image_body', 'figure']:
  359. result['type'] = 'image'
  360. # page JSON 需要完整的相对路径
  361. image_filename = content.get('image_path', '')
  362. result['img_path'] = f"images/{image_filename}" if image_filename else ''
  363. result['image_caption'] = self._ensure_list(content.get('caption', []))
  364. result['image_footnote'] = self._ensure_list(content.get('footnote', []))
  365. # 公式类型
  366. elif elem_type in ['interline_equation', 'inline_equation', 'equation']:
  367. result['type'] = 'equation'
  368. result['text'] = content.get('latex', '') if isinstance(content, dict) else ''
  369. result['text_format'] = 'latex'
  370. # 列表类型
  371. elif elem_type == 'list':
  372. result['type'] = 'list'
  373. result['sub_type'] = 'text'
  374. result['list_items'] = content.get('list_items', []) if isinstance(content, dict) else []
  375. # 页眉页脚
  376. elif elem_type in ['header', 'footer']:
  377. result['type'] = elem_type
  378. result['text'] = content.get('text', '') if isinstance(content, dict) else str(content)
  379. else:
  380. return None
  381. return result
  382. def _format_table_cells(self, cells: List[Dict]) -> List[Dict[str, Any]]:
  383. """
  384. 格式化表格单元格为 mineru_vllm_results_cell_bbox 格式
  385. 输出格式:
  386. {
  387. "type": "table_cell",
  388. "text": "单元格内容",
  389. "matched_text": "OCR匹配文本",
  390. "bbox": [x1, y1, x2, y2],
  391. "row": 1,
  392. "col": 1,
  393. "score": 100.0,
  394. "paddle_bbox_indices": [0, 1]
  395. }
  396. """
  397. formatted_cells = []
  398. for cell in cells:
  399. formatted_cell = {
  400. 'type': 'table_cell',
  401. 'text': cell.get('text', ''),
  402. 'matched_text': cell.get('matched_text', cell.get('text', '')),
  403. 'bbox': [float(x) for x in cell.get('bbox', [0, 0, 0, 0])[:4]],
  404. 'row': cell.get('row', 0),
  405. 'col': cell.get('col', 0),
  406. 'score': float(cell.get('score', 100.0)),
  407. 'paddle_bbox_indices': cell.get('paddle_bbox_indices',
  408. cell.get('paddle_indices', []))
  409. }
  410. formatted_cells.append(formatted_cell)
  411. return formatted_cells
  412. def _ensure_list(self, value) -> List:
  413. """确保值是列表"""
  414. if value is None:
  415. return []
  416. if isinstance(value, str):
  417. return [value] if value else []
  418. if isinstance(value, list):
  419. return value
  420. return [str(value)]
  421. # ==================== Markdown 生成 ====================
  422. def _save_markdown(
  423. self,
  424. results: Dict[str, Any],
  425. middle_json: Dict[str, Any],
  426. output_dir: Path,
  427. doc_name: str
  428. ) -> Path:
  429. """
  430. 保存 Markdown 文件
  431. 优先使用 MinerU union_make,降级使用自定义实现
  432. """
  433. md_path = output_dir / f"{doc_name}.md"
  434. if MINERU_AVAILABLE:
  435. try:
  436. # image_path 只保存文件名,vlm_union_make 会添加此前缀
  437. img_bucket_path = "images"
  438. markdown_content = vlm_union_make(
  439. middle_json['pdf_info'],
  440. MakeMode.MM_MD,
  441. img_bucket_path
  442. )
  443. if markdown_content:
  444. if isinstance(markdown_content, list):
  445. markdown_content = '\n\n'.join(markdown_content)
  446. header = self._generate_markdown_header(results)
  447. markdown_content = header + str(markdown_content)
  448. with open(md_path, 'w', encoding='utf-8') as f:
  449. f.write(markdown_content)
  450. logger.info(f"📝 Markdown saved (MinerU format): {md_path}")
  451. return md_path
  452. except Exception as e:
  453. logger.warning(f"MinerU union_make failed: {e}, falling back to custom implementation")
  454. # 降级方案
  455. markdown_content = self._generate_markdown_fallback(results)
  456. with open(md_path, 'w', encoding='utf-8') as f:
  457. f.write(markdown_content)
  458. logger.info(f"📝 Markdown saved (fallback): {md_path}")
  459. return md_path
  460. def _generate_markdown_header(self, results: Dict[str, Any]) -> str:
  461. """生成 Markdown 文件头"""
  462. return f"""---
  463. scene: {results.get('scene', 'unknown')}
  464. document: {results.get('document_path', '')}
  465. pages: {len(results.get('pages', []))}
  466. ---
  467. """
  468. def _generate_markdown_fallback(self, results: Dict[str, Any]) -> str:
  469. """降级方案:自定义 Markdown 生成"""
  470. md_lines = [
  471. f"---",
  472. f"scene: {results.get('scene', 'unknown')}",
  473. f"document: {results.get('document_path', '')}",
  474. f"pages: {len(results.get('pages', []))}",
  475. f"---",
  476. "",
  477. ]
  478. for page in results.get('pages', []):
  479. page_idx = page.get('page_idx', 0)
  480. for element in page.get('elements', []):
  481. elem_type = element.get('type', '')
  482. content = element.get('content', {})
  483. bbox = element.get('bbox', [])
  484. # 添加 bbox 注释
  485. if bbox:
  486. md_lines.append(f"<!-- bbox: {bbox} -->")
  487. if elem_type == 'title':
  488. text = content.get('text', '') if isinstance(content, dict) else str(content)
  489. level = element.get('level', 1)
  490. md_lines.append(f"{'#' * min(level, 6)} {text}")
  491. md_lines.append("")
  492. elif elem_type in ['text', 'ocr_text', 'ref_text']:
  493. text = content.get('text', '') if isinstance(content, dict) else str(content)
  494. if text:
  495. md_lines.append(text)
  496. md_lines.append("")
  497. elif elem_type in ['table', 'table_body']:
  498. # 表格标题
  499. for caption in self._ensure_list(content.get('table_caption', [])):
  500. md_lines.append(f"**{caption}**")
  501. html = content.get('html', '')
  502. if html:
  503. md_lines.append(f"\n{html}\n")
  504. md_lines.append("")
  505. elif elem_type in ['image', 'image_body', 'figure']:
  506. img_filename = content.get('image_path', '')
  507. if img_filename:
  508. md_lines.append(f"![](images/{img_filename})")
  509. md_lines.append("")
  510. elif elem_type in ['interline_equation', 'inline_equation', 'equation']:
  511. latex = content.get('latex', '')
  512. if latex:
  513. md_lines.append(f"$$\n{latex}\n$$")
  514. md_lines.append("")
  515. return '\n'.join(md_lines)
  516. # ==================== 图片保存 ====================
  517. def _save_image_elements(
  518. self,
  519. results: Dict[str, Any],
  520. images_dir: Path,
  521. doc_name: str
  522. ) -> List[str]:
  523. """保存图片元素"""
  524. saved_paths = []
  525. image_count = 0
  526. for page in results.get('pages', []):
  527. page_idx = page.get('page_idx', 0)
  528. for element in page.get('elements', []):
  529. if element.get('type') in ['image', 'image_body', 'figure']:
  530. content = element.get('content', {})
  531. image_data = content.get('image_data')
  532. if image_data is not None:
  533. image_count += 1
  534. image_filename = f"{doc_name}_page_{page_idx + 1}_image_{image_count}.png"
  535. image_path = images_dir / image_filename
  536. try:
  537. if isinstance(image_data, np.ndarray):
  538. cv2.imwrite(str(image_path), image_data)
  539. else:
  540. Image.fromarray(image_data).save(image_path)
  541. # 更新路径(只保存文件名,不含 images/ 前缀)
  542. # vlm_union_make 会自动添加 img_bucket_path 前缀
  543. content['image_path'] = image_filename
  544. content.pop('image_data', None)
  545. saved_paths.append(str(image_path))
  546. logger.debug(f"🖼️ Image saved: {image_path}")
  547. except Exception as e:
  548. logger.warning(f"Failed to save image: {e}")
  549. if image_count > 0:
  550. logger.info(f"🖼️ {image_count} images saved to: {images_dir}")
  551. return saved_paths
  552. # ==================== 表格 HTML ====================
  553. def _save_table_htmls(
  554. self,
  555. results: Dict[str, Any],
  556. output_dir: Path,
  557. doc_name: str
  558. ) -> Path:
  559. """保存表格 HTML 文件"""
  560. tables_dir = output_dir / 'tables'
  561. tables_dir.mkdir(exist_ok=True)
  562. table_count = 0
  563. for page in results.get('pages', []):
  564. page_idx = page.get('page_idx', 0)
  565. for element in page.get('elements', []):
  566. if element.get('type') in ['table', 'table_body']:
  567. table_count += 1
  568. content = element.get('content', {})
  569. html = content.get('html', '')
  570. cells = content.get('cells', [])
  571. if html:
  572. full_html = self._generate_table_html_with_styles(
  573. html, cells, doc_name, page_idx, table_count
  574. )
  575. html_path = tables_dir / f"{doc_name}_table_{table_count}_page_{page_idx + 1}.html"
  576. with open(html_path, 'w', encoding='utf-8') as f:
  577. f.write(full_html)
  578. if table_count > 0:
  579. logger.info(f"📊 {table_count} tables saved to: {tables_dir}")
  580. return tables_dir
  581. def _generate_table_html_with_styles(
  582. self,
  583. table_html: str,
  584. cells: List[Dict],
  585. doc_name: str,
  586. page_idx: int,
  587. table_idx: int
  588. ) -> str:
  589. """生成带样式的完整 HTML"""
  590. cells_json = json.dumps(cells, ensure_ascii=False, indent=2) if cells else "[]"
  591. return f"""<!DOCTYPE html>
  592. <html lang="zh-CN">
  593. <head>
  594. <meta charset="UTF-8">
  595. <meta name="viewport" content="width=device-width, initial-scale=1.0">
  596. <title>{doc_name} - Table {table_idx}</title>
  597. <style>
  598. body {{
  599. font-family: Arial, "Microsoft YaHei", sans-serif;
  600. margin: 20px;
  601. background-color: #f5f5f5;
  602. }}
  603. .container {{
  604. max-width: 1400px;
  605. margin: 0 auto;
  606. background-color: white;
  607. padding: 20px;
  608. box-shadow: 0 0 10px rgba(0,0,0,0.1);
  609. border-radius: 8px;
  610. }}
  611. .meta {{
  612. color: #666;
  613. font-size: 0.9em;
  614. margin-bottom: 20px;
  615. padding-bottom: 10px;
  616. border-bottom: 1px solid #ddd;
  617. }}
  618. table {{
  619. border-collapse: collapse;
  620. width: 100%;
  621. margin: 20px 0;
  622. }}
  623. th, td {{
  624. border: 1px solid #ddd;
  625. padding: 8px 12px;
  626. text-align: left;
  627. }}
  628. th {{
  629. background-color: #f2f2f2;
  630. font-weight: bold;
  631. }}
  632. tr:hover {{
  633. background-color: #f9f9f9;
  634. }}
  635. td[data-bbox], th[data-bbox] {{
  636. position: relative;
  637. }}
  638. td[data-bbox]:hover::after, th[data-bbox]:hover::after {{
  639. content: attr(data-bbox);
  640. position: absolute;
  641. bottom: 100%;
  642. left: 0;
  643. background: #333;
  644. color: white;
  645. padding: 2px 6px;
  646. font-size: 10px;
  647. border-radius: 3px;
  648. white-space: nowrap;
  649. z-index: 100;
  650. }}
  651. .cells-info {{
  652. margin-top: 30px;
  653. padding: 15px;
  654. background-color: #f8f9fa;
  655. border-radius: 5px;
  656. }}
  657. .cells-info summary {{
  658. cursor: pointer;
  659. font-weight: bold;
  660. color: #333;
  661. }}
  662. .cells-info pre {{
  663. background-color: #2d2d2d;
  664. color: #f8f8f2;
  665. padding: 15px;
  666. border-radius: 5px;
  667. overflow-x: auto;
  668. font-size: 12px;
  669. }}
  670. </style>
  671. </head>
  672. <body>
  673. <div class="container">
  674. <div class="meta">
  675. <p><strong>Document:</strong> {doc_name}</p>
  676. <p><strong>Page:</strong> {page_idx + 1}</p>
  677. <p><strong>Table:</strong> {table_idx}</p>
  678. <p><strong>Cells with coordinates:</strong> {len(cells)}</p>
  679. </div>
  680. {table_html}
  681. <div class="cells-info">
  682. <details>
  683. <summary>📍 单元格坐标数据 (JSON)</summary>
  684. <pre>{cells_json}</pre>
  685. </details>
  686. </div>
  687. </div>
  688. </body>
  689. </html>"""
  690. # ==================== Debug 可视化 ====================
  691. def _save_layout_images(
  692. self,
  693. results: Dict[str, Any],
  694. output_dir: Path,
  695. doc_name: str,
  696. draw_type_label: bool = True,
  697. draw_bbox_number: bool = True
  698. ) -> List[str]:
  699. """保存 Layout 可视化图片"""
  700. layout_paths = []
  701. for page in results.get('pages', []):
  702. page_idx = page.get('page_idx', 0)
  703. processed_image = page.get('original_image')
  704. if processed_image is None:
  705. processed_image = page.get('processed_image')
  706. if processed_image is None:
  707. logger.warning(f"Page {page_idx}: No image data found for layout visualization")
  708. continue
  709. if isinstance(processed_image, np.ndarray):
  710. image = Image.fromarray(processed_image).convert('RGB')
  711. elif isinstance(processed_image, Image.Image):
  712. image = processed_image.convert('RGB')
  713. else:
  714. continue
  715. draw = ImageDraw.Draw(image, 'RGBA')
  716. font = self._get_font(14)
  717. for idx, element in enumerate(page.get('elements', []), 1):
  718. elem_type = element.get('type', '')
  719. bbox = element.get('bbox', [0, 0, 0, 0])
  720. if len(bbox) < 4:
  721. continue
  722. x0, y0, x1, y1 = map(int, bbox[:4])
  723. color = self.COLOR_MAP.get(elem_type, (255, 0, 0))
  724. # 半透明填充
  725. overlay = Image.new('RGBA', image.size, (255, 255, 255, 0))
  726. overlay_draw = ImageDraw.Draw(overlay)
  727. overlay_draw.rectangle([x0, y0, x1, y1], fill=(*color, 50))
  728. image = Image.alpha_composite(image.convert('RGBA'), overlay).convert('RGB')
  729. draw = ImageDraw.Draw(image)
  730. # 边框
  731. draw.rectangle([x0, y0, x1, y1], outline=color, width=2)
  732. # 类型标签
  733. if draw_type_label:
  734. label = elem_type.replace('_', ' ').title()
  735. bbox_label = draw.textbbox((x0 + 2, y0 + 2), label, font=font)
  736. draw.rectangle(bbox_label, fill=color)
  737. draw.text((x0 + 2, y0 + 2), label, fill='white', font=font)
  738. # 序号
  739. if draw_bbox_number:
  740. number_text = str(idx)
  741. bbox_number = draw.textbbox((x1 - 25, y0 + 2), number_text, font=font)
  742. draw.rectangle(bbox_number, fill=(255, 0, 0))
  743. draw.text((x1 - 25, y0 + 2), number_text, fill='white', font=font)
  744. layout_path = output_dir / f"{doc_name}_page_{page_idx + 1}_layout.png"
  745. image.save(layout_path)
  746. layout_paths.append(str(layout_path))
  747. logger.info(f"🖼️ Layout image saved: {layout_path}")
  748. return layout_paths
  749. def _save_ocr_images(
  750. self,
  751. results: Dict[str, Any],
  752. output_dir: Path,
  753. doc_name: str
  754. ) -> List[str]:
  755. """保存 OCR 可视化图片"""
  756. ocr_paths = []
  757. for page in results.get('pages', []):
  758. page_idx = page.get('page_idx', 0)
  759. processed_image = page.get('original_image')
  760. if processed_image is None:
  761. processed_image = page.get('processed_image')
  762. if processed_image is None:
  763. logger.warning(f"Page {page_idx}: No image data found for OCR visualization")
  764. continue
  765. if isinstance(processed_image, np.ndarray):
  766. image = Image.fromarray(processed_image).convert('RGB')
  767. elif isinstance(processed_image, Image.Image):
  768. image = processed_image.convert('RGB')
  769. else:
  770. continue
  771. draw = ImageDraw.Draw(image)
  772. font = self._get_font(10)
  773. for element in page.get('elements', []):
  774. content = element.get('content', {})
  775. # OCR 文本框
  776. ocr_details = content.get('ocr_details', [])
  777. for ocr_item in ocr_details:
  778. ocr_bbox = ocr_item.get('bbox', [])
  779. if ocr_bbox:
  780. self._draw_polygon(draw, ocr_bbox, self.OCR_BOX_COLOR, width=1)
  781. # 表格单元格
  782. cells = content.get('cells', [])
  783. for cell in cells:
  784. cell_bbox = cell.get('bbox', [])
  785. if cell_bbox and len(cell_bbox) >= 4:
  786. x0, y0, x1, y1 = map(int, cell_bbox[:4])
  787. draw.rectangle([x0, y0, x1, y1], outline=self.CELL_BOX_COLOR, width=2)
  788. cell_text = cell.get('text', '')[:10]
  789. if cell_text:
  790. draw.text((x0 + 2, y0 + 2), cell_text, fill=self.CELL_BOX_COLOR, font=font)
  791. # OCR 框
  792. ocr_boxes = content.get('ocr_boxes', [])
  793. for ocr_box in ocr_boxes:
  794. bbox = ocr_box.get('bbox', [])
  795. if bbox:
  796. self._draw_polygon(draw, bbox, self.OCR_BOX_COLOR, width=1)
  797. ocr_path = output_dir / f"{doc_name}_page_{page_idx + 1}_ocr.png"
  798. image.save(ocr_path)
  799. ocr_paths.append(str(ocr_path))
  800. logger.info(f"🖼️ OCR image saved: {ocr_path}")
  801. return ocr_paths
  802. def _draw_polygon(
  803. self,
  804. draw: ImageDraw.Draw,
  805. bbox: List,
  806. color: tuple,
  807. width: int = 1
  808. ):
  809. """绘制多边形或矩形"""
  810. if isinstance(bbox[0], (list, tuple)):
  811. points = [(int(p[0]), int(p[1])) for p in bbox]
  812. points.append(points[0])
  813. draw.line(points, fill=color, width=width)
  814. elif len(bbox) >= 4:
  815. x0, y0, x1, y1 = map(int, bbox[:4])
  816. draw.rectangle([x0, y0, x1, y1], outline=color, width=width)
  817. def _get_font(self, size: int) -> ImageFont.FreeTypeFont:
  818. """获取字体"""
  819. font_paths = [
  820. "/System/Library/Fonts/Helvetica.ttc",
  821. "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
  822. "/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf",
  823. ]
  824. for font_path in font_paths:
  825. try:
  826. return ImageFont.truetype(font_path, size)
  827. except:
  828. continue
  829. return ImageFont.load_default()
  830. # ==================== 便捷函数 ====================
  831. def save_mineru_format(
  832. results: Dict[str, Any],
  833. output_dir: str,
  834. output_config: Dict[str, Any] = None
  835. ) -> Dict[str, Any]:
  836. """
  837. 便捷函数:保存为 MinerU 格式
  838. Args:
  839. results: pipeline 处理结果
  840. output_dir: 输出目录
  841. output_config: 输出配置
  842. Returns:
  843. 输出文件路径字典
  844. """
  845. if output_config is None:
  846. output_config = {
  847. 'save_json': True,
  848. 'save_page_json': True,
  849. 'save_markdown': True,
  850. 'save_html': True,
  851. 'save_layout_image': False,
  852. 'save_ocr_image': False,
  853. }
  854. formatter = OutputFormatterV2(output_dir)
  855. return formatter.save_results(results, output_config)
  856. if __name__ == "__main__":
  857. # 测试代码
  858. sample_results = {
  859. "document_path": "/path/to/sample.pdf",
  860. "scene": "bank_statement",
  861. "pages": [
  862. {
  863. "page_idx": 0,
  864. "image_shape": [1654, 2338, 3],
  865. "elements": [
  866. {
  867. "type": "title",
  868. "bbox": [100, 50, 800, 100],
  869. "content": {"text": "银行流水"},
  870. "confidence": 0.98
  871. },
  872. {
  873. "type": "table",
  874. "bbox": [100, 200, 800, 600],
  875. "content": {
  876. "html": "<table><tr><td>日期</td><td>金额</td></tr></table>",
  877. "cells": [
  878. {"text": "日期", "bbox": [100, 200, 200, 250], "row": 1, "col": 1},
  879. {"text": "金额", "bbox": [200, 200, 300, 250], "row": 1, "col": 2}
  880. ]
  881. }
  882. }
  883. ]
  884. }
  885. ]
  886. }
  887. output_files = save_mineru_format(
  888. sample_results,
  889. "./test_output_v2",
  890. {
  891. "save_json": True,
  892. "save_page_json": True,
  893. "save_markdown": True,
  894. "save_html": True,
  895. "save_layout_image": False,
  896. "save_ocr_image": False
  897. }
  898. )
  899. print("Generated files:", output_files)