ocr_platform/ocr_tools/remove_watermark_tool/remove_watermark.py

"""
银行流水水印去除工具

支持 PDF 和常见图片格式（jpg/png/tif/bmp/webp）。
- 输入 PDF → 输出去水印 PDF（扫描件）或直接复制（文字型）
- 输入图片 → 输出去水印图片（保持原格式）
适用于福建农信、邮储银行等带有半透明文字水印的银行流水单。

用法:
    # 处理单个 PDF 或图片
    python remove_watermark.py input.pdf
    python remove_watermark.py input.jpg

    # 指定输出路径
    python remove_watermark.py input.pdf -o output.pdf

    # 指定页面范围（支持 "1-5,7,9-12" 格式）
    python remove_watermark.py input.pdf --page-range 1-3

    # 调整去除阈值（默认 160，范围建议 140-180）
    python remove_watermark.py input.pdf --threshold 170

    # 批量处理目录下所有 PDF 和图片
    python remove_watermark.py /path/to/dir/ --batch

    # 预览单页/图片效果（不保存，直接展示对比图）
    python remove_watermark.py input.pdf --preview --page 0
    python remove_watermark.py input.jpg --preview
"""
import argparse
import sys
from pathlib import Path
from typing import Optional

# 将 ocr_platform 根目录加入 sys.path，以便导入 ocr_utils
_repo_root = Path(__file__).parents[2]
if str(_repo_root) not in sys.path:
    sys.path.insert(0, str(_repo_root))

from loguru import logger
from ocr_utils.watermark_utils import (
    detect_watermark,
    remove_watermark_from_image,
    scan_pdf_watermark_xobjs,
    remove_txt_pdf_watermark,
)

# 支持的图片后缀（小写）
IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".tif", ".tiff", ".bmp", ".webp"}


def _try_remove_txt_pdf_watermark(input_path: Path, output_path: Path) -> int:
    """
    对文字型 PDF 执行原生水印去除，保留文字可搜索性。

    内部委托给 watermark_utils.remove_txt_pdf_watermark() 完成内存流处理，
    有水印时将结果写入 output_path。

    流程：
    1. scan_pdf_watermark_xobjs() 快速扫描前 3 页，无水印直接返回 0
    2. remove_txt_pdf_watermark() 执行全量去除，返回 bytes 或 None
    3. 有水印时写 output_path

    Returns:
        1 表示去除成功，0 表示未发现水印
    """
    pdf_bytes = input_path.read_bytes()

    if not scan_pdf_watermark_xobjs(pdf_bytes, sample_pages=3):
        return 0

    cleaned = remove_txt_pdf_watermark(pdf_bytes)
    if cleaned is None:
        return 0

    output_path.write_bytes(cleaned)
    return 1



