ocr_platform/ocr_utils/output_formatter_v2.py

"""
统一输出格式化器 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
import numpy as np
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 NumpyEncoder(json.JSONEncoder):
    """自定义JSON编码器，处理numpy类型"""
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        return super().default(obj)


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（金额标准化已在 pipeline element_processors 中完成，此处不再重复）
        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, cls=NumpyEncoder)
            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)