beige/backend/app/services/ai_engine/image_postprocessor.py

"""
图片后处理（去AI化主路）
对齐大卫 xhs-tool/backend/infrastructure/imagePostProcess.js（运营实测去AI化版）。

主路 = 尺寸可选(±2%容差内不resize) + SynthID破除(可选) + 高保真重编码去元数据。

诚实声明：C2PA 元数据可去除；私有像素水印（如 SynthID）只能削弱，不保证 100% 清除。
"""
from __future__ import annotations

import io
import logging
import os

logger = logging.getLogger(__name__)

try:
    from PIL import Image, ImageEnhance, ImageOps
    _PILLOW_OK = True
except ImportError:
    _PILLOW_OK = False
    logger.warning("Pillow 未安装，image_postprocessor 不可用")

# 比例映射表，对齐大卫 RATIO_MAP。key 为字符串如 '3:4'
RATIO_MAP: dict[str, tuple[int, int]] = {
    "1:1":  (1024, 1024),
    "3:4":  (1024, 1536),   # gpt-image-2 原生尺寸，默认
    "4:3":  (1536, 1024),
    "9:16": (864,  1536),
    "16:9": (1536, 864),
}

# ±2% 容差内不做 resize，避免无谓重采样（对齐大卫 diff > 0.02 才 resize）
_RATIO_TOLERANCE = 0.02


def _need_resize(actual_w: int, actual_h: int, target_w: int, target_h: int) -> bool:
    """判断实际比例与目标比例差距是否超出容差。"""
    actual_ratio = actual_w / actual_h
    target_ratio = target_w / target_h
    diff = abs(actual_ratio - target_ratio) / target_ratio
    return diff > _RATIO_TOLERANCE


def process_image(
    image_bytes: bytes,
    aspect_ratio: str = "3:4",
    resample_strength: int = 1,  # 0=不重采样, 1=轻采样(默认), 2=重采样
) -> bytes:
    """
    处理单张图片。

    参数：
        image_bytes    — 原始图片 bytes（PNG/JPEG/WebP 等）
        aspect_ratio   — 目标比例，取 RATIO_MAP 的 key，默认 '3:4'=1024×1536
        resample_strength — 轻重采样削像素水印（0/1/2），默认 1=轻采样

    返回 JPEG bytes（无 EXIF/C2PA/XMP 元数据）。
    失败时降级返回原图 bytes，不抛异常（对齐大卫 catch 返回原图）。
    """
    if not _PILLOW_OK:
        logger.error("Pillow 未安装，跳过后处理，返回原图")
        return image_bytes

    try:
        img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        actual_w, actual_h = img.size
        target = RATIO_MAP.get(aspect_ratio)

        # --- Step1: 尺寸对齐（±2% 容差内跳过 resize）---
        if target:
            tw, th = target
            if _need_resize(actual_w, actual_h, tw, th):
                img = ImageOps.fit(img, (tw, th), method=Image.LANCZOS)
                logger.debug("resize %dx%d → %dx%d (ratio=%s)", actual_w, actual_h, tw, th, aspect_ratio)

        # --- Step2: resample_strength 削像素水印（可选，默认轻采样）---
        img = _apply_resample(img, resample_strength)

        # --- Step3: SynthID 破除（SYNTHID_HARD_MODE=1 才开，默认关）---
        if os.environ.get("SYNTHID_HARD_MODE") == "1" and target:
            img = _apply_synthid_break(img, target)

        # --- Step4: 高保真 JPEG 重编码，去所有元数据 ---
        buf = io.BytesIO()
        img.save(
            buf,
            format="JPEG",
            quality=100,
            subsampling=0,   # 4:4:4 chroma
            optimize=True,
            # 不传 exif/icc_profile/xmp = 不写入任何元数据
        )
        result = buf.getvalue()
        logger.debug("后处理完成 %d B → %d B (ratio=%s)", len(image_bytes), len(result), aspect_ratio)
        return result

    except Exception as exc:
        logger.warning("图片后处理失败，降级返回原图: %s", exc)
        return image_bytes


def _apply_resample(img: "Image.Image", strength: int) -> "Image.Image":
    """
    轻/重采样削像素级水印（resample_strength 控制）。
    0 — 不采样，仅靠重编码去元数据。
    1 — 轻采样：缩98%再回原尺寸，保视觉质量，削弱像素水印（对齐旧逻辑）。
    2 — 重采样：两次缩放，削弱更多，轻微质量损失。
    """
    if strength < 1:
        return img
    w, h = img.size
    img = img.resize((int(w * 0.98), int(h * 0.98)), Image.LANCZOS)
    img = img.resize((w, h), Image.LANCZOS)
    if strength >= 2:
        img = img.resize((int(w * 0.96), int(h * 0.96)), Image.LANCZOS)
        img = img.resize((w, h), Image.LANCZOS)
    return img


def _apply_synthid_break(img: "Image.Image", target: tuple[int, int]) -> "Image.Image":
    """
    SynthID 破除（SYNTHID_HARD_MODE=1 时调用）：
    对齐大卫逻辑 — 缩到(w-2,h-2)再裁掉1px边 + 亮度*1.005/饱和*0.998。
    诚实声明：只能削弱 SynthID，不保证 100% 清除。
    """
    tw, th = target
    img = ImageOps.fit(img, (tw - 2, th - 2), method=Image.LANCZOS)
    # 裁掉1px边，消除边缘水印残留
    img = img.crop((1, 1, tw - 3, th - 3))
    # 微调亮度/饱和（对齐大卫 modulate brightness/saturation）
    img = ImageEnhance.Brightness(img).enhance(1.005)
    img = ImageEnhance.Color(img).enhance(0.998)
    return img


def batch_process(
    images: list[bytes],
    aspect_ratio: str = "3:4",
    resample_strength: int = 1,
) -> list[dict]:
    """
    批量后处理。返回 [{index, data, error}]，单张失败不阻塞其余。
    """
    results = []
    for i, img_bytes in enumerate(images):
        try:
            processed = process_image(img_bytes, aspect_ratio=aspect_ratio,
                                      resample_strength=resample_strength)
            results.append({"index": i, "data": processed, "error": None})
        except Exception as exc:
            logger.error("图片[%d]后处理失败: %s", i, exc)
            results.append({"index": i, "data": img_bytes, "error": str(exc)})
    return results


async def gemini_rewatermark_fallback(
    client: "Any",  # GeminiClient，由 worker 注入
    image_bytes: bytes,
) -> bytes:
    """
    备选路：Gemini 重绘去水印。
    ⚠️ 对海报中文大字有改字风险，仅特殊场景启用。
    """
    prompt = (
        "Remove all watermarks, text overlays, and digital signatures from this image. "
        "Reconstruct any covered areas naturally to match the surrounding content. "
        "Return a clean version of the same image without any watermarks."
    )
    try:
        result = await client.gemini_generate(
            prompt, [image_bytes], "gemini-2.0-flash-preview-image-generation"
        )
        return result
    except Exception as exc:
        logger.error("Gemini 去水印失败，降级返回原图: %s", exc)
        return image_bytes