ocr_verify/merge_mineru_paddle_ocr.py

"""
合并 MinerU 和 PaddleOCR 的结果
使用 MinerU 的表格结构识别 + PaddleOCR 的文字框坐标
"""
import json
import re
from pathlib import Path
from typing import List, Dict, Tuple, Optional
from bs4 import BeautifulSoup
from fuzzywuzzy import fuzz


class MinerUPaddleOCRMerger:
    """合并 MinerU 和 PaddleOCR 的结果"""
    
    def __init__(self, look_ahead_window: int = 10, similarity_threshold: int = 80):
        """
        Args:
            look_ahead_window: 向前查找的窗口大小
            similarity_threshold: 文本相似度阈值
        """
        self.look_ahead_window = look_ahead_window
        self.similarity_threshold = similarity_threshold
    
    def merge_table_with_bbox(self, mineru_json_path: str, paddle_json_path: str, 
                              output_path: Optional[str] = None) -> Dict:
        """
        合并 MinerU 和 PaddleOCR 的结果
        
        Args:
            mineru_json_path: MinerU 输出的 JSON 路径
            paddle_json_path: PaddleOCR 输出的 JSON 路径
            output_path: 输出路径（可选）
        
        Returns:
            合并后的结果字典
        """
        # 加载数据
        with open(mineru_json_path, 'r', encoding='utf-8') as f:
            mineru_data = json.load(f)
        
        with open(paddle_json_path, 'r', encoding='utf-8') as f:
            paddle_data = json.load(f)
        
        # 提取 PaddleOCR 的文字框信息
        paddle_text_boxes = self._extract_paddle_text_boxes(paddle_data)
        
        # 处理 MinerU 的数据
        merged_data = self._process_mineru_data(mineru_data, paddle_text_boxes)
        
        # 保存结果
        if output_path:
            output_path = Path(output_path).resolve()
            output_path.parent.mkdir(parents=True, exist_ok=True)
            with open(str(output_path), 'w', encoding='utf-8') as f:
                json.dump(merged_data, f, ensure_ascii=False, indent=2)
        
        return merged_data
    
    def _extract_paddle_text_boxes(self, paddle_data: Dict) -> List[Dict]:
        """提取 PaddleOCR 的文字框信息"""
        text_boxes = []
        
        if 'overall_ocr_res' in paddle_data:
            ocr_res = paddle_data['overall_ocr_res']
            rec_texts = ocr_res.get('rec_texts', [])
            rec_polys = ocr_res.get('rec_polys', [])
            rec_scores = ocr_res.get('rec_scores', [])

            for i, (text, poly, score) in enumerate(zip(rec_texts, rec_polys, rec_scores)):
                if text and text.strip():
                    # 计算 bbox (x_min, y_min, x_max, y_max)
                    xs = [p[0] for p in poly]
                    ys = [p[1] for p in poly]
                    bbox = [min(xs), min(ys), max(xs), max(ys)]
                    
                    text_boxes.append({
                        'text': text,
                        'bbox': bbox,
                        'poly': poly,
                        'score': score,
                        'paddle_bbox_index': i,
                        'used': False  # 标记是否已被使用
                    })

        return text_boxes
    
    def _process_mineru_data(self, mineru_data: List[Dict], 
                            paddle_text_boxes: List[Dict]) -> List[Dict]:
        """处理 MinerU 数据，添加 bbox 信息"""
        merged_data = []
        paddle_pointer = 0  # PaddleOCR 文字框指针
        
        for item in mineru_data:
            if item['type'] == 'table':
                # 处理表格
                merged_item = item.copy()
                table_html = item.get('table_body', '')
                
                # 解析 HTML 表格并添加 bbox
                enhanced_html, paddle_pointer = self._enhance_table_html_with_bbox(
                    table_html, paddle_text_boxes, paddle_pointer
                )
                
                merged_item['table_body'] = enhanced_html
                merged_item['table_body_with_bbox'] = enhanced_html
                merged_item['bbox_mapping'] = 'merged_from_paddle_ocr'
                
                merged_data.append(merged_item)
            
            elif item['type'] in ['text', 'header']:
                # 处理普通文本
                merged_item = item.copy()
                text = item.get('text', '')
                
                # 查找匹配的 bbox
                matched_bbox, paddle_pointer = self._find_matching_bbox(
                    text, paddle_text_boxes, paddle_pointer
                )
                
                if matched_bbox:
                    merged_item['bbox'] = matched_bbox['bbox']
                    merged_item['bbox_source'] = 'paddle_ocr'
                    merged_item['text_score'] = matched_bbox['score']
                    # 标记为已使用
                    matched_bbox['used'] = True
                
                merged_data.append(merged_item)
            
            else:
                # 其他类型直接复制
                merged_data.append(item.copy())
        
        return merged_data
    
    def _enhance_table_html_with_bbox(self, html: str, paddle_text_boxes: List[Dict], 
                                      start_pointer: int) -> Tuple[str, int]:
        """
        为 HTML 表格添加 bbox 信息
        
        Args:
            html: 原始 HTML 表格
            paddle_text_boxes: PaddleOCR 文字框列表
            start_pointer: 起始指针位置
        
        Returns:
            (增强后的 HTML, 新的指针位置)
        """
        soup = BeautifulSoup(html, 'html.parser')
        current_pointer = start_pointer
        
        # 遍历所有单元格
        for cell in soup.find_all(['td', 'th']):
            cell_text = cell.get_text(strip=True)
            
            if not cell_text:
                continue
            
            # 查找匹配的 bbox
            matched_bbox, current_pointer = self._find_matching_bbox(
                cell_text, paddle_text_boxes, current_pointer
            )
            
            if matched_bbox:
                # 添加 data-bbox 属性
                bbox = matched_bbox['bbox']
                cell['data-bbox'] = f"[{bbox[0]},{bbox[1]},{bbox[2]},{bbox[3]}]"
                cell['data-score'] = f"{matched_bbox['score']:.4f}"
                cell['data-paddle-index'] = str(matched_bbox['paddle_bbox_index'])
                
                # 标记为已使用
                matched_bbox['used'] = True
        
        return str(soup), current_pointer
    
    def _find_matching_bbox(self, target_text: str, text_boxes: List[Dict], 
                           start_index: int) -> tuple[Optional[Dict], int]:
        """
        查找匹配的文字框
        
        Args:
            target_text: 目标文本
            text_boxes: 文字框列表
            start_index: 起始索引
        
        Returns:
            匹配的文字框信息，如果未找到返回 None
        """
        target_text = self._normalize_text(target_text)
        
        # 在窗口范围内查找
        search_end = min(start_index + self.look_ahead_window, len(text_boxes))
        
        best_match = None
        best_index = start_index
        best_similarity = 0
        
        for i in range(start_index, search_end):
            if text_boxes[i]['used']:
                continue
            
            box_text = self._normalize_text(text_boxes[i]['text'])
            
            # 计算相似度
            # similarity = fuzz.ratio(target_text, box_text)
            similarity = fuzz.token_set_ratio(target_text, box_text)
            
            # 精确匹配优先
            if target_text == box_text:
                return text_boxes[i], i + 1
            
            # 记录最佳匹配
            if similarity > best_similarity and similarity >= self.similarity_threshold:
                best_similarity = similarity
                best_match = text_boxes[i]
                best_index = i + 1

        return best_match, best_index

    def _normalize_text(self, text: str) -> str:
        """标准化文本（去除空格、标点等）"""
        # 移除所有空白字符
        text = re.sub(r'\s+', '', text)
        # 转换全角数字和字母为半角
        text = self._full_to_half(text)
        return text.lower()
    
    def _full_to_half(self, text: str) -> str:
        """全角转半角"""
        result = []
        for char in text:
            code = ord(char)
            if code == 0x3000:  # 全角空格
                code = 0x0020
            elif 0xFF01 <= code <= 0xFF5E:  # 全角字符
                code -= 0xFEE0
            result.append(chr(code))
        return ''.join(result)
    
    def generate_enhanced_markdown(self, merged_data: List[Dict], 
                                   output_path: Optional[str] = None) -> str:
        """
        生成增强的 Markdown（包含 bbox 信息的注释）
        
        Args:
            merged_data: 合并后的数据
            output_path: 输出路径（可选）
        
        Returns:
            Markdown 内容
        """
        md_lines = []
        
        for item in merged_data:
            if item['type'] == 'header':
                text = item.get('text', '')
                bbox = item.get('bbox', [])
                md_lines.append(f"<!-- bbox: {bbox} -->")
                md_lines.append(f"# {text}\n")
            
            elif item['type'] == 'text':
                text = item.get('text', '')
                bbox = item.get('bbox', [])
                if bbox:
                    md_lines.append(f"<!-- bbox: {bbox} -->")
                md_lines.append(f"{text}\n")
            
            elif item['type'] == 'table':
                md_lines.append("\n## 表格\n")
                md_lines.append("<!-- 表格单元格包含 data-bbox 属性 -->\n")
                md_lines.append(item.get('table_body_with_bbox', item.get('table_body', '')))
                md_lines.append("\n")
        
        markdown_content = '\n'.join(md_lines)
        
        if output_path:
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write(markdown_content)
        
        return markdown_content
    
    def extract_table_cells_with_bbox(self, merged_data: List[Dict]) -> List[Dict]:
        """
        提取所有表格单元格及其 bbox 信息
        
        Returns:
            单元格列表，每个包含 text, bbox, row, col 等信息
        """
        cells = []
        
        for item in merged_data:
            if item['type'] != 'table':
                continue
            
            html = item.get('table_body_with_bbox', item.get('table_body', ''))
            soup = BeautifulSoup(html, 'html.parser')
            
            # 遍历所有行
            for row_idx, row in enumerate(soup.find_all('tr')):
                # 遍历所有单元格
                for col_idx, cell in enumerate(row.find_all(['td', 'th'])):
                    cell_text = cell.get_text(strip=True)
                    bbox_str = cell.get('data-bbox', '')
                    
                    if bbox_str:
                        try:
                            bbox = json.loads(bbox_str)
                            cells.append({
                                'text': cell_text,
                                'bbox': bbox,
                                'row': row_idx,
                                'col': col_idx,
                                'score': float(cell.get('data-score', 0)),
                                'paddle_index': int(cell.get('data-paddle-index', -1))
                            })
                        except (json.JSONDecodeError, ValueError):
                            pass
        
        return cells


