module/text_classifier.py

"""场景分类器：多信号加权评分判断图片属于文本场景还是复杂场景

综合以下4个信号：
1. 文本面积占比（权重0.15）——OCR检测到的文本区域面积
2. 背景白度占比（权重0.50）——浅色像素比例
3. 颜色方差倒数（权重0.25）——颜色变化越小越可能是文档
4. 文本行密度（权重0.10）——单位面积内文本行数
"""

import cv2
import numpy as np

from module.config import (
    TEXT_AREA_WEIGHT, WHITENESS_WEIGHT, LOW_VARIANCE_WEIGHT,
    LINE_DENSITY_WEIGHT, DOC_SCORE_THRESHOLD,
)


def _cv_imread(path, flags=cv2.IMREAD_COLOR):
    """支持中文路径的cv2.imread"""
    return cv2.imdecode(np.fromfile(path, dtype=np.uint8), flags)


def _compute_whiteness(img, brightness_threshold=200):
    """计算图片背景白度占比

    白底文档的浅色像素（亮度>200）占比通常>60%

    Args:
        img: BGR格式的numpy数组
        brightness_threshold: 亮度阈值，高于此值视为'白'
    Returns:
        whiteness: 白色像素占比，范围[0,1]
    """
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    white_pixels = np.sum(gray >= brightness_threshold)
    total_pixels = gray.shape[0] * gray.shape[1]
    return white_pixels / total_pixels if total_pixels > 0 else 0.0


def _compute_color_variance_norm(img):
    """计算图片颜色方差的归一化值

    文档图片（白底黑字）颜色方差低，场景照片颜色方差高
    归一化到[0,1]，值越大表示颜色越丰富（越不像文档）

    Args:
        img: BGR格式的numpy数组
    Returns:
        variance_norm: 归一化颜色方差，范围[0,1]
    """
    variances = [np.var(img[:, :, c]) for c in range(3)]
    avg_var = np.mean(variances)
    # sigmoid归一化：方差2000约为0.5（经验值）
    # 文档方差通常<500，场景照片方差通常>3000
    variance_norm = 1.0 / (1.0 + np.exp(-(avg_var - 2000) / 800))
    return float(variance_norm)


def _compute_line_density(text_boxes, img_shape):
    """计算文本行密度

    即使文字分散，如果行数多，行密度仍然高
    营业执照虽然文字分散，但行数多（10+行），密度不会太低

    Args:
        text_boxes: OCR检测到的文本框列表
        img_shape: (h, w) 图片尺寸
    Returns:
        line_density: 归一化行密度，范围[0,1]
    """
    h, w = img_shape[:2]
    if not text_boxes or h == 0 or w == 0:
        return 0.0

    # 行数 / 图片面积（以1000*1000为参考）
    area_factor = (h * w) / (1000 * 1000)
    if area_factor == 0:
        return 0.0

    raw_density = len(text_boxes) / area_factor
    # sigmoid归一化：10行/百万像素约为0.5
    density_norm = 1.0 / (1.0 + np.exp(-(raw_density - 10) / 5))
    return float(density_norm)


def _ocr_image(image_path, ocr_engine):
    """仅执行 OCR 提取文本，不做场景分类

    Returns:
        full_text: 提取的文本字符串
    """
    try:
        result = ocr_engine.ocr(image_path, cls=True)
        if result and result[0]:
            return ' '.join(line[1][0] for line in result[0])
    except Exception as e:
        print(f"[OCR] 提取失败: {e}")
    return ''


