mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-10-14 05:20:46 +02:00
107 lines
4.0 KiB
Python
107 lines
4.0 KiB
Python
"""Contracts for the portable OpenAI/Google source-pattern runtime."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from hashlib import sha256
|
|
from typing import TYPE_CHECKING
|
|
|
|
import numpy as np
|
|
import pytest
|
|
from PIL import Image, PngImagePlugin
|
|
|
|
from remove_ai_watermarks import source_classify_v2 as model
|
|
|
|
if TYPE_CHECKING:
|
|
from pathlib import Path
|
|
|
|
|
|
def _synthetic_weights(path: Path, *, gate: float) -> None:
|
|
zero_mean = np.zeros((256, 256, 3), dtype=np.complex64)
|
|
zero_precision = np.zeros((256, 256, 3), dtype=np.float32)
|
|
unit = np.zeros((992, 992), dtype=np.float32)
|
|
unit[0, 0] = 1
|
|
fields: dict[str, np.ndarray] = {"schema": np.asarray(model.MODEL_SCHEMA)}
|
|
for name in model._FILTER_NAMES:
|
|
fields[f"{name}_mean"] = zero_mean
|
|
fields[f"{name}_precision"] = zero_precision
|
|
for route in model.ROUTES:
|
|
fields[f"openai_{route}_unit"] = unit
|
|
limits = dict.fromkeys(model._THRESHOLD_NAMES, 1.0)
|
|
limits["r183_openai_gate"] = gate
|
|
limits.update({f"r163_{group}": -0.5 for group in model.GROUPS})
|
|
fields.update({f"threshold_{name}": np.asarray(value, dtype=np.float64) for name, value in limits.items()})
|
|
np.savez_compressed(path, **fields)
|
|
|
|
|
|
def test_r163_gate_changes_real_prediction(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
image = tmp_path / "photo.png"
|
|
Image.new("RGB", (512, 512), (80, 110, 145)).save(image)
|
|
blocked = tmp_path / "blocked.npz"
|
|
allowed = tmp_path / "allowed.npz"
|
|
_synthetic_weights(blocked, gate=0.5)
|
|
_synthetic_weights(allowed, gate=-0.5)
|
|
|
|
def predict(path: Path) -> model.ImageSourceClassification:
|
|
monkeypatch.setenv(model.WEIGHTS_ENV, str(path))
|
|
monkeypatch.setattr(model, "MODEL_SHA256", sha256(path.read_bytes()).hexdigest())
|
|
return model.classify_image_source(image)
|
|
|
|
no_gate = predict(blocked)
|
|
with_gate = predict(allowed)
|
|
|
|
assert no_gate.label == "unknown"
|
|
assert no_gate.reason == "abstained"
|
|
assert with_gate.label == "openai"
|
|
assert with_gate.reason == "classified"
|
|
assert with_gate.watermark_truth == "unknown"
|
|
|
|
|
|
def test_top_level_api_is_lazy_and_points_to_new_runtime() -> None:
|
|
import remove_ai_watermarks as raiw
|
|
|
|
assert raiw.classify_image_source is model.classify_image_source
|
|
assert raiw.ImageSourceClassification is model.ImageSourceClassification
|
|
assert model.PROVIDERS == ("openai", "google")
|
|
assert len(model.WEIGHTS_REVISION) == 40
|
|
assert all(character in "0123456789abcdef" for character in model.WEIGHTS_REVISION)
|
|
|
|
|
|
def test_metadata_does_not_change_pixel_decision(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
weights = tmp_path / model.MODEL_FILE
|
|
_synthetic_weights(weights, gate=-0.5)
|
|
monkeypatch.setenv(model.WEIGHTS_ENV, str(weights))
|
|
monkeypatch.setattr(model, "MODEL_SHA256", sha256(weights.read_bytes()).hexdigest())
|
|
pixels = Image.new("RGB", (512, 512), (80, 110, 145))
|
|
clean = tmp_path / "clean.png"
|
|
tagged = tmp_path / "tagged.png"
|
|
pixels.save(clean)
|
|
metadata = PngImagePlugin.PngInfo()
|
|
metadata.add_text("Comment", "arbitrary metadata")
|
|
pixels.save(tagged, pnginfo=metadata)
|
|
|
|
first = model.classify_image_source(clean)
|
|
second = model.classify_image_source(tagged)
|
|
|
|
assert first.label == second.label == "openai"
|
|
assert first.scores == second.scores
|
|
|
|
|
|
def test_rejects_wrong_artifact_hash(tmp_path: Path) -> None:
|
|
weights = tmp_path / model.MODEL_FILE
|
|
_synthetic_weights(weights, gate=0.5)
|
|
|
|
with pytest.raises(ValueError, match="SHA-256 mismatch"):
|
|
model._load_model(weights, "0" * 64)
|
|
|
|
|
|
def test_small_image_abstains_without_watermark_claim(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
|
|
monkeypatch.setattr(model, "_model_path", lambda: pytest.fail("small image loaded weights"))
|
|
image = tmp_path / "small.png"
|
|
Image.new("RGB", (128, 128), "white").save(image)
|
|
|
|
result = model.classify_image_source(image)
|
|
|
|
assert result.label == "unknown"
|
|
assert result.reason == "feature_unavailable"
|
|
assert result.watermark_truth == "unknown"
|