def merge_mineru_paddle_batch(mineru_dir: str, paddle_dir: str, output_dir: str):
    """
    批量合并 MinerU 和 PaddleOCR 的结果
    
    Args:
        mineru_dir: MinerU 结果目录
        paddle_dir: PaddleOCR 结果目录
        output_dir: 输出目录
    """
    mineru_path = Path(mineru_dir)
    paddle_path = Path(paddle_dir)
    output_path = Path(output_dir)
    output_path.mkdir(parents=True, exist_ok=True)
    
    merger = MinerUPaddleOCRMerger(look_ahead_window=10, similarity_threshold=80)
    
    # 查找所有 MinerU 的 JSON 文件, page_001.json
    mineru_files = list(mineru_path.glob('*_page_*[0-9].json'))
    mineru_files.sort()
    
    print(f"找到 {len(mineru_files)} 个 MinerU 文件")
    
    for mineru_file in mineru_files:
        # 查找对应的 PaddleOCR 文件
        paddle_file = paddle_path / mineru_file.name
        
        if not paddle_file.exists():
            print(f"⚠️ 未找到对应的 PaddleOCR 文件: {paddle_file}")
            continue
        
        print(f"处理: {mineru_file.name}")
        
        # 输出文件路径
        merged_json_path = output_path / f"{mineru_file.stem}_merged.json"
        merged_md_path = output_path / f"{mineru_file.stem}_merged.md"
        cells_json_path = output_path / f"{mineru_file.stem}_cells.json"
        
        try:
            # 合并数据
            merged_data = merger.merge_table_with_bbox(
                str(mineru_file),
                str(paddle_file),
                str(merged_json_path)
            )
            
            # 生成 Markdown
            merger.generate_enhanced_markdown(merged_data, str(merged_md_path))
            
            # 提取单元格信息
            cells = merger.extract_table_cells_with_bbox(merged_data)
            
            with open(cells_json_path, 'w', encoding='utf-8') as f:
                json.dump(cells, f, ensure_ascii=False, indent=2)
            
            print(f"  ✅ 合并完成")
            print(f"  - 提取了 {len(cells)} 个表格单元格")
            
        except Exception as e:
            print(f"  ❌ 处理失败: {e}")
            import traceback
            traceback.print_exc()


