ocr_platform/ocr_utils/watermark/contrast.py

"""水印 对比度增强（由 ocr_utils.watermark_utils 迁入）。"""
from __future__ import annotations

import json
import re
from pathlib import Path
from typing import Any, Dict, Optional, Tuple, Union

import cv2
import numpy as np
from loguru import logger
from PIL import Image

def _enhance_text_restore(
    gray: np.ndarray,
    *,
    background_threshold: int = 248,
    text_lo_percentile: float = 1.0,
    text_hi_percentile: float = 99.0,
    text_black_target: int = 85,
) -> np.ndarray:
    """
    仅对非背景像素做动态范围压缩，将最深笔画拉向 text_black_target（默认 ~85，接近扫描件原图）。

    背景（>= background_threshold）保持白色，避免整图 gamma 导致背景发灰。
    """
    result = gray.copy()
    bg_th = int(np.clip(background_threshold, 200, 255))
    text_mask = gray < bg_th
    if not np.any(text_mask):
        return result

    vals = gray[text_mask].astype(np.float32)
    lo = float(np.percentile(vals, text_lo_percentile))
    hi = float(np.percentile(vals, text_hi_percentile))
    target = int(np.clip(text_black_target, 10, 200))
    if hi <= lo + 1.0:
        return result

    stretched = (vals - lo) * target / (hi - lo)
    result[text_mask] = np.clip(stretched, 0, 255).astype(np.uint8)
    return result


def enhance_document_contrast(
    gray: np.ndarray,
    method: str = "text_restore",
    *,
    clip_limit: float = 2.0,
    tile_grid_size: int = 8,
    gamma: float = 0.85,
    black_percentile: float = 2.0,
    white_percentile: float = 98.0,
    background_threshold: int = 248,
    text_lo_percentile: float = 1.0,
    text_hi_percentile: float = 99.0,
    text_black_target: int = 85,
) -> np.ndarray:
    """
    文档灰度图对比度增强（常用于去水印后恢复笔画深度）。

    Args:
        gray: 单通道 uint8 灰度图
        method: text_restore | clahe | gamma | linear
        clip_limit: CLAHE 对比度限制
        tile_grid_size: CLAHE 分块大小
        gamma: gamma 校正指数，<1 加深文字（去水印后发浅时适用）
        black_percentile: linear 拉伸下分位（映射到 0）
        white_percentile: linear 拉伸上分位（映射到 255）
        background_threshold: text_restore 背景阈值（>= 视为白底不处理）
        text_lo_percentile: text_restore 笔画下分位
        text_hi_percentile: text_restore 笔画上分位（映射到 text_black_target）
        text_black_target: text_restore 最深笔画目标灰度（越小越深，建议 75~95）

    Returns:
        增强后的灰度图
    """
    if gray is None or gray.size == 0:
        return gray
    if gray.ndim != 2:
        raise ValueError("enhance_document_contrast expects single-channel grayscale image")

    method = (method or "text_restore").lower().strip()

    if method == "text_restore":
        return _enhance_text_restore(
            gray,
            background_threshold=background_threshold,
            text_lo_percentile=text_lo_percentile,
            text_hi_percentile=text_hi_percentile,
            text_black_target=text_black_target,
        )

    if method == "gamma":
        gamma = max(0.1, min(float(gamma), 3.0))
        inv_gamma = 1.0 / gamma
        table = np.array(
            [((i / 255.0) ** inv_gamma) * 255 for i in range(256)],
            dtype=np.uint8,
        )
        return cv2.LUT(gray, table)

    if method == "linear":
        p_low = float(np.percentile(gray, black_percentile))
        p_high = float(np.percentile(gray, white_percentile))
        if p_high <= p_low + 1.0:
            return gray
        stretched = (gray.astype(np.float32) - p_low) * 255.0 / (p_high - p_low)
        return np.clip(stretched, 0, 255).astype(np.uint8)

    # 默认 CLAHE：局部对比度，适合扫描件
    tile = max(2, int(tile_grid_size))
    clahe = cv2.createCLAHE(
        clipLimit=max(0.1, float(clip_limit)),
        tileGridSize=(tile, tile),
    )
    return clahe.apply(gray)


def apply_contrast_enhancement_config(
    gray: np.ndarray,
    contrast_cfg: Optional[Dict[str, Any]],
) -> np.ndarray:
    """按配置字典应用对比度增强；未启用时原样返回。"""
    if not contrast_cfg or not contrast_cfg.get("enabled", False):
        return gray
    return enhance_document_contrast(
        gray,
        method=contrast_cfg.get("method", "text_restore"),
        clip_limit=contrast_cfg.get("clip_limit", 2.0),
        tile_grid_size=contrast_cfg.get("tile_grid_size", 8),
        gamma=contrast_cfg.get("gamma", 0.85),
        black_percentile=contrast_cfg.get("black_percentile", 2.0),
        white_percentile=contrast_cfg.get("white_percentile", 98.0),
        background_threshold=contrast_cfg.get("background_threshold", 248),
        text_lo_percentile=contrast_cfg.get("text_lo_percentile", 1.0),
        text_hi_percentile=contrast_cfg.get("text_hi_percentile", 99.0),
        text_black_target=contrast_cfg.get("text_black_target", 75),
    )