def process_document(
    input_path: Path,
    output_path: Path,
    threshold: int = 160,
    morph_close_kernel: int = 0,
    dpi: int = 200,
    page_range: Optional[str] = None,
    force_image: bool = False,
) -> int:
    """
    统一处理函数：支持 PDF（扫描件）和图片，去除水印后保存。

    使用 PDFUtils.load_and_classify_document 加载并分类：
    - 文字型 PDF（pdf_type='txt'）：优先尝试原生 XObject 水印去除（保留可搜索性）；
      失败时自动回退图像化处理，或 force_image=True 时直接走图像处理
    - 扫描件 PDF（pdf_type='ocr'）：逐页去水印后重新打包为 PDF
    - 图片：检测水印后去除并保存

    Args:
        input_path: 输入文件路径（PDF 或图片）
        output_path: 输出文件路径
        threshold: 灰度阈值（140-180），越大保守，越小激进
        morph_close_kernel: 形态学闭运算核大小，0 跳过
        dpi: PDF 渲染分辨率
        page_range: 页面范围字符串，如 "1-5,7,9-12"（从 1 开始，仅对 PDF 有效）
        force_image: 强制对文字型 PDF 使用图像化处理（会失去文字可搜索性，
                     但能处理水印嵌在内容流中的情况）

    Returns:
        实际处理的页/图片数
    """
    import shutil
    import numpy as np
    from io import BytesIO
    from PIL import Image
    from ocr_utils.pdf_utils import PDFUtils

    is_pdf = input_path.suffix.lower() == ".pdf"

    # 统一加载 + 分类（PDF 用 MinerU pdf_classify，图片直接读取）
    images, pdf_type, pdf_doc, renderer = PDFUtils.load_and_classify_document(
        input_path, dpi=dpi, page_range=page_range
    )

    # _known_has_wm: 当 txt 分支已确认有水印时设为 True，避免公共段用更严格阈值误判
    _known_has_wm: Optional[bool] = None

    # 文字型 PDF：优先尝试原生 XObject 水印去除，保留可搜索性
    if is_pdf and pdf_type == "txt" and not force_image:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        removed = _try_remove_txt_pdf_watermark(input_path, output_path)
        if removed > 0:
            logger.info(
                f"✅ 文字型 PDF '{input_path.name}'：删除 {removed} 个水印 XObject，"
                "保留文字可搜索性，已保存。"
            )
            return removed

        # XObject 扫描无结果，用较低阈值（0.5%）做图像水印检测二次确认
        # 文字 PDF 背景干净，降低阈值以检测稀疏文字水印
        first_np = np.array(images[0]["img_pil"])
        if detect_watermark(first_np, ratio_threshold=0.005):
            logger.warning(
                f"⚠️  文字型 PDF '{input_path.name}'：未找到 XObject 水印，"
                "但图像检测发现水印（内联内容流水印），"
                "回退为图像化处理（输出将失去文字可搜索性）。"
            )
            _known_has_wm = True  # 明确检测到水印，跳过公共段二次检测
        else:
            logger.info(
                f"✅ 文字型 PDF '{input_path.name}'：未检测到水印，直接复制。"
            )
            shutil.copy2(str(input_path), str(output_path))
            return 0
    elif is_pdf and pdf_type == "txt" and force_image:
        logger.warning(
            f"⚠️  文字型 PDF '{input_path.name}'：--force-image 模式，"
            "强制图像化处理（输出将失去文字可搜索性）。"
        )
        _known_has_wm = True  # force_image 模式不再检测，直接去除

    logger.info(
        f"{'📄' if is_pdf else '🖼️ '} 处理: {input_path.name}  "
        f"共 {len(images)} {'页' if is_pdf else '张'}  threshold={threshold}"
    )

    # 水印检测（仅用第一页/图判断，同一文档水印通常一致）
    # _known_has_wm 已在 txt 分支设置时，跳过重复检测
    if _known_has_wm is not None:
        has_wm = _known_has_wm
        logger.info("🔍 检测到水印，启动去水印处理" if has_wm else "✅ 未检测到水印，跳过")
    else:
        first_np = np.array(images[0]["img_pil"])
        # 扫描件/图片路径：使用宽松一档的中间调阈值（2.5%）以避免边界误判，
        # 斜向直线验证仍作为双重保险防止误报
        has_wm = detect_watermark(first_np, ratio_threshold=0.025)
        if has_wm:
            logger.info("🔍 检测到水印，启动去水印处理")
        else:
            logger.info("✅ 未检测到水印，跳过去水印处理")
            if not is_pdf:
                # 图片无水印：直接复制
                output_path.parent.mkdir(parents=True, exist_ok=True)
                shutil.copy2(str(input_path), str(output_path))
                return 1

    output_path.parent.mkdir(parents=True, exist_ok=True)

    if is_pdf:
        # 逐页处理后重新打包为 PDF
        try:
            import fitz
        except ImportError:
            raise ImportError("请安装 PyMuPDF: pip install PyMuPDF")

        new_doc = fitz.open()
        for i, img_dict in enumerate(images):
            pil_img = img_dict["img_pil"]
            img_np = np.array(pil_img)

            if has_wm:
                cleaned_gray = remove_watermark_from_image(
                    img_np, threshold=threshold,
                    morph_close_kernel=morph_close_kernel, return_pil=False,
                )
                out_pil = Image.fromarray(cleaned_gray).convert("RGB")
            else:
                out_pil = pil_img

            buf = BytesIO()
            out_pil.save(buf, format="PNG", optimize=False)
            buf.seek(0)

            # 按渲染图尺寸创建新页面（保持原始 DPI 尺寸）
            w_px, h_px = out_pil.size
            new_page = new_doc.new_page(width=w_px * 72 / dpi, height=h_px * 72 / dpi)
            new_page.insert_image(new_page.rect, stream=buf.read())

            if (i + 1) % 10 == 0 or i == len(images) - 1:
                logger.info(f"  进度: {i + 1}/{len(images)}")

        new_doc.save(str(output_path), garbage=4, deflate=True)
    else:
        # 图片：有水印则去除后保存
        img_np = np.array(images[0]["img_pil"])
        cleaned_gray = remove_watermark_from_image(
            img_np, threshold=threshold,
            morph_close_kernel=morph_close_kernel, return_pil=False,
        )
        Image.fromarray(cleaned_gray, mode="L").save(str(output_path))

    logger.info(f"✅ 保存到: {output_path}")
    return len(images)