def classify_scene(image_path, ocr_engine, doc_score_threshold=DOC_SCORE_THRESHOLD,
                   precomputed_ocr_result=None):
    """多信号加权评分判断图片属于文本场景还是复杂场景

    快速预判：先用廉价的视觉信号（白度+方差）做预判，只有中间模糊区才调用 OCR。
    典型情况下可省掉 OCR，从 2.6s 降到 0.1s。

    Args:
        image_path: 图片路径
        ocr_engine: PaddleOCR引擎实例
        doc_score_threshold: 文档场景判定阈值，默认0.55
        precomputed_ocr_result: 预计算的OCR结果，避免重复OCR

    Returns:
        scene_type: 'text'(文本场景) 或 'complex'(复杂场景)
        doc_score: 文档评分
        detail: 各信号的详细分数字典
        ocr_result: OCR原始结果（可复用给下游）
    """
    img = _cv_imread(image_path)
    if img is None:
        return 'complex', 0.0, {}, None

    h, w = img.shape[:2]
    total_area = h * w

    # === 第1步：廉价的视觉信号（<5ms） ===
    whiteness = _compute_whiteness(img)
    variance_norm = _compute_color_variance_norm(img)
    low_variance = 1.0 - variance_norm

    # === 第2步：快速预判（省 OCR）===
    # 白度>0.7 且 颜色方差低 → 文档（白底文档的典型特征）
    # 白度<0.2 且 颜色方差高 → 场景照片（彩色实景的典型特征）
    # 仅当落在模糊区才调用完整 OCR
    needs_ocr = True
    fast_scene = None
    fast_doc_score = None

    # 白底文档：白度极高 + 方差极低
    if whiteness >= 0.70 and low_variance >= 0.7:
        # 视觉信号已足够判断为 text
        fast_scene = 'text'
        # 估算 doc_score：whiteness 权重0.5 + low_variance 权重0.25 = 0.75
        # text_area 和 line_density 假设中等值（0.3, 0.3）
        fast_doc_score = (
            TEXT_AREA_WEIGHT * 0.3
            + WHITENESS_WEIGHT * whiteness
            + LOW_VARIANCE_WEIGHT * low_variance
            + LINE_DENSITY_WEIGHT * 0.3
        )
        needs_ocr = False
    # 彩色实景：白度极低 + 方差极高
    elif whiteness <= 0.20 and low_variance <= 0.3:
        fast_scene = 'complex'
        fast_doc_score = (
            TEXT_AREA_WEIGHT * 0.0
            + WHITENESS_WEIGHT * whiteness
            + LOW_VARIANCE_WEIGHT * low_variance
            + LINE_DENSITY_WEIGHT * 0.0
        )
        needs_ocr = False

    if not needs_ocr:
        detail = {
            'text_area_ratio': 0.0 if fast_scene == 'complex' else 0.3,
            'whiteness': round(whiteness, 4),
            'low_variance': round(low_variance, 4),
            'line_density': 0.0 if fast_scene == 'complex' else 0.3,
            'text_box_count': 0,
            'full_text': '',
            'fast_classified': True,
        }
        return fast_scene, round(fast_doc_score, 4), detail, None

    # === 第3步：OCR（仅在模糊区执行，~1.5-2s） ===
    if precomputed_ocr_result is not None:
        result = precomputed_ocr_result
    else:
        result = ocr_engine.ocr(image_path, cls=True)

    text_area = 0
    text_boxes = []
    full_text = ''
    if result and result[0]:
        for line in result[0]:
            box = line[0]
            text_boxes.append(box)
            n = len(box)
            area = 0
            for i in range(n):
                j = (i + 1) % n
                area += box[i][0] * box[j][1]
                area -= box[j][0] * box[i][1]
            text_area += abs(area) / 2
        full_text = ' '.join(line[1][0] for line in result[0])

    text_area_ratio = text_area / total_area if total_area > 0 else 0
    line_density = _compute_line_density(text_boxes, img.shape)

    doc_score = (
        TEXT_AREA_WEIGHT * text_area_ratio
        + WHITENESS_WEIGHT * whiteness
        + LOW_VARIANCE_WEIGHT * low_variance
        + LINE_DENSITY_WEIGHT * line_density
    )

    scene_type = 'text' if doc_score >= doc_score_threshold else 'complex'

    detail = {
        'text_area_ratio': round(text_area_ratio, 4),
        'whiteness': round(whiteness, 4),
        'low_variance': round(low_variance, 4),
        'line_density': round(line_density, 4),
        'text_box_count': len(text_boxes),
        'full_text': full_text,
    }

    return scene_type, round(doc_score, 4), detail, result