mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-19 20:17:12 +02:00
91 lines
3.2 KiB
Python
91 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from PIL import Image
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "scripts"))
|
|
|
|
import synthid_color_space_probe as probe
|
|
import synthid_ensemble_attack as attack
|
|
|
|
|
|
def _write_image(path: Path, *, phase: float, seed: int) -> None:
|
|
height = width = 64
|
|
rng = np.random.default_rng(seed)
|
|
pixels = 100.0 + rng.normal(0.0, 2.0, size=(height, width, 3))
|
|
yy, xx = np.mgrid[:height, :width]
|
|
wave = np.cos(2.0 * np.pi * (7.0 * yy / height + 5.0 * xx / width) + phase)
|
|
pixels[:, :, 0] += 14.0 * wave
|
|
pixels[:, :, 1] -= 10.0 * wave
|
|
Image.fromarray(np.clip(np.rint(pixels), 0, 255).astype(np.uint8), mode="RGB").save(path)
|
|
|
|
|
|
def test_projection_removes_only_positive_phase_component() -> None:
|
|
height = width = 64
|
|
yy, xx = np.mgrid[:height, :width]
|
|
channel = np.cos(2.0 * np.pi * (7.0 * yy / height + 5.0 * xx / width) + 0.4)
|
|
before = np.fft.fft2(channel)[7, 5]
|
|
phase = np.angle(before)
|
|
|
|
projected = attack._remove_positive_projection(
|
|
channel,
|
|
rows=np.asarray([7]),
|
|
columns=np.asarray([5]),
|
|
phases=np.asarray([phase]),
|
|
strength=1.0,
|
|
)
|
|
after = np.fft.fft2(projected)[7, 5]
|
|
|
|
assert np.real(after * np.exp(-1j * phase)) == pytest.approx(0.0, abs=1e-10)
|
|
assert np.max(np.abs(np.imag(np.fft.ifft2(np.fft.fft2(projected))))) < 1e-12
|
|
|
|
|
|
def test_alternating_projection_reduces_rgb_and_sv_evidence(tmp_path: Path) -> None:
|
|
positives: list[Path] = []
|
|
for index in range(3):
|
|
path = tmp_path / f"positive-{index}.png"
|
|
_write_image(path, phase=0.4, seed=index)
|
|
positives.append(path)
|
|
source = tmp_path / "source.png"
|
|
_write_image(source, phase=0.4, seed=10)
|
|
bins = np.asarray([(7, 5, channel) for channel in range(3)], dtype=np.int32)
|
|
rgb_model = probe.discover_model(positives, color_space="rgb", candidate_bins=bins, peak_count=3)
|
|
hsv_model = probe.discover_model(positives, color_space="hsv", candidate_bins=bins, peak_count=3)
|
|
before_rgb = probe.score_image(source, rgb_model)
|
|
before_hsv = probe.score_image(source, hsv_model)
|
|
|
|
with Image.open(source) as image:
|
|
source_pixels = np.asarray(image.convert("RGB"), dtype=np.uint8)
|
|
projected = attack.alternating_projection(
|
|
source_pixels,
|
|
rgb_model,
|
|
hsv_model,
|
|
strength=1.0,
|
|
iterations=2,
|
|
)
|
|
output = tmp_path / "projected.png"
|
|
Image.fromarray(projected, mode="RGB").save(output)
|
|
after_rgb = probe.score_image(output, rgb_model)
|
|
after_hsv = probe.score_image(output, hsv_model)
|
|
|
|
assert after_rgb.evidence_score < before_rgb.evidence_score
|
|
assert sum(after_hsv.channel_evidence[1:]) < sum(before_hsv.channel_evidence[1:])
|
|
|
|
|
|
@pytest.mark.parametrize("strength", [-0.1, 1.1])
|
|
def test_projection_rejects_out_of_range_strength(strength: float) -> None:
|
|
channel = np.zeros((64, 64))
|
|
|
|
with pytest.raises(ValueError, match="between zero and one"):
|
|
attack._remove_positive_projection(
|
|
channel,
|
|
rows=np.asarray([7]),
|
|
columns=np.asarray([5]),
|
|
phases=np.asarray([0.0]),
|
|
strength=strength,
|
|
)
|