kocr/core/recognizer.py

#!/usr/bin/env python3
"""KOCR v4 — OCR 识别模块（TableMaster + DB+SVTR 双管线）"""

import os
import sys
import json


def _build_ocr_params(config: dict, photo_type: str = 'normal'):
    """根据配置和照片类型构建 PaddleOCR 参数
    
    Args:
        config: config.yaml 中 ocr 章节
        photo_type: 照片类型（用于覆盖参数）
    
    Returns:
        params dict
    """
    params = {
        'use_angle_cls': config.get('use_angle_cls', True),
        'lang': config.get('lang', 'ch'),
        'use_gpu': config.get('use_gpu', False),
        'show_log': config.get('show_log', False),
        'det_db_thresh': config.get('det_db_thresh', 0.3),
        'det_db_box_thresh': config.get('det_db_box_thresh', 0.5),
        'det_db_unclip_ratio': config.get('det_db_unclip_ratio', 2.2),
        'det_limit_side_len': config.get('max_side_len', 1600),
        'det_db_score_mode': config.get('det_db_score_mode', 'slow'),
        'min_size': config.get('min_size', 8),
        'rec_drop_score': config.get('rec_drop_score', 0.5),
    }
    
    # 自定义字典路径（相对于项目根目录）
    rec_dict = config.get('rec_char_dict_path')
    if rec_dict:
        params['rec_char_dict_path'] = os.path.abspath(rec_dict)
    
    # 照片类型参数覆盖
    overrides = config.get('photo_type_overrides', {}).get(photo_type, {})
    params.update(overrides)
    
    return params


def _check_gpu_available() -> bool:
    """检查 GPU 是否对 PaddleOCR 可用（不仅看 nvidia-smi）"""
    try:
        # 先看 nvidia-smi
        import subprocess
        result = subprocess.run(
            ['nvidia-smi'],
            capture_output=True,
            text=True,
            timeout=5
        )
        if result.returncode != 0:
            return False
        # 再看 PaddlePaddle 能否真正加载 GPU
        import paddle
        return paddle.is_compiled_with_cuda()
    except Exception:
        return False


def run_ocr_ppstructure(img, config: dict, verbose=False):
    """尝试 PPStructure TableMaster 表格识别
    
    Args:
        img: 预处理后的图像
        config: OCR 配置
        verbose: 是否输出详细信息
    
    Returns:
        lines 列表或 None（失败）
    """
    table_algorithm = config.get('table_algorithm', 'TableMaster')
    use_gpu = _check_gpu_available() and config.get('use_gpu', False)
    
    if verbose:
        print(f"   📋 PPStructure 识别 ({table_algorithm})...", file=sys.stderr)
    
    try:
        from paddleocr import PPStructure
        
        pp_args = {
            'show_log': False,
            'lang': 'ch',
            'layout': True,
            'ocr': True,
            'table_algorithm': table_algorithm,
            'use_gpu': use_gpu,
        }
        
        engine = PPStructure(**pp_args)
        table_result = engine(img)
        
        for res in table_result:
            if res['type'] == 'table':
                html = res['res']['html']
                if '<td>' in html and html.count('<tr>') >= 2:
                    lines = parse_table_html_to_lines(html)
                    if lines and len(lines) >= 4:
                        if verbose:
                            print(f"   ✅ TableMaster 识别成功", file=sys.stderr)
                        return lines
    except Exception as e:
        if verbose:
            print(f"   ⚠️  TableMaster 不可用: {e}", file=sys.stderr)
    
    return None


def run_ocr_det_rec(img, config: dict, photo_type: str = 'normal', verbose=False):
    """DB 检测 + SVTR 文本识别（回退方案）
    
    Args:
        img: 预处理后的图像
        config: OCR 配置
        photo_type: 照片类型
        verbose: 是否输出详细信息
    
    Returns:
        lines 列表
    """
    params = _build_ocr_params(config, photo_type)
    use_gpu = _check_gpu_available() and config.get('use_gpu', False)
    params['use_gpu'] = use_gpu
    
    if verbose:
        print(f"   🔍 DB+SVTR 识别 (GPU={use_gpu})", file=sys.stderr)
    
    from .postprocessor import char_correct
    from paddleocr import PaddleOCR
    
    ocr = PaddleOCR(**params)
    result = ocr.ocr(img, cls=True)
    
    # 防御：PaddleOCR 可能返回 None 或 [None]
    if not result or result[0] is None:
        if verbose:
            print(f"   ⚠️  PaddleOCR 未检测到文字", file=sys.stderr)
        return []
    
    lines = []
    for line_group in result:
        for line in line_group:
            box = line[0]
            text = char_correct(line[1][0])
            conf = line[1][1]
            xs = [p[0] for p in box]
            ys = [p[1] for p in box]
            cx = sum(xs) / 4
            cy = sum(ys) / 4
            lines.append((text, conf, cx, cy, box))
    
    return lines


def parse_table_html_to_lines(html):
    """解析 PPStructure 表格 HTML 为结构化行数据"""
    from bs4 import BeautifulSoup
    lines = []
    try:
        soup = BeautifulSoup(html, 'html.parser')
        tables = soup.find_all('table')
        if not tables:
            return []
        rows = tables[0].find_all('tr')
        for ri, row in enumerate(rows):
            cells = row.find_all(['td', 'th'])
            row_texts = []
            for ci, cell in enumerate(cells):
                text = cell.get_text(strip=True)
                if text:
                    x = float(ci * 300 + 100)
                    y = float(ri * 200 + 100)
                    lines.append((text, 0.95, x, y, None))
                    row_texts.append(text)
            if not row_texts:
                continue
    except Exception:
        return []
    return lines


def run_ocr(img, config: dict, photo_type: str = 'normal', verbose=False):
    """完整 OCR 识别入口：先 TableMaster，失败回退 DB+SVTR
    
    Args:
        img: 预处理后的图像
        config: OCR 配置
        photo_type: 照片类型
        verbose: 是否输出详细信息
    
    Returns:
        识别结果行列表
    """
    # 1. 尝试 TableMaster 表格识别
    lines = run_ocr_ppstructure(img, config, verbose=verbose)
    if lines:
        return lines
    
    # 2. 回退 DB+SVTR
    lines = run_ocr_det_rec(img, config, photo_type, verbose=verbose)
    return lines