mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-19 12:07:13 +02:00
244 lines
8.1 KiB
Python
244 lines
8.1 KiB
Python
"""Build non-generative spatial-fragmentation SynthID attack candidates.
|
|
|
|
The candidates combine deterministic smooth local warps with mild global
|
|
resampling, color changes, and codec round-trips. These are pixel transforms,
|
|
not semantic reconstruction or generative inpainting. Provider-oracle results
|
|
must be evaluated in a frozen batch with a source-positive control.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from dataclasses import asdict
|
|
from pathlib import Path
|
|
|
|
import click
|
|
import cv2
|
|
import numpy as np
|
|
from PIL import Image
|
|
from synthid_pixel_attack import (
|
|
jpeg_round_trip,
|
|
load_rgb,
|
|
measure,
|
|
norm_matched_noise,
|
|
resize_squeeze,
|
|
smooth_warp,
|
|
)
|
|
from synthid_v3_codebook_probe import load_v3_model, score_image
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
def bounded_smooth_warp(
|
|
pixels: np.ndarray,
|
|
*,
|
|
max_displacement: float,
|
|
sigma: float,
|
|
seed: int,
|
|
) -> np.ndarray:
|
|
"""Apply a smooth warp whose per-axis displacement is absolutely bounded."""
|
|
if max_displacement < 0.0 or sigma <= 0.0:
|
|
raise ValueError("max_displacement must be nonnegative and sigma positive")
|
|
height, width = pixels.shape[:2]
|
|
rng = np.random.default_rng(seed)
|
|
fields: list[np.ndarray] = []
|
|
for _ in range(2):
|
|
noise = rng.normal(size=(height, width)).astype(np.float32)
|
|
field = cv2.GaussianBlur(noise, (0, 0), sigmaX=sigma, sigmaY=sigma)
|
|
maximum = float(np.max(np.abs(field)))
|
|
fields.append(np.zeros_like(field) if maximum == 0.0 else field * (max_displacement / maximum))
|
|
yy, xx = np.mgrid[:height, :width].astype(np.float32)
|
|
return cv2.remap(
|
|
pixels,
|
|
xx + fields[0],
|
|
yy + fields[1],
|
|
interpolation=cv2.INTER_LANCZOS4,
|
|
borderMode=cv2.BORDER_REFLECT_101,
|
|
)
|
|
|
|
|
|
def affine_combo(pixels: np.ndarray, *, rotation_degrees: float, zoom: float) -> np.ndarray:
|
|
"""Apply one centered rotation-and-zoom resampling operation."""
|
|
if zoom < 0.0:
|
|
raise ValueError("zoom must be nonnegative")
|
|
height, width = pixels.shape[:2]
|
|
matrix = cv2.getRotationMatrix2D(
|
|
center=((width - 1) / 2.0, (height - 1) / 2.0),
|
|
angle=rotation_degrees,
|
|
scale=1.0 + zoom,
|
|
)
|
|
return cv2.warpAffine(
|
|
pixels,
|
|
matrix,
|
|
(width, height),
|
|
flags=cv2.INTER_LANCZOS4,
|
|
borderMode=cv2.BORDER_REFLECT_101,
|
|
)
|
|
|
|
|
|
def color_nudge(
|
|
pixels: np.ndarray,
|
|
*,
|
|
brightness: float,
|
|
contrast: float,
|
|
saturation: float,
|
|
hue_degrees: float,
|
|
) -> np.ndarray:
|
|
"""Apply bounded global RGB contrast and HSV saturation/hue changes."""
|
|
rgb = pixels.astype(np.float32) / 255.0
|
|
rgb = np.clip((rgb - 0.5) * (1.0 + contrast) + 0.5 + brightness, 0.0, 1.0)
|
|
hsv = cv2.cvtColor(rgb, cv2.COLOR_RGB2HSV)
|
|
hsv[:, :, 0] = np.mod(hsv[:, :, 0] + hue_degrees, 360.0)
|
|
hsv[:, :, 1] = np.clip(hsv[:, :, 1] * (1.0 + saturation), 0.0, 1.0)
|
|
result = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB)
|
|
return np.clip(np.rint(result * 255.0), 0, 255).astype(np.uint8)
|
|
|
|
|
|
def jpeg_chain(pixels: np.ndarray, qualities: tuple[int, ...]) -> np.ndarray:
|
|
"""Apply sequential JPEG round-trips at QUALITIES."""
|
|
result = pixels
|
|
for quality in qualities:
|
|
result = jpeg_round_trip(result, quality)
|
|
return result
|
|
|
|
|
|
def build_candidates(source: np.ndarray) -> dict[str, np.ndarray]:
|
|
"""Build the frozen spatial-fragmentation ladder for SOURCE."""
|
|
candidates: dict[str, np.ndarray] = {
|
|
"control": source.copy(),
|
|
"elastic-075": smooth_warp(source, amplitude=0.75, sigma=56.0, seed=20260812),
|
|
"elastic-125": smooth_warp(source, amplitude=1.25, sigma=52.0, seed=20260813),
|
|
"bounded-100": bounded_smooth_warp(
|
|
source,
|
|
max_displacement=1.0,
|
|
sigma=56.0,
|
|
seed=20260817,
|
|
),
|
|
"bounded-180": bounded_smooth_warp(
|
|
source,
|
|
max_displacement=1.8,
|
|
sigma=56.0,
|
|
seed=20260818,
|
|
),
|
|
"bounded-280": bounded_smooth_warp(
|
|
source,
|
|
max_displacement=2.8,
|
|
sigma=44.0,
|
|
seed=20260819,
|
|
),
|
|
}
|
|
|
|
balanced = smooth_warp(source, amplitude=0.75, sigma=56.0, seed=20260814)
|
|
balanced = affine_combo(balanced, rotation_degrees=0.2, zoom=0.004)
|
|
balanced = resize_squeeze(balanced, 0.94)
|
|
balanced = color_nudge(
|
|
balanced,
|
|
brightness=0.004,
|
|
contrast=0.006,
|
|
saturation=-0.005,
|
|
hue_degrees=0.15,
|
|
)
|
|
balanced = jpeg_chain(balanced, (94, 90))
|
|
candidates["fragment-balanced"] = balanced
|
|
|
|
strong = smooth_warp(source, amplitude=1.5, sigma=48.0, seed=20260815)
|
|
strong = affine_combo(strong, rotation_degrees=0.4, zoom=0.01)
|
|
strong = resize_squeeze(strong, 0.88)
|
|
strong = color_nudge(
|
|
strong,
|
|
brightness=0.008,
|
|
contrast=0.012,
|
|
saturation=-0.01,
|
|
hue_degrees=0.3,
|
|
)
|
|
strong = jpeg_chain(strong, (92, 88))
|
|
candidates["fragment-strong"] = strong
|
|
candidates["sham-strong-rms"] = norm_matched_noise(source, strong, seed=20260816)
|
|
|
|
bounded_balanced = bounded_smooth_warp(
|
|
source,
|
|
max_displacement=1.8,
|
|
sigma=56.0,
|
|
seed=20260820,
|
|
)
|
|
bounded_balanced = resize_squeeze(bounded_balanced, 0.98)
|
|
bounded_balanced = color_nudge(
|
|
bounded_balanced,
|
|
brightness=0.002,
|
|
contrast=0.003,
|
|
saturation=-0.003,
|
|
hue_degrees=0.1,
|
|
)
|
|
bounded_balanced = jpeg_chain(bounded_balanced, (96,))
|
|
candidates["bounded-fragment-balanced"] = bounded_balanced
|
|
|
|
bounded_strong = bounded_smooth_warp(
|
|
source,
|
|
max_displacement=2.8,
|
|
sigma=44.0,
|
|
seed=20260821,
|
|
)
|
|
bounded_strong = affine_combo(bounded_strong, rotation_degrees=0.2, zoom=0.004)
|
|
bounded_strong = resize_squeeze(bounded_strong, 0.94)
|
|
bounded_strong = color_nudge(
|
|
bounded_strong,
|
|
brightness=0.004,
|
|
contrast=0.006,
|
|
saturation=-0.005,
|
|
hue_degrees=0.15,
|
|
)
|
|
bounded_strong = jpeg_chain(bounded_strong, (94, 90))
|
|
candidates["bounded-fragment-strong"] = bounded_strong
|
|
candidates["sham-bounded-strong-rms"] = norm_matched_noise(source, bounded_strong, seed=20260822)
|
|
return candidates
|
|
|
|
|
|
@click.command()
|
|
@click.argument("codebook", type=click.Path(exists=True, dir_okay=False, path_type=Path))
|
|
@click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path))
|
|
@click.argument("output_dir", type=click.Path(file_okay=False, path_type=Path))
|
|
@click.option("--height", type=click.IntRange(min=64), required=True)
|
|
@click.option("--width", type=click.IntRange(min=64), required=True)
|
|
@click.option("--peak-count", type=click.IntRange(min=1), default=256, show_default=True)
|
|
def main(codebook: Path, source: Path, output_dir: Path, height: int, width: int, peak_count: int) -> None:
|
|
"""Write a frozen spatial-fragmentation batch for SOURCE."""
|
|
reference = load_rgb(source)
|
|
if reference.shape != (height, width, 3):
|
|
raise click.BadParameter("source geometry does not match --height and --width")
|
|
model = load_v3_model(codebook, height=height, width=width, peak_count=peak_count)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
rows: list[dict[str, object]] = []
|
|
for name, pixels in build_candidates(reference).items():
|
|
path = output_dir / f"{name}.png"
|
|
Image.fromarray(pixels, mode="RGB").save(path)
|
|
rows.append(
|
|
{
|
|
**asdict(measure(reference, pixels, name=name, path=path)),
|
|
**asdict(score_image(path, model)),
|
|
}
|
|
)
|
|
report_path = output_dir / "report.json"
|
|
report_path.write_text(
|
|
json.dumps(
|
|
{
|
|
"source": str(source),
|
|
"codebook": str(codebook),
|
|
"height": height,
|
|
"width": width,
|
|
"peak_count": peak_count,
|
|
"variants": rows,
|
|
},
|
|
indent=2,
|
|
)
|
|
+ "\n",
|
|
encoding="utf-8",
|
|
)
|
|
log.info("Wrote %d frozen fragmentation candidates: %s", len(rows), report_path)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
|
main()
|