stamp.py

#!/usr/bin/env python3
"""平铺水印命令行工具。

用法示例:
    # 位置参数指定输入/输出，其余用缩写选项（图片或视频均可）
    python signature.py in.png out.jpg -s "版权所有" -m 10 -r 45
    python signature.py in.mp4 out.mp4 -s "版权所有" -m 10 -r 45

    # 指定字体大小与间隙（均支持 px / pt 单位）
    python signature.py in.png out.jpg -s "版权所有" --font-size 64px -m 8pt

    # 等价的全写选项形式
    python signature.py --input in.png --output out.jpg --signature "版权所有"
"""

import argparse
import os
import platform
import shutil
import subprocess
import sys
import tempfile

from PIL import Image, ImageDraw, ImageFont

VIDEO_EXTS = {
    ".mp4", ".mkv", ".mov", ".avi", ".wmv", ".flv", ".webm", ".m4v",
    ".mpg", ".mpeg", ".ts", ".m2ts", ".3gp", ".ogv",
}


def parse_color(value: str) -> tuple[int, int, int]:
    """解析十六进制颜色字符串（如 #123DF / #123DEF），大小写不敏感。

    返回 (r, g, b)。格式非法时抛 ValueError。
    """
    s = value.strip().lstrip("#")
    if len(s) not in (3, 6):
        raise ValueError(f"颜色值必须是 #RGB 或 #RRGGBB 格式（含或不含 #），当前为: {value!r}")
    try:
        if len(s) == 3:
            s = "".join(ch * 2 for ch in s)  # #RGB -> #RRGGBB
        rgb = bytes.fromhex(s)
    except ValueError:
        raise ValueError(f"颜色值含非法十六进制字符: {value!r}")
    return rgb[0], rgb[1], rgb[2]


def parse_size(value: str, dpi: float = 96.0) -> int:
    """解析长度数值，支持 px 或 pt 单位。

    - "64" / "64px"：64 像素
    - "48pt"：按 dpi 换算为像素，px = pt * dpi / 72，默认 dpi=96
    非法格式或 ≤0 抛 ValueError。
    """
    s = value.strip().lower()
    try:
        if s.endswith("px"):
            num = float(s[:-2])
        elif s.endswith("pt"):
            num = float(s[:-2]) * dpi / 72.0
        else:
            num = float(s)
    except ValueError:
        raise ValueError(f"无法解析数值: {value!r}")
    px = round(num)
    if px <= 0:
        raise ValueError(f"数值必须为正数: {value!r}")
    return px


def load_font(path: str | None, size: int) -> ImageFont.FreeTypeFont:
    """加载水印字体。未指定时按平台选默认字体。

    Windows 优先微软雅黑（PIL 用 Arial 渲染中文会乱码），其它系统用 Arial。
    """
    if path:
        if not os.path.exists(path):
            sys.exit(f"错误：指定的字体文件不存在: {path}")
        return ImageFont.truetype(path, size)

    candidates: list[str]
    if platform.system() == "Windows":
        candidates = [
            "C:/Windows/Fonts/msyh.ttc",  # 微软雅黑
            "C:/Windows/Fonts/msyh.ttf",
            "C:/Windows/Fonts/simhei.ttf",  # 黑体
            "C:/Windows/Fonts/simsun.ttc",  # 宋体
        ]
    else:
        candidates = ["/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf", "/System/Library/Fonts/Helvetica.ttc"]

    for cand in candidates:
        if os.path.exists(cand):
            return ImageFont.truetype(cand, size)

    return ImageFont.load_default(size)


