R524809
123 lines
5.2 KiB
Python
123 lines
5.2 KiB
Python
"""生成图水印：AI 出图返回后、落盘前的后处理合成（不经过生图模型）。
|
||
|
||
样式复刻 ozonSeller「图表处理」的默认水印：
|
||
- 图片水印：徽章图中心裁方 → 圆形遮罩 → 宽度为图宽 15%，右下角，边距约 1% 图宽；
|
||
- 文字水印：字号为图宽 6%（下限 12px），白色填充 + 黑色描边（alpha 0.55，
|
||
描边宽 fontSize/8），加粗无衬线。
|
||
容错原则：字体/资产缺失或合成异常时 log 警告并返回原图，绝不阻断生图。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import io
|
||
import logging
|
||
from pathlib import Path
|
||
|
||
from PIL import Image, ImageDraw, ImageFont
|
||
|
||
from config import get_settings
|
||
|
||
log = logging.getLogger("suite.watermark")
|
||
|
||
# 尺寸比例（与 ozonSeller app.js 常量一致）
|
||
BADGE_SCALE = 0.15 # 图片水印直径 / 图宽
|
||
BADGE_MARGIN = 0.01 # 图片水印边距 / 图宽（ozonSeller 固定 10px，按比例更稳）
|
||
TEXT_SCALE = 0.06 # 文字字号 / 图宽
|
||
TEXT_MIN_SIZE = 12
|
||
STROKE_ALPHA = 0.55
|
||
STROKE_RATIO = 1 / 8 # 描边宽 / 字号
|
||
|
||
# CJK/西文字体回退链（macOS 本地服务）；命中后模块级缓存
|
||
_FONT_CANDIDATES = [
|
||
"/System/Library/Fonts/PingFang.ttc",
|
||
"/System/Library/Fonts/Hiragino Sans GB.ttc",
|
||
"/System/Library/Fonts/STHeiti Light.ttc",
|
||
"/Library/Fonts/Arial Unicode.ttf",
|
||
]
|
||
_font_path: str | None = None
|
||
|
||
|
||
def _load_font(size: int) -> ImageFont.FreeTypeFont | ImageFont.ImageFont:
|
||
global _font_path
|
||
if _font_path is None:
|
||
_font_path = next((p for p in _FONT_CANDIDATES if Path(p).is_file()), "")
|
||
if _font_path:
|
||
return ImageFont.truetype(_font_path, size)
|
||
log.warning("未找到系统字体（%s），文字水印退化为 Pillow 默认字体，中文可能乱码", _FONT_CANDIDATES)
|
||
return ImageFont.load_default(size) if size >= 10 else ImageFont.load_default()
|
||
|
||
|
||
def _circular_badge(size: int) -> Image.Image | None:
|
||
"""徽章资产 → 指定直径的圆形 RGBA 贴片；资产缺失返回 None。"""
|
||
path = get_settings().watermark_image_path
|
||
try:
|
||
badge = Image.open(path).convert("RGBA")
|
||
except Exception as exc: # noqa: BLE001
|
||
log.warning("水印图片加载失败（%s），跳过图片水印: %s", path, exc)
|
||
return None
|
||
side = min(badge.size) # 中心裁方
|
||
left, top = (badge.width - side) // 2, (badge.height - side) // 2
|
||
square = badge.crop((left, top, left + side, top + side)).resize((size, size))
|
||
mask = Image.new("L", (size, size), 0)
|
||
ImageDraw.Draw(mask).ellipse((0, 0, size - 1, size - 1), fill=255)
|
||
square.putalpha(mask)
|
||
return square
|
||
|
||
|
||
def _apply_image_watermark(canvas: Image.Image, opacity: float) -> None:
|
||
size = max(24, round(canvas.width * BADGE_SCALE))
|
||
badge = _circular_badge(size)
|
||
if badge is None:
|
||
return
|
||
badge.putalpha(badge.getchannel("A").point(lambda a: round(a * opacity)))
|
||
margin = max(10, round(canvas.width * BADGE_MARGIN))
|
||
canvas.alpha_composite(badge, (canvas.width - size - margin, canvas.height - size - margin))
|
||
|
||
|
||
def _apply_text_watermark(canvas: Image.Image, text: str, opacity: float) -> None:
|
||
text = (text or "").strip()
|
||
if not text:
|
||
return
|
||
font_size = max(TEXT_MIN_SIZE, round(canvas.width * TEXT_SCALE))
|
||
font = _load_font(font_size)
|
||
layer = Image.new("RGBA", canvas.size, (0, 0, 0, 0))
|
||
draw = ImageDraw.Draw(layer)
|
||
bbox = draw.textbbox((0, 0), text, font=font, stroke_width=max(1, round(font_size * STROKE_RATIO)))
|
||
tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
|
||
if tw >= canvas.width: # 文案比图还宽：按比例缩字号重排一次
|
||
font_size = max(TEXT_MIN_SIZE, round(font_size * canvas.width / tw * 0.94))
|
||
font = _load_font(font_size)
|
||
bbox = draw.textbbox((0, 0), text, font=font, stroke_width=max(1, round(font_size * STROKE_RATIO)))
|
||
tw, th = bbox[2] - bbox[0], bbox[3] - bbox[1]
|
||
margin = max(10, round(canvas.width * BADGE_MARGIN))
|
||
x = canvas.width - tw - margin - bbox[0]
|
||
y = canvas.height - th - margin - bbox[1]
|
||
stroke = (0, 0, 0, round(255 * STROKE_ALPHA))
|
||
fill = (255, 255, 255, 255)
|
||
draw.text((x, y), text, font=font, fill=fill, stroke_width=max(1, round(font_size * STROKE_RATIO)),
|
||
stroke_fill=stroke)
|
||
layer.putalpha(layer.getchannel("A").point(lambda a: round(a * opacity)))
|
||
canvas.alpha_composite(layer)
|
||
|
||
|
||
def apply_watermark(data: bytes, opts: dict) -> bytes:
|
||
"""给图片字节加水印，返回同格式字节；opts: {type, text, opacity(0-100)}。"""
|
||
is_png = data[:8] == b"\x89PNG\r\n\x1a\n"
|
||
fmt = "PNG" if is_png else "JPEG"
|
||
try:
|
||
img = Image.open(io.BytesIO(data))
|
||
canvas = img.convert("RGBA")
|
||
opacity = min(100, max(1, int(opts.get("opacity") or 30))) / 100
|
||
if opts.get("type") == "text":
|
||
_apply_text_watermark(canvas, opts.get("text") or "", opacity)
|
||
else:
|
||
_apply_image_watermark(canvas, opacity)
|
||
out = io.BytesIO()
|
||
if fmt == "PNG":
|
||
canvas.save(out, format="PNG")
|
||
else:
|
||
canvas.convert("RGB").save(out, format="JPEG", quality=95)
|
||
return out.getvalue()
|
||
except Exception: # noqa: BLE001
|
||
log.exception("水印合成失败，返回原图")
|
||
return data
|