""" 统一输出格式化器 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 输出(带坐标信息) 6. 金额数字标准化(全角→半角转换) 模块结构: - json_formatters.py: JSON 格式化工具 - markdown_generator.py: Markdown 生成器 - html_generator.py: HTML 生成器 - visualization_utils.py: 可视化工具 """ import json import sys 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 # 导入数字标准化工具 from .normalize_financial_numbers import normalize_markdown_table, normalize_json_table class OutputFormatterV2: """ 统一输出格式化器 严格遵循 MinerU mineru_vllm_results_cell_bbox 格式: - middle.json: MinerU 标准格式,用于生成 Markdown - page_xxx.json: 每页独立的 JSON,包含 table_cells - Markdown: 带 bbox 注释 - 表格: HTML 格式,带 data-bbox 属性 命名规则: - PDF输入: 文件名_page_001.*(按页编号) - 图片输入: 文件名.*(不加页码后缀) """ # 颜色映射(导出供其他模块使用) 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) @staticmethod def is_pdf_input(results: Dict[str, Any]) -> bool: """ 判断输入是否为 PDF Args: results: 处理结果 Returns: True 如果输入是 PDF,否则 False """ doc_path = results.get('document_path', '') if doc_path: return Path(doc_path).suffix.lower() == '.pdf' # 如果没有 document_path,检查 metadata input_type = results.get('metadata', {}).get('input_type', '') return input_type == 'pdf' @staticmethod def get_page_name(doc_name: str, page_idx: int, is_pdf: bool, total_pages: int = 1) -> str: """ 获取页面名称 Args: doc_name: 文档名称 page_idx: 页码索引(从0开始) is_pdf: 是否为 PDF 输入 total_pages: 总页数 Returns: 页面名称(不含扩展名) """ if is_pdf or total_pages > 1: # PDF 或多页输入:添加页码后缀 return f"{doc_name}_page_{page_idx + 1:03d}" else: # 单个图片:不添加页码后缀 return doc_name def save_results( self, results: Dict[str, Any], output_config: Dict[str, Any] ) -> Dict[str, Any]: """ 保存处理结果 命名规则: - PDF输入: 文件名_page_001.*(按页编号) - 图片输入: 文件名.*(不加页码后缀) Args: results: 处理结果 output_config: 输出配置,支持以下选项: - create_subdir: 是否在输出目录下创建文档名子目录(默认 False) - ... 其他选项见 save_mineru_format 函数 Returns: 输出文件路径字典 """ output_paths: Dict[str, Any] = { 'images': [], 'json_pages': [], } # 创建文档输出目录 doc_name = Path(results['document_path']).stem # 是否创建子目录(默认不创建,直接使用指定的输出目录) create_subdir = output_config.get('create_subdir', False) if create_subdir: doc_output_dir = self.output_dir / doc_name else: doc_output_dir = self.output_dir doc_output_dir.mkdir(parents=True, exist_ok=True) # 判断输入类型 is_pdf = self.is_pdf_input(results) total_pages = len(results.get('pages', [])) # 创建 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, is_pdf=is_pdf ) 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" json_content = json.dumps(middle_json, ensure_ascii=False, indent=2) # 金额数字标准化 normalize_numbers = output_config.get('normalize_numbers', True) if normalize_numbers: original_content = json_content json_content = normalize_json_table(json_content) # 检查是否有变化 if json_content != original_content: # 保存原始文件 original_path = doc_output_dir / f"{doc_name}_middle_original.json" with open(original_path, 'w', encoding='utf-8') as f: f.write(original_content) logger.info(f"📄 Original middle JSON saved: {original_path}") output_paths['middle_json_original'] = str(original_path) with open(json_path, 'w', encoding='utf-8') as f: f.write(json_content) 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): normalize_numbers = output_config.get('normalize_numbers', True) page_json_paths = JSONFormatters.save_page_jsons( results, doc_output_dir, doc_name, is_pdf=is_pdf, normalize_numbers=normalize_numbers ) output_paths['json_pages'] = page_json_paths # 5. 保存 Markdown(完整版) if output_config.get('save_markdown', True): normalize_numbers = output_config.get('normalize_numbers', True) md_path, original_md_path = MarkdownGenerator.save_markdown( results, middle_json, doc_output_dir, doc_name, normalize_numbers=normalize_numbers ) output_paths['markdown'] = str(md_path) if original_md_path: output_paths['markdown_original'] = str(original_md_path) # 5.5 保存每页独立的 Markdown if output_config.get('save_page_markdown', True): normalize_numbers = output_config.get('normalize_numbers', True) page_md_paths = MarkdownGenerator.save_page_markdowns( results, doc_output_dir, doc_name, is_pdf=is_pdf, normalize_numbers=normalize_numbers ) 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, is_pdf=is_pdf ) 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), is_pdf=is_pdf ) 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, is_pdf=is_pdf ) 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: 输出配置,支持以下选项: - create_subdir: 在输出目录下创建文档名子目录(默认 False) - save_json: 保存 middle.json - save_page_json: 保存每页 JSON - save_markdown: 保存完整 Markdown - save_page_markdown: 保存每页 Markdown - save_html: 保存表格 HTML - save_layout_image: 保存布局可视化图 - save_ocr_image: 保存 OCR 可视化图 - normalize_numbers: 标准化金额数字(全角→半角) Returns: 输出文件路径字典 """ if output_config is None: output_config = { 'create_subdir': False, # 默认不创建子目录,直接使用指定目录 '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, 'normalize_numbers': True, # 默认启用数字标准化 } 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": "
日期金额
", "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)