if __name__ == "__main__":
    # 示例用法
    mineru_dir = "/Users/zhch158/workspace/data/流水分析/A用户_单元格扫描流水/mineru-vlm-2.5.3_Results"
    paddle_dir = "/Users/zhch158/workspace/data/流水分析/A用户_单元格扫描流水/data_PPStructureV3_Results"
    output_dir = "/Users/zhch158/workspace/data/流水分析/A用户_单元格扫描流水/merged_results"
    
    merge_mineru_paddle_batch(mineru_dir, paddle_dir, output_dir)

    # 示例：合并1个文件
    # mineru_json = "/Users/zhch158/workspace/data/流水分析/A用户_单元格扫描流水/mineru-vlm-2.5.3_Results/A用户_单元格扫描流水_page_001.json"
    # paddle_json = "/Users/zhch158/workspace/data/流水分析/A用户_单元格扫描流水/data_PPStructureV3_Results/A用户_单元格扫描流水_page_001.json"
    # output_json = "/Users/zhch158/workspace/data/流水分析/A用户_单元格扫描流水/merged_results/A用户_单元格扫描流水_page_001.json"

    # merger = MinerUPaddleOCRMerger(look_ahead_window=10, similarity_threshold=80)
    # merger.merge_table_with_bbox(mineru_json, paddle_json, output_json)