def preview_page(
    input_path: Path,
    page_idx: int = 0,
    threshold: int = 160,
    dpi: int = 200,
):
    """展示单页原图与去水印对比（需要 matplotlib）。支持 PDF 和图片文件。"""
    try:
        import numpy as np
        import matplotlib.pyplot as plt
        import matplotlib
        matplotlib.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'SimHei', 'sans-serif']
        matplotlib.rcParams['axes.unicode_minus'] = False
    except ImportError as e:
        raise ImportError(f"预览需要 matplotlib: {e}")

    suffix = input_path.suffix.lower()

    if suffix == ".pdf":
        try:
            import fitz
        except ImportError:
            raise ImportError("PDF 预览需要 PyMuPDF: pip install PyMuPDF")
        doc = fitz.open(str(input_path))
        if page_idx >= len(doc):
            raise ValueError(f"页码 {page_idx} 超出范围（共 {len(doc)} 页）")
        mat = fitz.Matrix(dpi / 72, dpi / 72)
        page = doc[page_idx]
        pix = page.get_pixmap(matrix=mat, alpha=False)
        img_np = np.frombuffer(pix.samples, dtype=np.uint8).reshape(pix.h, pix.w, 3)
        title_orig = f"原图  第 {page_idx + 1} 页"
    elif suffix in IMAGE_SUFFIXES:
        from PIL import Image
        img_np = np.array(Image.open(str(input_path)).convert("RGB"))
        title_orig = f"原图  {input_path.name}"
    else:
        raise ValueError(f"不支持的文件格式: {suffix}")

    cleaned = remove_watermark_from_image(img_np, threshold=threshold, return_pil=False)

    fig, axes = plt.subplots(1, 2, figsize=(20, 14))
    axes[0].imshow(img_np)
    axes[0].set_title(title_orig, fontsize=14)
    axes[0].axis('off')

    axes[1].imshow(cleaned, cmap='gray')
    axes[1].set_title(f"去水印后  threshold={threshold}", fontsize=14)
    axes[1].axis('off')

    plt.tight_layout()
    plt.show()


