output_formatter_v2.py 26 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723
  1. """
  2. 增强版输出格式化器 v2
  3. 支持:
  4. 1. MinerU 标准 middle.json 格式
  5. 2. Markdown 输出
  6. 3. Debug 模式:layout 图片、OCR 图片、单元格坐标图片
  7. 4. 表格 HTML 输出(带坐标信息)
  8. """
  9. import json
  10. import os
  11. import sys
  12. from pathlib import Path
  13. from typing import Dict, Any, List, Optional, Union
  14. from loguru import logger
  15. import numpy as np
  16. from PIL import Image, ImageDraw, ImageFont
  17. import cv2
  18. # 导入 MinerU 组件
  19. mineru_path = Path(__file__).parents[3]
  20. if str(mineru_path) not in sys.path:
  21. sys.path.insert(0, str(mineru_path))
  22. try:
  23. from mineru.backend.vlm.vlm_middle_json_mkcontent import union_make as vlm_union_make
  24. from mineru.utils.enum_class import MakeMode, BlockType, ContentType
  25. MINERU_AVAILABLE = True
  26. except ImportError as e:
  27. logger.warning(f"MinerU components not available: {e}")
  28. MINERU_AVAILABLE = False
  29. class OutputFormatterV2:
  30. """增强版输出格式化器"""
  31. # 颜色映射(与 MinerU 保持一致)
  32. COLOR_MAP = {
  33. 'title': (102, 102, 255), # 蓝色
  34. 'text': (153, 0, 76), # 深红
  35. 'image': (153, 255, 51), # 绿色
  36. 'image_body': (153, 255, 51),
  37. 'image_caption': (102, 178, 255),
  38. 'image_footnote': (255, 178, 102),
  39. 'table': (204, 204, 0), # 黄色
  40. 'table_body': (204, 204, 0),
  41. 'table_caption': (255, 255, 102),
  42. 'table_footnote': (229, 255, 204),
  43. 'interline_equation': (0, 255, 0), # 亮绿
  44. 'inline_equation': (0, 200, 0),
  45. 'list': (40, 169, 92),
  46. 'code': (102, 0, 204), # 紫色
  47. 'header': (128, 128, 128), # 灰色
  48. 'footer': (128, 128, 128),
  49. 'ref_text': (180, 180, 180),
  50. 'ocr_text': (153, 0, 76),
  51. 'error': (255, 0, 0), # 红色
  52. }
  53. # OCR 框颜色
  54. OCR_BOX_COLOR = (0, 255, 0) # 绿色
  55. CELL_BOX_COLOR = (255, 165, 0) # 橙色
  56. def __init__(self, output_dir: str):
  57. """
  58. 初始化格式化器
  59. Args:
  60. output_dir: 输出目录
  61. """
  62. self.output_dir = Path(output_dir)
  63. self.output_dir.mkdir(parents=True, exist_ok=True)
  64. def save_results(
  65. self,
  66. results: Dict[str, Any],
  67. output_config: Dict[str, Any]
  68. ) -> Dict[str, str]:
  69. """
  70. 保存处理结果
  71. Args:
  72. results: 处理结果
  73. output_config: 输出配置
  74. Returns:
  75. 输出文件路径字典
  76. """
  77. output_paths = {}
  78. # 创建文档输出目录
  79. doc_name = Path(results['document_path']).stem
  80. doc_output_dir = self.output_dir / doc_name
  81. doc_output_dir.mkdir(parents=True, exist_ok=True)
  82. # 1. 转换为 MinerU 标准 middle.json
  83. middle_json = self._convert_to_middle_json(results)
  84. # 2. 保存 middle.json
  85. if output_config.get('save_json', True):
  86. json_path = doc_output_dir / f"{doc_name}_middle.json"
  87. with open(json_path, 'w', encoding='utf-8') as f:
  88. json.dump(middle_json, f, ensure_ascii=False, indent=2)
  89. output_paths['middle_json'] = str(json_path)
  90. logger.info(f"📄 Middle JSON saved: {json_path}")
  91. # 3. 保存增强版 JSON(包含单元格坐标)
  92. enhanced_json_path = doc_output_dir / f"{doc_name}_enhanced.json"
  93. with open(enhanced_json_path, 'w', encoding='utf-8') as f:
  94. json.dump(results, f, ensure_ascii=False, indent=2, default=self._json_serializer)
  95. output_paths['enhanced_json'] = str(enhanced_json_path)
  96. logger.info(f"📄 Enhanced JSON saved: {enhanced_json_path}")
  97. # 4. 保存 Markdown
  98. if output_config.get('save_markdown', True):
  99. md_path = self._save_markdown(results, doc_output_dir, doc_name)
  100. output_paths['markdown'] = str(md_path)
  101. # 5. 保存表格 HTML
  102. if output_config.get('save_html', True):
  103. html_dir = self._save_table_htmls(results, doc_output_dir, doc_name)
  104. output_paths['table_htmls'] = str(html_dir)
  105. # 6. Debug 模式:保存可视化图片
  106. if output_config.get('save_layout_image', False):
  107. layout_paths = self._save_layout_images(
  108. results, doc_output_dir, doc_name,
  109. draw_type_label=output_config.get('draw_type_label', True),
  110. draw_bbox_number=output_config.get('draw_bbox_number', True)
  111. )
  112. output_paths['layout_images'] = layout_paths
  113. if output_config.get('save_ocr_image', False):
  114. ocr_paths = self._save_ocr_images(results, doc_output_dir, doc_name)
  115. output_paths['ocr_images'] = ocr_paths
  116. logger.info(f"✅ All results saved to: {doc_output_dir}")
  117. return output_paths
  118. def _json_serializer(self, obj):
  119. """JSON 序列化辅助函数"""
  120. if isinstance(obj, np.ndarray):
  121. return obj.tolist()
  122. if isinstance(obj, np.integer):
  123. return int(obj)
  124. if isinstance(obj, np.floating):
  125. return float(obj)
  126. if isinstance(obj, Path):
  127. return str(obj)
  128. return str(obj)
  129. def _convert_to_middle_json(self, results: Dict[str, Any]) -> Dict[str, Any]:
  130. """转换为 MinerU 标准 middle.json 格式"""
  131. middle_json = {
  132. "pdf_info": [],
  133. "_backend": "vlm",
  134. "_scene": results.get('scene', 'unknown'),
  135. "_version_name": "2.5.0"
  136. }
  137. for page in results.get('pages', []):
  138. page_info = {
  139. 'page_idx': page['page_idx'],
  140. 'page_size': list(page.get('image_shape', [0, 0])[:2][::-1]),
  141. 'angle': page.get('angle', 0),
  142. 'para_blocks': [],
  143. 'discarded_blocks': []
  144. }
  145. for element in page.get('elements', []):
  146. block = self._element_to_block(element, page_info['page_size'])
  147. if block:
  148. elem_type = element.get('type', '')
  149. if elem_type in ['header', 'footer', 'page_number', 'aside_text']:
  150. page_info['discarded_blocks'].append(block)
  151. else:
  152. page_info['para_blocks'].append(block)
  153. middle_json['pdf_info'].append(page_info)
  154. return middle_json
  155. def _element_to_block(
  156. self,
  157. element: Dict[str, Any],
  158. page_size: List[int]
  159. ) -> Dict[str, Any]:
  160. """将元素转换为 MinerU block 格式"""
  161. elem_type = element.get('type', '')
  162. bbox = element.get('bbox', [0, 0, 0, 0])
  163. content = element.get('content', {})
  164. block = {
  165. 'type': elem_type,
  166. 'bbox': bbox,
  167. 'angle': element.get('angle', 0),
  168. 'lines': []
  169. }
  170. # 文本类型
  171. if elem_type in ['text', 'title', 'ref_text', 'header', 'footer', 'ocr_text']:
  172. text = content.get('text', '') if isinstance(content, dict) else str(content)
  173. if text:
  174. block['lines'] = [{
  175. 'bbox': bbox,
  176. 'spans': [{
  177. 'bbox': bbox,
  178. 'type': 'text',
  179. 'content': text
  180. }]
  181. }]
  182. # 表格类型
  183. elif elem_type in ['table', 'table_body']:
  184. table_html = content.get('html', '')
  185. cells = content.get('cells', [])
  186. block['blocks'] = [{
  187. 'type': 'table_body',
  188. 'bbox': bbox,
  189. 'angle': 0,
  190. 'lines': [{
  191. 'bbox': bbox,
  192. 'spans': [{
  193. 'bbox': bbox,
  194. 'type': 'table',
  195. 'html': table_html,
  196. 'cells': cells # 增强:包含单元格坐标
  197. }]
  198. }]
  199. }]
  200. # 公式类型
  201. elif elem_type in ['interline_equation', 'inline_equation', 'equation']:
  202. latex = content.get('latex', '')
  203. block['lines'] = [{
  204. 'bbox': bbox,
  205. 'spans': [{
  206. 'bbox': bbox,
  207. 'type': 'interline_equation' if 'interline' in elem_type else 'inline_equation',
  208. 'content': latex
  209. }]
  210. }]
  211. return block
  212. def _save_markdown(
  213. self,
  214. results: Dict[str, Any],
  215. output_dir: Path,
  216. doc_name: str
  217. ) -> Path:
  218. """保存 Markdown 文件"""
  219. md_path = output_dir / f"{doc_name}.md"
  220. md_lines = [
  221. f"---",
  222. f"scene: {results.get('scene', 'unknown')}",
  223. f"document: {results.get('document_path', '')}",
  224. f"pages: {len(results.get('pages', []))}",
  225. f"---",
  226. "",
  227. ]
  228. for page in results.get('pages', []):
  229. page_idx = page.get('page_idx', 0)
  230. md_lines.append(f"\n## Page {page_idx + 1}\n")
  231. for element in page.get('elements', []):
  232. elem_type = element.get('type', '')
  233. content = element.get('content', {})
  234. bbox = element.get('bbox', [])
  235. # 添加 bbox 注释
  236. bbox_comment = f"<!-- bbox: {bbox} -->"
  237. if elem_type == 'title':
  238. text = content.get('text', '') if isinstance(content, dict) else str(content)
  239. level = element.get('level', 1)
  240. md_lines.append(f"{'#' * min(level, 6)} {text}")
  241. md_lines.append(bbox_comment)
  242. md_lines.append("")
  243. elif elem_type in ['text', 'ocr_text', 'ref_text']:
  244. text = content.get('text', '') if isinstance(content, dict) else str(content)
  245. if text:
  246. md_lines.append(text)
  247. md_lines.append(bbox_comment)
  248. md_lines.append("")
  249. elif elem_type in ['table', 'table_body']:
  250. html = content.get('html', '')
  251. if html:
  252. # 转换 HTML 表格为 Markdown
  253. md_table = self._html_to_markdown_table(html)
  254. md_lines.append(md_table)
  255. md_lines.append(bbox_comment)
  256. # 添加单元格坐标信息
  257. cells = content.get('cells', [])
  258. if cells:
  259. md_lines.append("")
  260. md_lines.append("<details>")
  261. md_lines.append("<summary>单元格坐标信息</summary>")
  262. md_lines.append("")
  263. md_lines.append("```json")
  264. md_lines.append(json.dumps(cells, ensure_ascii=False, indent=2))
  265. md_lines.append("```")
  266. md_lines.append("</details>")
  267. md_lines.append("")
  268. elif elem_type in ['interline_equation', 'inline_equation', 'equation']:
  269. latex = content.get('latex', '')
  270. if latex:
  271. md_lines.append(f"$$\n{latex}\n$$")
  272. md_lines.append(bbox_comment)
  273. md_lines.append("")
  274. with open(md_path, 'w', encoding='utf-8') as f:
  275. f.write('\n'.join(md_lines))
  276. logger.info(f"📝 Markdown saved: {md_path}")
  277. return md_path
  278. def _html_to_markdown_table(self, html: str) -> str:
  279. """将 HTML 表格转换为 Markdown 格式"""
  280. try:
  281. from bs4 import BeautifulSoup
  282. soup = BeautifulSoup(html, 'html.parser')
  283. rows = soup.find_all('tr')
  284. if not rows:
  285. return html
  286. md_rows = []
  287. for row_idx, row in enumerate(rows):
  288. cells = row.find_all(['td', 'th'])
  289. cell_texts = [cell.get_text(strip=True) for cell in cells]
  290. md_row = '| ' + ' | '.join(cell_texts) + ' |'
  291. md_rows.append(md_row)
  292. # 添加表头分隔符
  293. if row_idx == 0:
  294. separator = '| ' + ' | '.join(['---'] * len(cells)) + ' |'
  295. md_rows.append(separator)
  296. return '\n'.join(md_rows)
  297. except Exception as e:
  298. logger.warning(f"HTML to Markdown conversion failed: {e}")
  299. return html
  300. def _save_table_htmls(
  301. self,
  302. results: Dict[str, Any],
  303. output_dir: Path,
  304. doc_name: str
  305. ) -> Path:
  306. """保存表格 HTML 文件(带坐标信息)"""
  307. tables_dir = output_dir / 'tables'
  308. tables_dir.mkdir(exist_ok=True)
  309. table_count = 0
  310. for page in results.get('pages', []):
  311. page_idx = page.get('page_idx', 0)
  312. for element in page.get('elements', []):
  313. if element.get('type') in ['table', 'table_body']:
  314. table_count += 1
  315. content = element.get('content', {})
  316. html = content.get('html', '')
  317. cells = content.get('cells', [])
  318. if html:
  319. # 生成带样式的 HTML
  320. full_html = self._generate_table_html_with_styles(
  321. html, cells, doc_name, page_idx, table_count
  322. )
  323. html_path = tables_dir / f"{doc_name}_table_{table_count}_page_{page_idx + 1}.html"
  324. with open(html_path, 'w', encoding='utf-8') as f:
  325. f.write(full_html)
  326. if table_count > 0:
  327. logger.info(f"📊 {table_count} tables saved to: {tables_dir}")
  328. return tables_dir
  329. def _generate_table_html_with_styles(
  330. self,
  331. table_html: str,
  332. cells: List[Dict],
  333. doc_name: str,
  334. page_idx: int,
  335. table_idx: int
  336. ) -> str:
  337. """生成带样式和坐标信息的完整 HTML"""
  338. cells_json = json.dumps(cells, ensure_ascii=False, indent=2) if cells else "[]"
  339. return f"""<!DOCTYPE html>
  340. <html lang="zh-CN">
  341. <head>
  342. <meta charset="UTF-8">
  343. <meta name="viewport" content="width=device-width, initial-scale=1.0">
  344. <title>{doc_name} - Table {table_idx}</title>
  345. <style>
  346. body {{
  347. font-family: Arial, "Microsoft YaHei", sans-serif;
  348. margin: 20px;
  349. background-color: #f5f5f5;
  350. }}
  351. .container {{
  352. max-width: 1400px;
  353. margin: 0 auto;
  354. background-color: white;
  355. padding: 20px;
  356. box-shadow: 0 0 10px rgba(0,0,0,0.1);
  357. border-radius: 8px;
  358. }}
  359. .meta {{
  360. color: #666;
  361. font-size: 0.9em;
  362. margin-bottom: 20px;
  363. padding-bottom: 10px;
  364. border-bottom: 1px solid #ddd;
  365. }}
  366. table {{
  367. border-collapse: collapse;
  368. width: 100%;
  369. margin: 20px 0;
  370. }}
  371. th, td {{
  372. border: 1px solid #ddd;
  373. padding: 8px 12px;
  374. text-align: left;
  375. }}
  376. th {{
  377. background-color: #f2f2f2;
  378. font-weight: bold;
  379. }}
  380. tr:hover {{
  381. background-color: #f9f9f9;
  382. }}
  383. td[data-bbox], th[data-bbox] {{
  384. position: relative;
  385. }}
  386. td[data-bbox]:hover::after, th[data-bbox]:hover::after {{
  387. content: attr(data-bbox);
  388. position: absolute;
  389. bottom: 100%;
  390. left: 0;
  391. background: #333;
  392. color: white;
  393. padding: 2px 6px;
  394. font-size: 10px;
  395. border-radius: 3px;
  396. white-space: nowrap;
  397. z-index: 100;
  398. }}
  399. .cells-info {{
  400. margin-top: 30px;
  401. padding: 15px;
  402. background-color: #f8f9fa;
  403. border-radius: 5px;
  404. }}
  405. .cells-info summary {{
  406. cursor: pointer;
  407. font-weight: bold;
  408. color: #333;
  409. }}
  410. .cells-info pre {{
  411. background-color: #2d2d2d;
  412. color: #f8f8f2;
  413. padding: 15px;
  414. border-radius: 5px;
  415. overflow-x: auto;
  416. font-size: 12px;
  417. }}
  418. </style>
  419. </head>
  420. <body>
  421. <div class="container">
  422. <div class="meta">
  423. <p><strong>Document:</strong> {doc_name}</p>
  424. <p><strong>Page:</strong> {page_idx + 1}</p>
  425. <p><strong>Table:</strong> {table_idx}</p>
  426. <p><strong>Cells with coordinates:</strong> {len(cells)}</p>
  427. </div>
  428. {table_html}
  429. <div class="cells-info">
  430. <details>
  431. <summary>📍 单元格坐标数据 (JSON)</summary>
  432. <pre>{cells_json}</pre>
  433. </details>
  434. </div>
  435. </div>
  436. </body>
  437. </html>"""
  438. def _save_layout_images(
  439. self,
  440. results: Dict[str, Any],
  441. output_dir: Path,
  442. doc_name: str,
  443. draw_type_label: bool = True,
  444. draw_bbox_number: bool = True
  445. ) -> List[str]:
  446. """保存 Layout 可视化图片"""
  447. layout_paths = []
  448. for page in results.get('pages', []):
  449. page_idx = page.get('page_idx', 0)
  450. processed_image = page.get('processed_image')
  451. if processed_image is None:
  452. continue
  453. # 转换为 PIL Image
  454. if isinstance(processed_image, np.ndarray):
  455. image = Image.fromarray(processed_image).convert('RGB')
  456. elif isinstance(processed_image, Image.Image):
  457. image = processed_image.convert('RGB')
  458. else:
  459. continue
  460. # 创建绘图对象
  461. draw = ImageDraw.Draw(image, 'RGBA')
  462. # 加载字体
  463. font = self._get_font(14)
  464. # 绘制每个元素
  465. for idx, element in enumerate(page.get('elements', []), 1):
  466. elem_type = element.get('type', '')
  467. bbox = element.get('bbox', [0, 0, 0, 0])
  468. if len(bbox) < 4:
  469. continue
  470. x0, y0, x1, y1 = map(int, bbox[:4])
  471. color = self.COLOR_MAP.get(elem_type, (255, 0, 0))
  472. # 绘制半透明填充
  473. overlay = Image.new('RGBA', image.size, (255, 255, 255, 0))
  474. overlay_draw = ImageDraw.Draw(overlay)
  475. overlay_draw.rectangle([x0, y0, x1, y1], fill=(*color, 50))
  476. image = Image.alpha_composite(image.convert('RGBA'), overlay).convert('RGB')
  477. draw = ImageDraw.Draw(image)
  478. # 绘制边框
  479. draw.rectangle([x0, y0, x1, y1], outline=color, width=2)
  480. # 标注类型
  481. if draw_type_label:
  482. label = elem_type.replace('_', ' ').title()
  483. bbox_label = draw.textbbox((x0 + 2, y0 + 2), label, font=font)
  484. draw.rectangle(bbox_label, fill=color)
  485. draw.text((x0 + 2, y0 + 2), label, fill='white', font=font)
  486. # 标注序号
  487. if draw_bbox_number:
  488. number_text = str(idx)
  489. bbox_number = draw.textbbox((x1 - 25, y0 + 2), number_text, font=font)
  490. draw.rectangle(bbox_number, fill=(255, 0, 0))
  491. draw.text((x1 - 25, y0 + 2), number_text, fill='white', font=font)
  492. # 保存图片
  493. layout_path = output_dir / f"{doc_name}_page_{page_idx + 1}_layout.png"
  494. image.save(layout_path)
  495. layout_paths.append(str(layout_path))
  496. logger.info(f"🖼️ Layout image saved: {layout_path}")
  497. return layout_paths
  498. def _save_ocr_images(
  499. self,
  500. results: Dict[str, Any],
  501. output_dir: Path,
  502. doc_name: str
  503. ) -> List[str]:
  504. """保存 OCR 可视化图片(显示文本框和单元格坐标)"""
  505. ocr_paths = []
  506. for page in results.get('pages', []):
  507. page_idx = page.get('page_idx', 0)
  508. processed_image = page.get('processed_image')
  509. if processed_image is None:
  510. continue
  511. # 转换为 PIL Image
  512. if isinstance(processed_image, np.ndarray):
  513. image = Image.fromarray(processed_image).convert('RGB')
  514. elif isinstance(processed_image, Image.Image):
  515. image = processed_image.convert('RGB')
  516. else:
  517. continue
  518. draw = ImageDraw.Draw(image)
  519. font = self._get_font(10)
  520. # 遍历所有元素
  521. for element in page.get('elements', []):
  522. content = element.get('content', {})
  523. # 绘制 OCR 文本框
  524. ocr_details = content.get('ocr_details', [])
  525. for ocr_item in ocr_details:
  526. ocr_bbox = ocr_item.get('bbox', [])
  527. if ocr_bbox:
  528. self._draw_polygon(draw, ocr_bbox, self.OCR_BOX_COLOR, width=1)
  529. # 绘制表格单元格坐标
  530. cells = content.get('cells', [])
  531. for cell in cells:
  532. cell_bbox = cell.get('bbox', [])
  533. if cell_bbox and len(cell_bbox) >= 4:
  534. x0, y0, x1, y1 = map(int, cell_bbox[:4])
  535. draw.rectangle([x0, y0, x1, y1], outline=self.CELL_BOX_COLOR, width=2)
  536. # 标注单元格文本
  537. cell_text = cell.get('text', '')[:10]
  538. if cell_text:
  539. draw.text((x0 + 2, y0 + 2), cell_text, fill=self.CELL_BOX_COLOR, font=font)
  540. # 绘制表格 OCR 框
  541. ocr_boxes = content.get('ocr_boxes', [])
  542. for ocr_box in ocr_boxes:
  543. bbox = ocr_box.get('bbox', [])
  544. if bbox:
  545. self._draw_polygon(draw, bbox, self.OCR_BOX_COLOR, width=1)
  546. # 保存图片
  547. ocr_path = output_dir / f"{doc_name}_page_{page_idx + 1}_ocr.png"
  548. image.save(ocr_path)
  549. ocr_paths.append(str(ocr_path))
  550. logger.info(f"🖼️ OCR image saved: {ocr_path}")
  551. return ocr_paths
  552. def _draw_polygon(
  553. self,
  554. draw: ImageDraw.Draw,
  555. bbox: List,
  556. color: tuple,
  557. width: int = 1
  558. ):
  559. """绘制多边形或矩形"""
  560. # 4点坐标格式 [[x1,y1], [x2,y2], [x3,y3], [x4,y4]]
  561. if isinstance(bbox[0], (list, tuple)):
  562. points = [(int(p[0]), int(p[1])) for p in bbox]
  563. points.append(points[0]) # 闭合多边形
  564. draw.line(points, fill=color, width=width)
  565. # 4值坐标格式 [x1, y1, x2, y2]
  566. elif len(bbox) >= 4:
  567. x0, y0, x1, y1 = map(int, bbox[:4])
  568. draw.rectangle([x0, y0, x1, y1], outline=color, width=width)
  569. def _get_font(self, size: int) -> ImageFont.FreeTypeFont:
  570. """获取字体"""
  571. font_paths = [
  572. "/System/Library/Fonts/Helvetica.ttc",
  573. "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
  574. "/usr/share/fonts/truetype/liberation/LiberationSans-Regular.ttf",
  575. ]
  576. for font_path in font_paths:
  577. try:
  578. return ImageFont.truetype(font_path, size)
  579. except:
  580. continue
  581. return ImageFont.load_default()
  582. if __name__ == "__main__":
  583. # 测试代码
  584. sample_results = {
  585. "document_path": "/path/to/sample.pdf",
  586. "scene": "bank_statement",
  587. "pages": [
  588. {
  589. "page_idx": 0,
  590. "image_shape": [1654, 2338, 3],
  591. "elements": [
  592. {
  593. "type": "title",
  594. "bbox": [100, 50, 800, 100],
  595. "content": {"text": "银行流水"},
  596. "confidence": 0.98
  597. },
  598. {
  599. "type": "table",
  600. "bbox": [100, 200, 800, 600],
  601. "content": {
  602. "html": "<table><tr><td>日期</td><td>金额</td></tr></table>",
  603. "cells": [
  604. {"text": "日期", "bbox": [100, 200, 200, 250], "row": 1, "col": 1},
  605. {"text": "金额", "bbox": [200, 200, 300, 250], "row": 1, "col": 2}
  606. ]
  607. }
  608. }
  609. ]
  610. }
  611. ]
  612. }
  613. formatter = OutputFormatterV2("./test_output_v2")
  614. output_files = formatter.save_results(
  615. sample_results,
  616. {
  617. "save_json": True,
  618. "save_markdown": True,
  619. "save_html": True,
  620. "save_layout_image": False,
  621. "save_ocr_image": False
  622. }
  623. )
  624. print("Generated files:", output_files)