ocr_platform/ocr_tools/gan_experiments_lab/watermark_synthesis.py

"""
水印合成脚本：在clean图片上叠加斜向浅色文字水印，输出带水印图 + 精确mask。

用法:
    python watermark_synthesis.py                          # 默认参数演示
    python watermark_synthesis.py --input ./test_images/clean/   # 指定输入目录
    python watermark_synthesis.py --text "SAMPLE" --opacity 0.15 --angle 45
"""
from __future__ import annotations

import argparse
import math
from pathlib import Path
from typing import Optional

import cv2
import numpy as np
from loguru import logger
from PIL import Image, ImageDraw, ImageFont


def _find_font() -> str:
    """查找可用中文字体，找不到返回默认字体。"""
    candidates = [
        "/System/Library/Fonts/PingFang.ttc",
        "/System/Library/Fonts/STHeiti Light.ttc",
        "/System/Library/Fonts/Hiragino Sans GB.ttc",
        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
        "/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc",
    ]
    for fp in candidates:
        if Path(fp).exists():
            return fp
    logger.warning("未找到中文字体，使用PIL默认字体")
    return ""


def _text_size_to_font_size(text_height_px: int) -> int:
    """根据目标文字像素高度估算 font_size。"""
    return int(text_height_px * 1.15)


def _render_watermark_tile(
    pil_img: Image.Image,
    text: str,
    font_path: str,
    font_size: int,
    opacity: float,
    angle_deg: float,
    spacing_x: int,
    spacing_y: int,
) -> tuple[np.ndarray, np.ndarray]:
    """
    在图上平铺斜向水印文字，返回 (watermarked_np, mask_np)。

    mask_np: H×W bool, True=水印像素位置。
    """
    w, h = pil_img.size
    text_height = int(font_size / 1.15)
    gray_value = int(255 * (1 - opacity))

    # 创建水印文字mask（稍大画布以覆盖旋转后区域）
    diag = int(math.sqrt(w * w + h * h)) + text_height * 4
    tile_w = diag
    tile_h = diag

    tile = Image.new("L", (tile_w, tile_h), 0)
    draw = ImageDraw.Draw(tile)
    font = ImageFont.truetype(font_path, font_size) if font_path else ImageFont.load_default()

    # 步长取spacing + 文字大小，确保均匀分布
    step_x = text_height + spacing_x
    step_y = text_height + spacing_y

    for y in range(0, tile_h, step_y):
        for x in range(0, tile_w, step_x):
            draw.text((x, y), text, fill=255, font=font)

    # 旋转
    tile_rot = tile.rotate(angle_deg, expand=False, fillcolor=0)

    # 裁剪到原图大小（中心对齐）
    cx, cy = tile_rot.size[0] // 2, tile_rot.size[1] // 2
    left = cx - w // 2
    top = cy - h // 2
    watermark_tile = tile_rot.crop((left, top, left + w, top + h))

    mask_np = np.array(watermark_tile) > 0

    # 叠加到原图
    base = np.array(pil_img.convert("RGB"))
    alpha = opacity
    result = base.copy()
    result[mask_np] = (
        base[mask_np].astype(np.float32) * (1 - alpha)
        + np.array([gray_value, gray_value, gray_value], dtype=np.float32) * alpha
    ).astype(np.uint8)

    return result, mask_np


def synthesize_watermark(
    input_path: Path,
    output_dir: Path,
    *,
    text: str = "SAMPLE",
    font_path: str = "",
    text_height_px: int = 36,
    opacity: float = 0.12,
    angle_deg: float = 45.0,
    spacing_x: int = 180,
    spacing_y: int = 180,
    save_mask: bool = True,
) -> Path:
    """
    在输入图片上合成水印，输出到 output_dir。

    Returns:
        合成后的图片路径
    """
    output_dir.mkdir(parents=True, exist_ok=True)
    pil_img = Image.open(str(input_path)).convert("RGB")

    fp = font_path or _find_font()
    font_size = _text_size_to_font_size(text_height_px)

    logger.info(
        f"合成水印: {input_path.name} | "
        f"text='{text}' font_size={font_size} opacity={opacity} angle={angle_deg}"
    )

    result_np, mask_np = _render_watermark_tile(
        pil_img, text, fp, font_size, opacity, angle_deg, spacing_x, spacing_y
    )

    out_name = f"{input_path.stem}_watermarked{input_path.suffix}"
    out_path = output_dir / out_name
    Image.fromarray(result_np).save(str(out_path))
    logger.info(f"  水印图: {out_path}")

    if save_mask:
        mask_path = output_dir / f"{input_path.stem}_mask.png"
        cv2.imwrite(str(mask_path), (mask_np.astype(np.uint8) * 255))
        logger.info(f"  mask:   {mask_path}")

    return out_path


def main():
    parser = argparse.ArgumentParser(description="水印合成工具")
    parser.add_argument("--input", type=Path, default=None,
                        help="输入图片或目录（默认: test_images/clean/）")
    parser.add_argument("--output", type=Path, default=None,
                        help="输出目录（默认: test_images/synthetic/）")
    parser.add_argument("--text", type=str, default="行内内部使用",
                        help="水印文字内容")
    parser.add_argument("--text-height", type=int, default=48,
                        help="文字像素高度（默认48）")
    parser.add_argument("--opacity", type=float, default=0.10,
                        help="水印透明度 0~1（默认0.10）")
    parser.add_argument("--angle", type=float, default=45.0,
                        help="水印倾斜角度（默认45°）")
    parser.add_argument("--spacing-x", type=int, default=250,
                        help="水印文字水平间距（默认250px）")
    parser.add_argument("--spacing-y", type=int, default=250,
                        help="水印文字垂直间距（默认250px）")
    parser.add_argument("--font", type=str, default="",
                        help="字体文件路径")
    parser.add_argument("--no-mask", action="store_true",
                        help="不保存mask")
    parser.add_argument("--demo", action="store_true",
                        help="使用input目录下第一张测试图生成演示图")
    args = parser.parse_args()

    root = Path(__file__).parent
    input_dir = args.input or (root / "test_images" / "clean")
    output_dir = args.output or (root / "test_images" / "synthetic")

    if args.demo:
        # 无clean图时，直接用input目录的水印图再加一层合成水印做演示
        img_files = sorted(root.glob("test_images/input/*"))
        if not img_files:
            logger.error("test_images/input/ 下没有测试图片，请放入图片后重试")
            return
        input_dir = root / "test_images" / "input"
        output_dir = root / "test_images" / "synthetic"

    input_dir = Path(input_dir)
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)

    if input_dir.is_dir():
        img_files = sorted([
            f for f in input_dir.iterdir()
            if f.suffix.lower() in {".png", ".jpg", ".jpeg", ".bmp", ".tif", ".tiff", ".webp"}
        ])
    elif input_dir.is_file():
        img_files = [input_dir]
    else:
        logger.error(f"输入路径不存在: {input_dir}")
        return

    if not img_files:
        logger.warning(f"{input_dir} 下没有图片文件")
        return

    for f in img_files:
        synthesize_watermark(
            f, output_dir,
            text=args.text,
            font_path=args.font,
            text_height_px=args.text_height,
            opacity=args.opacity,
            angle_deg=args.angle,
            spacing_x=args.spacing_x,
            spacing_y=args.spacing_y,
            save_mask=not args.no_mask,
        )


if __name__ == "__main__":
    main()