def resolve_font_size(args, width: int, height: int, dpi: float = 96.0) -> int:
    """确定字体像素大小。

    若指定了 --font-size（px/pt），则解析它并返回；
    否则按宽度自适应（width // 20，下限 16px）。
    """
    if args.font_size:
        try:
            return parse_size(args.font_size, dpi)
        except ValueError as e:
            raise ValueError(f"--font-size {e}") from e
    return max(16, width // 20)


def resolve_margin(args, dpi: float = 96.0) -> int:
    """确定水印间隙（像素），--margin 支持 px/pt 单位。"""
    try:
        return parse_size(args.margin, dpi)
    except ValueError as e:
        raise ValueError(f"--margin {e}") from e


def build_stamp(text: str, font: ImageFont.FreeTypeFont, color_rgb, alpha: int, rotation: float) -> Image.Image:
    """创建并返回一个旋转后的水印 stamp（含透明通道）。"""
    bbox = font.getbbox(text)
    text_w = bbox[2] - bbox[0]
    text_h = bbox[3] - bbox[1]
    stamp = Image.new("RGBA", (max(text_w, 1), max(text_h, 1)), (0, 0, 0, 0))
    draw = ImageDraw.Draw(stamp)
    draw.text((-bbox[0], -bbox[1]), text, font=font, fill=(*color_rgb, alpha))
    return stamp.rotate(-rotation, expand=True, resample=Image.BICUBIC)


def render_tiled_watermark(
    width: int,
    height: int,
    text: str,
    font: ImageFont.FreeTypeFont,
    color_rgb,
    alpha: int,
    rotation: float,
    margin: int,
) -> Image.Image:
    """在 (width, height) 的画布上平铺旋转水印，返回透明 RGBA overlay。

    采用交错砖墙错位排布，填补旋转产生的空隙。
    """
    stamp = build_stamp(text, font, color_rgb, alpha, rotation)
    overlay = Image.new("RGBA", (width, height), (0, 0, 0, 0))
    step_x = stamp.width + margin
    step_y = stamp.height + margin
    if step_x <= 0 or step_y <= 0:
        raise ValueError("水印尺寸与 margin 计算出的步长无效")

    row = 0
    y = -stamp.height
    while y < height:
        offset = (step_x / 2) if (row % 2) else 0
        x = -stamp.width + offset
        while x < width:
            overlay.paste(stamp, (int(x), int(y)), stamp)
            x += step_x
        y += step_y
        row += 1
    return overlay


def save_without_compression(im: Image.Image, path: str) -> None:
    """按输出扩展名选择最接近"不压缩"的保存参数。

    - PNG: compress_level=0（zlib 不压缩，store）
    - JPEG: 无损模式不存在，尽量用 quality=100 + 无子采样，减少损失
    - TGA: 支持透明通道，且 PIL 默认即不压缩，保留 RGBA 直接保存
    - BMP / 其它: 交给 PIL 默认处理

    JPEG 不支持透明通道，若输入是 RGBA 则先合成到白底再转 RGB。
    """
    ext = os.path.splitext(path)[1].lower()
    if ext in (".jpg", ".jpeg"):
        if im.mode == "RGBA":
            bg = Image.new("RGB", im.size, (255, 255, 255))
            bg.paste(im, mask=im.split()[3])
            im = bg
        elif im.mode != "RGB":
            im = im.convert("RGB")
        im.save(path, quality=100, subsampling=0)
    elif ext in (".png", ".apng"):
        im.save(path, compress_level=0)
    elif ext in (".tga",):
        im.save(path)  # TGA 原生支持 alpha，且默认不压缩
    else:
        im.save(path)


def is_video(path: str) -> bool:
    """根据扩展名判断输入是否为视频文件。"""
    return os.path.splitext(path)[1].lower() in VIDEO_EXTS


def probe_video_size(path: str) -> tuple[int, int] | None:
    """用 ffprobe 获取视频首帧尺寸 (width, height)。失败返回 None。"""
    ffprobe = shutil.which("ffprobe")
    if not ffprobe:
        return None
    cmd = [
        ffprobe, "-v", "error", "-select_streams", "v:0",
        "-show_entries", "stream=width,height", "-of", "csv=p=0:s=x", path,
    ]
    try:
        out = subprocess.run(cmd, capture_output=True, text=True, check=True).stdout.strip()
        w, h = out.split("x")
        return int(w), int(h)
    except Exception:
        return None


def run_ffmpeg_overlay(video_in: str, wm_png: str, video_out: str) -> None:
    """调用 ffmpeg 将平铺水印 PNG 叠加到视频上。

    输出用 libx264 CRF 0 无损编码，音频直接复制。
    """
    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        sys.exit("错误：未找到 ffmpeg 命令，无法处理视频")

    filter_complex = (
        f"[1:v]format=rgba[wm];"
        f"[0:v][wm]overlay=0:0:format=auto[v]"
    )
    cmd = [
        ffmpeg, "-y",
        "-i", video_in,
        "-i", wm_png,
        "-filter_complex", filter_complex,
        "-map", "[v]", "-map", "0:a?",
        "-c:v", "libx264", "-crf", "0", "-preset", "slow",
        "-pix_fmt", "yuv420p",
        "-c:a", "copy",
        video_out,
    ]
    print("调用 ffmpeg 处理视频…")
    subprocess.run(cmd, check=True)


def process_video(args) -> None:
    """视频输入分支：生成平铺水印 PNG，再交给 ffmpeg 叠加。"""
    ffprobe = shutil.which("ffprobe")
    if not ffprobe:
        sys.exit("错误：未找到 ffprobe 命令，无法读取视频尺寸")

    size = probe_video_size(args.input)
    if size is None:
        sys.exit(f"错误：无法用 ffprobe 读取视频尺寸: {args.input}")
    width, height = size

    font_size = resolve_font_size(args, width, height)
    font = load_font(args.font, font_size)
    try:
        margin = resolve_margin(args)
        overlay = render_tiled_watermark(
            width, height, args.signature, font, args.color_rgb, args.alpha,
            args.rotation, margin,
        )
    except ValueError as e:
        sys.exit(f"错误：{e}")

    # 生成临时水印 PNG
    fd, wm_path = tempfile.mkstemp(suffix=".png")
    os.close(fd)
    try:
        overlay.save(wm_path, compress_level=0)
        out_dir = os.path.dirname(os.path.abspath(args.output))
        os.makedirs(out_dir, exist_ok=True)
        run_ffmpeg_overlay(args.input, wm_path, args.output)
    finally:
        os.remove(wm_path)
    print(f"完成：已输出到 {args.output}")


def main() -> None:
    parser = argparse.ArgumentParser(description="在图片上平铺旋转水印")
    parser.add_argument("input", nargs="?", default=None, help="输入图片路径（位置参数）")
    parser.add_argument("output", nargs="?", default=None, help="输出路径（位置参数，父目录不存在会自动创建）")
    parser.add_argument("-i", "--input-path", dest="input_path", default=None, help="输入图片路径（选项形式，与位置参数 input 二选一）")
    parser.add_argument("-o", "--output-path", dest="output_path", default=None, help="输出路径（选项形式，与位置参数 output 二选一）")
    parser.add_argument("-s", "--signature", required=True, help="水印文本")
    parser.add_argument("-m", "--margin", default="10", help="每个水印 BoundingBox 之间的间隙，支持 px/pt，如 10 / 10px / 8pt，默认 10px")
    parser.add_argument("-r", "--rotation", type=float, default=45.0, help="水印旋转角度（度），默认 45")
    parser.add_argument("-c", "--color", default="#FFFFFF", help="水印颜色，十六进制（#RGB 或 #RRGGBB），大小写不敏感，默认 #FFFFFF")
    parser.add_argument("-p", "--opacity", type=int, default=80, help="水印不透明度 1-100，默认 80")
    parser.add_argument("-f", "--font", default=None, help="水印字体文件路径（默认按平台自动选择）")
    parser.add_argument("--font-size", default=None, help="字体大小，支持像素或磅：如 64 / 64px / 48pt（不指定则按宽度自适应）")
    args = parser.parse_args()

    # 解析输入：选项形式优先，否则用位置参数
    input_path = args.input_path or args.input
    output_path = args.output_path or args.output
    if not input_path:
        parser.error("缺少输入图片（可用位置参数，或 -i/--input-path）")
    if not output_path:
        parser.error("缺少输出路径（可用位置参数，或 -o/--output-path）")
    args.input = input_path
    args.output = output_path

    # 校验颜色
    try:
        color_rgb = parse_color(args.color)
    except ValueError as e:
        parser.error(str(e))
    args.color_rgb = color_rgb

    # 校验不透明度
    if not (1 <= args.opacity <= 100):
        parser.error(f"--opacity 必须在 1-100 之间，当前为: {args.opacity}")
    args.alpha = round(args.opacity / 100 * 255)

    # 校验输入
    if not os.path.exists(args.input):
        parser.error(f"输入文件不存在: {args.input}")

    # 视频输入：转交给 ffmpeg 处理
    if is_video(args.input):
        process_video(args)
        return

    # 加载图片
    im = Image.open(args.input)
    has_alpha = "A" in im.mode
    im = im.convert("RGBA")

    width, height = im.size
    dpi = (im.info.get("dpi") or (96.0,))[0]
    try:
        font_size = resolve_font_size(args, width, height, dpi)
    except ValueError as e:
        parser.error(str(e))
    font = load_font(args.font, font_size)
    try:
        margin = resolve_margin(args, dpi)
        overlay = render_tiled_watermark(
            width, height, args.signature, font, args.color_rgb, args.alpha,
            args.rotation, margin,
        )
    except ValueError as e:
        parser.error(str(e))

    result = Image.alpha_composite(im, overlay)
    if not has_alpha:
        result = result.convert("RGB")

    # 自动创建输出目录
    out_dir = os.path.dirname(os.path.abspath(args.output))
    os.makedirs(out_dir, exist_ok=True)
    save_without_compression(result, args.output)
    print(f"完成：已输出到 {args.output}")


if __name__ == "__main__":
    main()