src/veil_pgd/robustness/preprocess.py

"""Transforms that simulate what a scraper/training pipeline does to an image.

`scraper_sim` (JPEG Q85 + light blur) is baked into fitness so the overlay is
selected to survive realistic preprocessing. `hard_transforms` (heavy blur,
downscale, rotation) are the robustness *ceiling* used only in the eval harness.
"""

from __future__ import annotations

import io

from PIL import Image, ImageFilter

from veil_pgd.config import RobustnessSim, get_settings


def jpeg_recompress(image: Image.Image, quality: int) -> Image.Image:
    buf = io.BytesIO()
    image.convert("RGB").save(buf, format="JPEG", quality=quality)
    buf.seek(0)
    return Image.open(buf).convert("RGB")


def scraper_sim(image: Image.Image, cfg: RobustnessSim | None = None) -> Image.Image:
    cfg = cfg or get_settings().robustness
    out = image
    if cfg.gaussian_blur_radius > 0:
        out = out.filter(ImageFilter.GaussianBlur(cfg.gaussian_blur_radius))
    out = jpeg_recompress(out, cfg.jpeg_quality)
    return out


def hard_transforms(image: Image.Image) -> dict[str, Image.Image]:
    """The robustness ceiling: transforms known to weaken typographic overlays."""
    w, h = image.size
    return {
        "heavy_blur": image.filter(ImageFilter.GaussianBlur(2.5)),
        "downscale_half": image.resize((max(1, w // 2), max(1, h // 2))).resize((w, h)),
        "rotate_15": image.rotate(15, resample=Image.BICUBIC, expand=False),
        "jpeg_q50": jpeg_recompress(image, 50),
    }