ocr_platform/ocr_utils/json_formatters.py

"""
JSON 格式化工具模块

提供 JSON 输出格式化功能：
- MinerU middle.json 格式转换
- mineru_vllm_results_cell_bbox 格式转换
- 表格单元格格式化
- 金额数字标准化（全角→半角）
"""
import json
import sys
import numpy as np
from pathlib import Path
from typing import Dict, Any, List, Optional
from loguru import logger

# 导入数字标准化工具
from .normalize_financial_numbers import normalize_json_table


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 JSONFormatters:
    """JSON 格式化工具类"""
    
    @staticmethod
    def convert_to_middle_json(results: Dict[str, Any]) -> Dict[str, Any]:
        """
        转换为 MinerU 标准 middle.json 格式
        
        用于 vlm_union_make 生成 Markdown
        
        Args:
            results: 处理结果
            
        Returns:
            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 = JSONFormatters._element_to_middle_block(element)
                if block:
                    elem_type = element.get('type', '')
                    if elem_type in ['header', 'footer', 'page_number', 'aside_text', 'abandon', 'discarded']:
                        page_info['discarded_blocks'].append(block)
                    else:
                        page_info['para_blocks'].append(block)
            
            # 处理丢弃元素（从 discarded_blocks 字段）
            for element in page.get('discarded_blocks', []):
                block = JSONFormatters._element_to_middle_block(element)
                if block:
                    page_info['discarded_blocks'].append(block)
            
            middle_json['pdf_info'].append(page_info)
        
        return middle_json
    
    @staticmethod
    def _element_to_middle_block(element: Dict[str, Any]) -> Optional[Dict[str, Any]]:
        """
        将元素转换为 MinerU middle.json block 格式
        
        MinerU 期望的嵌套结构:
        - image 类型: { type: "image", blocks: [{ type: "image_body", lines: [...] }] }
        - table 类型: { type: "table", blocks: [{ type: "table_body", lines: [...] }] }
        """
        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),
            'reading_order': element.get('reading_order', 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['type'] = 'table'
            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 ['image', 'image_body', 'figure']:
            block['type'] = 'image'
            block['blocks'] = [{
                'type': 'image_body',
                'bbox': bbox,
                'angle': element.get('angle', 0),
                'lines': [{
                    'bbox': bbox,
                    'spans': [{
                        'bbox': bbox,
                        'type': 'image',
                        'image_path': content.get('image_path', ''),
                        'description': content.get('description', '')
                    }]
                }]
            }]
        
        # 公式类型
        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
                }]
            }]
        
        # 表格/图片附属文本
        elif elem_type in ['table_caption', 'table_footnote', 'image_caption', 'image_footnote']:
            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 ['abandon', 'discarded']:
            block['type'] = 'abandon'
            text = content.get('text', '') if isinstance(content, dict) else str(content)
            if text:
                block['lines'] = [{
                    'bbox': bbox,
                    'spans': [{
                        'bbox': bbox,
                        'type': 'text',
                        'content': text
                    }]
                }]
        
        return block
    
    @staticmethod
    def save_page_jsons(
        results: Dict[str, Any],
        output_dir: Path,
        doc_name: str,
        is_pdf: bool = True,
        normalize_numbers: bool = True
    ) -> List[str]:
        """
        保存每页独立的 JSON（mineru_vllm_results_cell_bbox 格式）
        
        命名规则:
        - PDF输入: 文件名_page_001.json
        - 图片输入（单页）: 文件名.json
        
        Args:
            results: 处理结果
            output_dir: 输出目录
            doc_name: 文档名称
            is_pdf: 是否为 PDF 输入
            normalize_numbers: 是否标准化金额数字（全角→半角）
            
        Returns:
            保存的文件路径列表
        """
        saved_paths = []
        total_pages = len(results.get('pages', []))
        
        for page in results.get('pages', []):
            page_idx = page.get('page_idx', 0)
            
            # 根据输入类型决定命名
            if is_pdf or total_pages > 1:
                page_name = f"{doc_name}_page_{page_idx + 1:03d}"
            else:
                page_name = doc_name
            
            # 获取页面旋转角度
            page_rotation_angle = float(page.get('angle', 0))
            
            # 转换为 mineru_vllm_results_cell_bbox 格式
            page_elements = []
            for element in page.get('elements', []):
                converted = JSONFormatters._element_to_cell_bbox_format(element, page_idx, page_rotation_angle)
                if converted:
                    page_elements.append(converted)
            
            # 添加丢弃元素
            for element in page.get('discarded_blocks', []):
                converted = JSONFormatters._element_to_cell_bbox_format(element, page_idx, page_rotation_angle)
                if converted:
                    page_elements.append(converted)
            
            # 转换为 JSON 字符串
            json_content = json.dumps(page_elements, ensure_ascii=False, indent=2, cls=NumpyEncoder)
            
            # 金额数字标准化
            if normalize_numbers:
                original_content = json_content
                json_content = normalize_json_table(json_content)
                
                if json_content != original_content:
                    original_path = output_dir / f"{page_name}_original.json"
                    with open(original_path, 'w', encoding='utf-8') as f:
                        f.write(original_content)
                    logger.debug(f"📄 Original page JSON saved: {original_path}")
            
            # 保存 JSON
            json_path = output_dir / f"{page_name}.json"
            with open(json_path, 'w', encoding='utf-8') as f:
                f.write(json_content)
            
            saved_paths.append(str(json_path))
            logger.debug(f"📄 Page JSON saved: {json_path}")
        
        if saved_paths:
            logger.info(f"📄 {len(saved_paths)} page JSONs saved")
        
        return saved_paths
    
    @staticmethod
    def _element_to_cell_bbox_format(
        element: Dict[str, Any],
        page_idx: int,
        page_rotation_angle: float = 0.0
    ) -> Optional[Dict[str, Any]]:
        """
        将元素转换为 mineru_vllm_results_cell_bbox 格式
        
        Args:
            element: 元素字典
            page_idx: 页面索引
            page_rotation_angle: 页面旋转角度（0, 90, 180, 270）
        """
        elem_type = element.get('type', '')
        bbox = element.get('bbox', [0, 0, 0, 0])
        content = element.get('content', {})
        
        # 确保 bbox 是整数列表
        bbox = [int(x) for x in bbox[:4]] if bbox else [0, 0, 0, 0]
        
        result = {
            'bbox': bbox,
            'page_idx': page_idx,
            'page_rotation_angle': page_rotation_angle,
            'reading_order': element.get('reading_order', 0)
        }
        
        # 文本类型
        if elem_type in ['text', 'title', 'ref_text', 'ocr_text']:
            text = content.get('text', '') if isinstance(content, dict) else str(content)
            result['type'] = 'text' if elem_type != 'title' else 'title'
            result['text'] = text
            if elem_type == 'title':
                result['text_level'] = element.get('level', 1)
        
        # 表格类型
        elif elem_type in ['table', 'table_body']:
            result['type'] = 'table'
            result['img_path'] = content.get('table_image_path', '')
            result['table_caption'] = JSONFormatters._ensure_list(content.get('table_caption', []))
            result['table_footnote'] = JSONFormatters._ensure_list(content.get('table_footnote', []))
            result['table_body'] = content.get('html', '')
            
            # 关键：table_cells 数组
            cells = content.get('cells', [])
            if cells:
                result['table_cells'] = JSONFormatters.format_table_cells(cells)
            
            # 旋转和倾斜信息
            if 'table_angle' in content:
                result['image_rotation_angle'] = float(content['table_angle'])
            if 'skew_angle' in content:
                result['skew_angle'] = float(content['skew_angle'])
        
        # 图片类型
        elif elem_type in ['image', 'image_body', 'figure']:
            result['type'] = 'image'
            image_filename = content.get('image_path', '')
            result['img_path'] = f"images/{image_filename}" if image_filename else ''
            result['image_caption'] = JSONFormatters._ensure_list(content.get('caption', []))
            result['image_footnote'] = JSONFormatters._ensure_list(content.get('footnote', []))
        
        # 公式类型
        elif elem_type in ['interline_equation', 'inline_equation', 'equation']:
            result['type'] = 'equation'
            result['text'] = content.get('latex', '') if isinstance(content, dict) else ''
            result['text_format'] = 'latex'
        
        # 列表类型
        elif elem_type == 'list':
            result['type'] = 'list'
            result['sub_type'] = 'text'
            result['list_items'] = content.get('list_items', []) if isinstance(content, dict) else []
        
        # 页眉页脚
        elif elem_type in ['header', 'footer']:
            result['type'] = elem_type
            result['text'] = content.get('text', '') if isinstance(content, dict) else str(content)
        
        # 表格/图片附属文本
        elif elem_type in ['table_caption', 'table_footnote', 'image_caption', 'image_footnote']:
            result['type'] = elem_type
            result['text'] = content.get('text', '') if isinstance(content, dict) else str(content)
        
        # 丢弃元素
        elif elem_type in ['discarded', 'abandon']:
            result['type'] = 'discarded'
            result['original_category'] = element.get('original_category', 'unknown')
            result['text'] = content.get('text', '') if isinstance(content, dict) else ''
        
        else:
            return None
        
        return result
    
    @staticmethod
    def format_table_cells(cells: List[Dict]) -> List[Dict[str, Any]]:
        """
        格式化表格单元格为 mineru_vllm_results_cell_bbox 格式
        
        输出格式:
        {
            "type": "table_cell",
            "text": "单元格内容",
            "matched_text": "OCR匹配文本",
            "bbox": [x1, y1, x2, y2],
            "row": 1,
            "col": 1,
            "score": 100.0,
            "paddle_bbox_indices": [0, 1]
        }
        """
        formatted_cells = []
        
        for cell in cells:
            formatted_cell = {
                'type': 'table_cell',
                'text': cell.get('text', ''),
                'matched_text': cell.get('matched_text', cell.get('text', '')),
                'bbox': [float(x) for x in cell.get('bbox', [0, 0, 0, 0])[:4]],
                'row': cell.get('row', 0),
                'col': cell.get('col', 0),
                'score': float(cell.get('score', 100.0)),
                'paddle_bbox_indices': cell.get('paddle_bbox_indices', 
                                                cell.get('paddle_indices', []))
            }
            formatted_cells.append(formatted_cell)
        
        return formatted_cells
    
    @staticmethod
    def _ensure_list(value) -> List:
        """确保值是列表"""
        if value is None:
            return []
        if isinstance(value, str):
            return [value] if value else []
        if isinstance(value, list):
            return value
        return [str(value)]