"""
增强版输出格式化器 v2
支持:
1. MinerU 标准 middle.json 格式
2. Markdown 输出
3. Debug 模式:layout 图片、OCR 图片、单元格坐标图片
4. 表格 HTML 输出(带坐标信息)
"""
import json
import os
import sys
from pathlib import Path
from typing import Dict, Any, List, Optional, Union
from loguru import logger
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import cv2
# 导入 MinerU 组件
mineru_path = Path(__file__).parents[3]
if str(mineru_path) not in sys.path:
sys.path.insert(0, str(mineru_path))
try:
from mineru.backend.vlm.vlm_middle_json_mkcontent import union_make as vlm_union_make
from mineru.utils.enum_class import MakeMode, BlockType, ContentType
MINERU_AVAILABLE = True
except ImportError as e:
logger.warning(f"MinerU components not available: {e}")
MINERU_AVAILABLE = False
class OutputFormatterV2:
"""增强版输出格式化器"""
# 颜色映射(与 MinerU 保持一致)
COLOR_MAP = {
'title': (102, 102, 255), # 蓝色
'text': (153, 0, 76), # 深红
'image': (153, 255, 51), # 绿色
'image_body': (153, 255, 51),
'image_caption': (102, 178, 255),
'image_footnote': (255, 178, 102),
'table': (204, 204, 0), # 黄色
'table_body': (204, 204, 0),
'table_caption': (255, 255, 102),
'table_footnote': (229, 255, 204),
'interline_equation': (0, 255, 0), # 亮绿
'inline_equation': (0, 200, 0),
'list': (40, 169, 92),
'code': (102, 0, 204), # 紫色
'header': (128, 128, 128), # 灰色
'footer': (128, 128, 128),
'ref_text': (180, 180, 180),
'ocr_text': (153, 0, 76),
'error': (255, 0, 0), # 红色
}
# OCR 框颜色
OCR_BOX_COLOR = (0, 255, 0) # 绿色
CELL_BOX_COLOR = (255, 165, 0) # 橙色
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, str]:
"""
保存处理结果
Args:
results: 处理结果
output_config: 输出配置
Returns:
输出文件路径字典
"""
output_paths = {}
# 创建文档输出目录
doc_name = Path(results['document_path']).stem
doc_output_dir = self.output_dir / doc_name
doc_output_dir.mkdir(parents=True, exist_ok=True)
# 1. 转换为 MinerU 标准 middle.json
middle_json = self._convert_to_middle_json(results)
# 2. 保存 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}")
# 3. 保存增强版 JSON(包含单元格坐标)
enhanced_json_path = doc_output_dir / f"{doc_name}_enhanced.json"
with open(enhanced_json_path, 'w', encoding='utf-8') as f:
json.dump(results, f, ensure_ascii=False, indent=2, default=self._json_serializer)
output_paths['enhanced_json'] = str(enhanced_json_path)
logger.info(f"📄 Enhanced JSON saved: {enhanced_json_path}")
# 4. 保存 Markdown
if output_config.get('save_markdown', True):
md_path = self._save_markdown(results, doc_output_dir, doc_name)
output_paths['markdown'] = str(md_path)
# 5. 保存表格 HTML
if output_config.get('save_html', True):
html_dir = self._save_table_htmls(results, doc_output_dir, doc_name)
output_paths['table_htmls'] = str(html_dir)
# 6. Debug 模式:保存可视化图片
if output_config.get('save_layout_image', False):
layout_paths = self._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 = self._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 _json_serializer(self, obj):
"""JSON 序列化辅助函数"""
if isinstance(obj, np.ndarray):
return obj.tolist()
if isinstance(obj, np.integer):
return int(obj)
if isinstance(obj, np.floating):
return float(obj)
if isinstance(obj, Path):
return str(obj)
return str(obj)
def _convert_to_middle_json(self, results: Dict[str, Any]) -> Dict[str, Any]:
"""转换为 MinerU 标准 middle.json 格式"""
middle_json = {
"pdf_info": [],
"_backend": "vlm",
"_scene": results.get('scene', 'unknown'),
"_version_name": "2.5.0"
}
for page in results.get('pages', []):
page_info = {
'page_idx': page['page_idx'],
'page_size': list(page.get('image_shape', [0, 0])[:2][::-1]),
'angle': page.get('angle', 0),
'para_blocks': [],
'discarded_blocks': []
}
for element in page.get('elements', []):
block = self._element_to_block(element, page_info['page_size'])
if block:
elem_type = element.get('type', '')
if elem_type in ['header', 'footer', 'page_number', 'aside_text']:
page_info['discarded_blocks'].append(block)
else:
page_info['para_blocks'].append(block)
middle_json['pdf_info'].append(page_info)
return middle_json
def _element_to_block(
self,
element: Dict[str, Any],
page_size: List[int]
) -> Dict[str, Any]:
"""将元素转换为 MinerU block 格式"""
elem_type = element.get('type', '')
bbox = element.get('bbox', [0, 0, 0, 0])
content = element.get('content', {})
block = {
'type': elem_type,
'bbox': bbox,
'angle': element.get('angle', 0),
'lines': []
}
# 文本类型
if elem_type in ['text', 'title', 'ref_text', 'header', 'footer', 'ocr_text']:
text = content.get('text', '') if isinstance(content, dict) else str(content)
if text:
block['lines'] = [{
'bbox': bbox,
'spans': [{
'bbox': bbox,
'type': 'text',
'content': text
}]
}]
# 表格类型
elif elem_type in ['table', 'table_body']:
table_html = content.get('html', '')
cells = content.get('cells', [])
block['blocks'] = [{
'type': 'table_body',
'bbox': bbox,
'angle': 0,
'lines': [{
'bbox': bbox,
'spans': [{
'bbox': bbox,
'type': 'table',
'html': table_html,
'cells': cells # 增强:包含单元格坐标
}]
}]
}]
# 公式类型
elif elem_type in ['interline_equation', 'inline_equation', 'equation']:
latex = content.get('latex', '')
block['lines'] = [{
'bbox': bbox,
'spans': [{
'bbox': bbox,
'type': 'interline_equation' if 'interline' in elem_type else 'inline_equation',
'content': latex
}]
}]
return block
def _save_markdown(
self,
results: Dict[str, Any],
output_dir: Path,
doc_name: str
) -> Path:
"""保存 Markdown 文件"""
md_path = output_dir / f"{doc_name}.md"
md_lines = [
f"---",
f"scene: {results.get('scene', 'unknown')}",
f"document: {results.get('document_path', '')}",
f"pages: {len(results.get('pages', []))}",
f"---",
"",
]
for page in results.get('pages', []):
page_idx = page.get('page_idx', 0)
md_lines.append(f"\n## Page {page_idx + 1}\n")
for element in page.get('elements', []):
elem_type = element.get('type', '')
content = element.get('content', {})
bbox = element.get('bbox', [])
# 添加 bbox 注释
bbox_comment = f""
if elem_type == 'title':
text = content.get('text', '') if isinstance(content, dict) else str(content)
level = element.get('level', 1)
md_lines.append(f"{'#' * min(level, 6)} {text}")
md_lines.append(bbox_comment)
md_lines.append("")
elif elem_type in ['text', 'ocr_text', 'ref_text']:
text = content.get('text', '') if isinstance(content, dict) else str(content)
if text:
md_lines.append(text)
md_lines.append(bbox_comment)
md_lines.append("")
elif elem_type in ['table', 'table_body']:
html = content.get('html', '')
if html:
# 转换 HTML 表格为 Markdown
md_table = self._html_to_markdown_table(html)
md_lines.append(md_table)
md_lines.append(bbox_comment)
# 添加单元格坐标信息
cells = content.get('cells', [])
if cells:
md_lines.append("")
md_lines.append("单元格坐标信息
")
md_lines.append("")
md_lines.append("```json")
md_lines.append(json.dumps(cells, ensure_ascii=False, indent=2))
md_lines.append("```")
md_lines.append("
{cells_json}
| 日期 | 金额 |