def main():
    parser = argparse.ArgumentParser(
        description="银行流水水印去除工具",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    parser.add_argument("input", type=Path, help="输入 PDF / 图片文件或目录（批量模式）")
    parser.add_argument("-o", "--output", type=Path, default=None,
                        help="输出路径（单文件模式；默认在原文件名后加 _cleaned）")
    parser.add_argument("--threshold", type=int, default=160,
                        help="灰度阈值 (140-180)，默认 160")
    parser.add_argument("--morph-kernel", type=int, default=2,
                        help="形态学闭运算核大小，0 跳过，默认 2")
    parser.add_argument("--dpi", type=int, default=200,
                        help="渲染 DPI，默认 200")
    parser.add_argument("--batch", action="store_true",
                        help="批量模式：处理目录下所有 PDF 和图片")
    parser.add_argument("--preview", action="store_true",
                        help="预览模式：展示单页对比图（不保存）")
    parser.add_argument("--page", type=int, default=0,
                        help="预览页码（0-based），默认第 0 页")
    parser.add_argument("--page-range", type=str, default=None,
                        help="处理页面范围，如 '1-3,5,7-9'（从 1 开始，仅对 PDF 有效）")
    parser.add_argument("--force-image", action="store_true",
                        help="强制对文字型 PDF 使用图像化处理（会失去可搜索性，适用于 XObject 方法无法去除的内联水印）")

    args = parser.parse_args()

    if args.preview:
        preview_page(
            args.input,
            page_idx=args.page,
            threshold=args.threshold,
            dpi=args.dpi,
        )
        return

    if args.batch:
        # 批量模式：处理目录下所有 PDF 和图片
        input_dir = args.input
        if not input_dir.is_dir():
            logger.error(f"批量模式需要传入目录: {input_dir}")
            sys.exit(1)

        # 收集所有支持的文件
        all_files: list[Path] = sorted(input_dir.glob("*.pdf"))
        for ext in IMAGE_SUFFIXES:
            all_files.extend(sorted(input_dir.glob(f"*{ext}")))
            all_files.extend(sorted(input_dir.glob(f"*{ext.upper()}")))
        all_files = sorted(set(all_files))

        if not all_files:
            logger.warning(f"目录中没有可处理的文件（PDF/图片）: {input_dir}")
            return
        out_dir = args.output or input_dir / "cleaned"
        out_dir.mkdir(parents=True, exist_ok=True)
        for file in all_files:
            out_file = out_dir / f"{file.stem}_cleaned{file.suffix}"
            try:
                process_document(file, out_file, args.threshold, args.morph_kernel, args.dpi, args.page_range, args.force_image)
            except Exception as e:
                logger.error(f"❌ 处理失败 {file.name}: {e}")
        logger.info(f"✅ 批量处理完成，共 {len(all_files)} 个文件 -> {out_dir}")
    else:
        # 单文件模式
        input_path = args.input
        if not input_path.is_file():
            logger.error(f"文件不存在: {input_path}")
            sys.exit(1)
        output_path = args.output or input_path.with_name(
            f"{input_path.stem}_cleaned{input_path.suffix}"
        )
        suffix = input_path.suffix.lower()
        if suffix == ".pdf" or suffix in IMAGE_SUFFIXES:
            process_document(input_path, output_path, args.threshold, args.morph_kernel, args.dpi, args.page_range, args.force_image)
        else:
            logger.error(f"不支持的文件格式: {suffix}，支持 PDF 和 {IMAGE_SUFFIXES}")
            sys.exit(1)


if __name__ == "__main__":
    if len(sys.argv) == 1:
        print("ℹ️  未提供命令行参数，使用默认配置运行...")

        # 默认配置（用于开发测试）
        default_config = {
            # 测试输入
            # "input": "/Users/zhch158/workspace/data/流水分析/杨万益_福建农信.pdf",
            # "input": "Users/zhch158/workspace/data/流水分析/提取自杨万益_福建农信.png",
            
            # 文字PDF测试
            # "input": "/Users/zhch158/workspace/data/流水分析/提取自赤峰黄金2023年报.pdf",
            # "input": "/Users/zhch158/workspace/data/测试文字PDF-水印.pdf",
            "input": "/Users/zhch158/workspace/data/非结构化文档识别统一平台（ocr_platform）-交易流水识别，财报识别.pdf",
            # "output": "./output/杨万益_福建农信",
            # 页面范围（可选，支持 "1-5,7" 语法，仅对 PDF 有效）
            # "page_range": "3",  # 仅处理第 1 页（对应 --page-range 参数）
            "dpi": 200,
            "threshold": 160,
            "morph_kernel": 0,  # 遮罩替换模式下不需要闭运算
            # "preview": True,
        }

        # 构造参数（注意 input 是位置参数，morph_kernel 对应 --morph-kernel）
        sys.argv = [sys.argv[0], default_config["input"]]
        skip_keys = {"input"}
        for key, value in default_config.items():
            if key in skip_keys:
                continue
            # 将下划线转换为连字符（如 morph_kernel -> morph-kernel）
            flag = f"--{key.replace('_', '-')}"
            if isinstance(value, bool):
                if value:
                    sys.argv.append(flag)
            else:
                sys.argv.extend([flag, str(value)])

    sys.exit(main())