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 不可用")

# 北哥暖棕调色层（A4），自带开关与降级
from .beige_color_grade import grade as _beige_grade, is_enabled as _beige_enabled

# 比例映射表，对齐大卫 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+3: 去AI化破水印（硬路默认开，倩倩姐2026-06-26拍板，对齐大卫xhs实测版）---
        # 硬路 = 缩2px裁1px边错位采样 + 亮度/饱和微调，专破 SynthID 像素水印。
        # 「像素不能有任何压缩」红线：大卫硬路只动边缘2px，没有全图来回缩放；
        # Clover 原 Step2（缩98%再放大）是额外的全图 LANCZOS 重采样、会软化全图像素，
        # 故硬路开启时直接跳过 Step2，只走硬路局部错位——既破水印又不压全图像素。
        hard_mode = os.environ.get("SYNTHID_HARD_MODE", "1") != "0"
        if hard_mode:
            # 硬路只动边缘(缩2px裁1px)，用图片当前尺寸即可，不依赖 RATIO_MAP；
            # 故非标准比例(target=None)也走硬路，绝不降级到全图 LANCZOS 有损重采样。
            img = _apply_synthid_break(img)
        else:
            # 仅硬路显式关闭时，才用全图轻采样兜底破水印
            img = _apply_resample(img, resample_strength)

        # --- Step3.5: 北哥暖棕调色（BEIGE_COLOR_GRADE 默认开，去AI化核心）---
        # 在重编码前以 bytes 往返：grade 内部自带开关/降级，关则原样透传
        if _beige_enabled():
            buf0 = io.BytesIO()
            img.save(buf0, format="JPEG", quality=100, subsampling=0)
            graded = _beige_grade(buf0.getvalue())
            img = Image.open(io.BytesIO(graded)).convert("RGB")

        # --- 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") -> "Image.Image":
    """
    SynthID 破除（硬路，默认开）：对齐大卫 xhs 实测版。
    顺序对齐大卫：先微调亮度/饱和(modulate) → 再缩2px裁1px边错位采样。
    只动边缘像素、按图片当前尺寸算，不依赖目标比例，不做全图重采样、不压全图像素。
    诚实声明：只能削弱 SynthID，不保证 100% 清除。
    """
    # 先 modulate（亮度*1.005/饱和*0.998），对齐大卫 modulate 在缩裁之前
    img = ImageEnhance.Brightness(img).enhance(1.005)
    img = ImageEnhance.Color(img).enhance(0.998)
    # 按当前尺寸缩2px再裁1px边，错位采样破边缘水印残留（w-4 × h-4）
    w, h = img.size
    img = ImageOps.fit(img, (w - 2, h - 2), method=Image.LANCZOS)
    img = img.crop((1, 1, w - 3, h - 3))
    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