| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218 |
- """
- 统一输出格式化器 v2
- 严格遵循 MinerU mineru_vllm_results_cell_bbox 格式
- 支持:
- 1. MinerU 标准 middle.json 格式(用于 union_make 生成 Markdown)
- 2. mineru_vllm_results_cell_bbox 格式(每页独立 JSON)
- 3. Markdown 输出(复用 MinerU union_make)
- 4. Debug 模式:layout 图片、OCR 图片
- 5. 表格 HTML 输出(带坐标信息)
- 模块结构:
- - json_formatters.py: JSON 格式化工具
- - markdown_generator.py: Markdown 生成器
- - html_generator.py: HTML 生成器
- - visualization_utils.py: 可视化工具
- """
- import json
- from pathlib import Path
- from typing import Dict, Any, List, Optional
- from loguru import logger
- # 导入子模块
- from .json_formatters import JSONFormatters
- from .markdown_generator import MarkdownGenerator
- from .html_generator import HTMLGenerator
- from .visualization_utils import VisualizationUtils
- class OutputFormatterV2:
- """
- 统一输出格式化器
-
- 严格遵循 MinerU mineru_vllm_results_cell_bbox 格式:
- - middle.json: MinerU 标准格式,用于生成 Markdown
- - page_xxx.json: 每页独立的 JSON,包含 table_cells
- - Markdown: 带 bbox 注释
- - 表格: HTML 格式,带 data-bbox 属性
- """
-
- # 颜色映射(导出供其他模块使用)
- COLOR_MAP = VisualizationUtils.COLOR_MAP
- OCR_BOX_COLOR = VisualizationUtils.OCR_BOX_COLOR
- CELL_BOX_COLOR = VisualizationUtils.CELL_BOX_COLOR
-
- def __init__(self, output_dir: str):
- """
- 初始化格式化器
-
- Args:
- output_dir: 输出目录
- """
- self.output_dir = Path(output_dir)
- self.output_dir.mkdir(parents=True, exist_ok=True)
-
- def save_results(
- self,
- results: Dict[str, Any],
- output_config: Dict[str, Any]
- ) -> Dict[str, Any]:
- """
- 保存处理结果
-
- Args:
- results: 处理结果
- output_config: 输出配置
-
- Returns:
- 输出文件路径字典
- """
- output_paths: Dict[str, Any] = {
- 'images': [],
- 'json_pages': [],
- }
-
- # 创建文档输出目录
- doc_name = Path(results['document_path']).stem
- doc_output_dir = self.output_dir / doc_name
- doc_output_dir.mkdir(parents=True, exist_ok=True)
-
- # 创建 images 子目录
- images_dir = doc_output_dir / 'images'
- images_dir.mkdir(exist_ok=True)
-
- # 1. 首先保存图片元素(设置 image_path)
- image_paths = VisualizationUtils.save_image_elements(results, images_dir, doc_name)
- if image_paths:
- output_paths['images'] = image_paths
-
- # 2. 转换为 MinerU middle.json 格式
- middle_json = JSONFormatters.convert_to_middle_json(results)
-
- # 3. 保存 middle.json
- if output_config.get('save_json', True):
- json_path = doc_output_dir / f"{doc_name}_middle.json"
- with open(json_path, 'w', encoding='utf-8') as f:
- json.dump(middle_json, f, ensure_ascii=False, indent=2)
- output_paths['middle_json'] = str(json_path)
- logger.info(f"📄 Middle JSON saved: {json_path}")
-
- # 4. 保存每页独立的 mineru_vllm_results_cell_bbox 格式 JSON
- if output_config.get('save_page_json', True):
- page_json_paths = JSONFormatters.save_page_jsons(results, doc_output_dir, doc_name)
- output_paths['json_pages'] = page_json_paths
-
- # 5. 保存 Markdown(完整版)
- if output_config.get('save_markdown', True):
- md_path = MarkdownGenerator.save_markdown(results, middle_json, doc_output_dir, doc_name)
- output_paths['markdown'] = str(md_path)
-
- # 5.5 保存每页独立的 Markdown
- if output_config.get('save_page_markdown', True):
- page_md_paths = MarkdownGenerator.save_page_markdowns(results, doc_output_dir, doc_name)
- output_paths['markdown_pages'] = page_md_paths
-
- # 6. 保存表格 HTML
- if output_config.get('save_html', True):
- html_dir = HTMLGenerator.save_table_htmls(results, doc_output_dir, doc_name)
- output_paths['table_htmls'] = str(html_dir)
-
- # 7. Debug 模式:保存可视化图片
- if output_config.get('save_layout_image', False):
- layout_paths = VisualizationUtils.save_layout_images(
- results, doc_output_dir, doc_name,
- draw_type_label=output_config.get('draw_type_label', True),
- draw_bbox_number=output_config.get('draw_bbox_number', True)
- )
- output_paths['layout_images'] = layout_paths
-
- if output_config.get('save_ocr_image', False):
- ocr_paths = VisualizationUtils.save_ocr_images(results, doc_output_dir, doc_name)
- output_paths['ocr_images'] = ocr_paths
-
- logger.info(f"✅ All results saved to: {doc_output_dir}")
- return output_paths
- # ==================== 便捷函数 ====================
- def save_mineru_format(
- results: Dict[str, Any],
- output_dir: str,
- output_config: Optional[Dict[str, Any]] = None
- ) -> Dict[str, Any]:
- """
- 便捷函数:保存为 MinerU 格式
-
- Args:
- results: pipeline 处理结果
- output_dir: 输出目录
- output_config: 输出配置
-
- Returns:
- 输出文件路径字典
- """
- if output_config is None:
- output_config = {
- 'save_json': True,
- 'save_page_json': True,
- 'save_markdown': True,
- 'save_page_markdown': True,
- 'save_html': True,
- 'save_layout_image': False,
- 'save_ocr_image': False,
- }
-
- formatter = OutputFormatterV2(output_dir)
- return formatter.save_results(results, output_config)
- if __name__ == "__main__":
- # 测试代码
- sample_results = {
- "document_path": "/path/to/sample.pdf",
- "scene": "bank_statement",
- "pages": [
- {
- "page_idx": 0,
- "image_shape": [1654, 2338, 3],
- "elements": [
- {
- "type": "title",
- "bbox": [100, 50, 800, 100],
- "content": {"text": "银行流水"},
- "confidence": 0.98
- },
- {
- "type": "table",
- "bbox": [100, 200, 800, 600],
- "content": {
- "html": "<table><tr><td>日期</td><td>金额</td></tr></table>",
- "cells": [
- {"text": "日期", "bbox": [100, 200, 200, 250], "row": 1, "col": 1},
- {"text": "金额", "bbox": [200, 200, 300, 250], "row": 1, "col": 2}
- ]
- }
- }
- ]
- }
- ]
- }
-
- output_files = save_mineru_format(
- sample_results,
- "./test_output_v2",
- {
- "save_json": True,
- "save_page_json": True,
- "save_markdown": True,
- "save_page_markdown": True,
- "save_html": True,
- "save_layout_image": False,
- "save_ocr_image": False
- }
- )
-
- print("Generated files:", output_files)
|