""" 银行流水水印去除工具 支持 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())