mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-19 20:17:12 +02:00
260 lines
11 KiB
Python
260 lines
11 KiB
Python
"""Discover a polarity-invariant spectral carrier from low-texture groups.
|
|
|
|
This research harness is intentionally separate from the shipping detector. It
|
|
uses repeated, independently generated low-texture images to find Fourier bins
|
|
whose phase is stable within a content group and whose phase axis is stable
|
|
across groups. Treat its output as a carrier hypothesis until it passes the
|
|
provider-oracle and hard-negative gates in the SynthID research plan.
|
|
|
|
Usage:
|
|
uv run --extra pixels python scripts/synthid_consensus_probe.py discover \
|
|
refs/black refs/white refs/red --limit 5 \
|
|
--model-out .local-eval/synthid/consensus.npz
|
|
|
|
uv run --extra pixels python scripts/synthid_consensus_probe.py score \
|
|
.local-eval/synthid/consensus.npz images/*.png
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from dataclasses import asdict, dataclass
|
|
from pathlib import Path
|
|
|
|
import click
|
|
import numpy as np
|
|
from PIL import Image, ImageFilter
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
IMAGE_SUFFIXES = {".jpeg", ".jpg", ".png", ".webp"}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ConsensusScore:
|
|
"""One image's alignment with a frozen carrier hypothesis."""
|
|
|
|
path: str
|
|
score: float
|
|
active_weight_fraction: float
|
|
peak_count: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ConsensusModel:
|
|
"""A compact, pickle-free carrier hypothesis."""
|
|
|
|
size: int
|
|
peaks: np.ndarray
|
|
axial_phase: np.ndarray
|
|
weights: np.ndarray
|
|
expected_magnitude: np.ndarray
|
|
|
|
|
|
def _image_paths(directory: Path, limit: int | None) -> list[Path]:
|
|
"""Return a deterministic list of supported images in DIRECTORY."""
|
|
paths = sorted(path for path in directory.iterdir() if path.suffix.lower() in IMAGE_SUFFIXES)
|
|
if limit is not None:
|
|
paths = paths[:limit]
|
|
if not paths:
|
|
raise ValueError(f"no supported images in {directory}")
|
|
return paths
|
|
|
|
|
|
def _high_pass_rgb(path: Path, size: int, blur_radius: float) -> np.ndarray:
|
|
"""Decode PATH, canonicalize its geometry, and remove local image content."""
|
|
with Image.open(path) as source:
|
|
image = source.convert("RGB").resize((size, size), Image.Resampling.LANCZOS)
|
|
pixels = np.asarray(image, dtype=np.float64)
|
|
blurred = np.asarray(image.filter(ImageFilter.GaussianBlur(radius=blur_radius)), dtype=np.float64)
|
|
return pixels - blurred
|
|
|
|
|
|
def _spectrum(path: Path, size: int, blur_radius: float) -> np.ndarray:
|
|
"""Return a centered channel-wise spectrum for PATH."""
|
|
residual = _high_pass_rgb(path, size, blur_radius)
|
|
return np.fft.fftshift(np.fft.fft2(residual, axes=(0, 1)), axes=(0, 1))
|
|
|
|
|
|
def _group_statistics(paths: list[Path], size: int, blur_radius: float) -> tuple[np.ndarray, np.ndarray]:
|
|
"""Return phase coherence and mean magnitude for one reference group."""
|
|
spectra = [_spectrum(path, size, blur_radius) for path in paths]
|
|
units = [spectrum / (np.abs(spectrum) + 1e-12) for spectrum in spectra]
|
|
mean_unit = np.mean(units, axis=0)
|
|
coherence = np.abs(mean_unit)
|
|
mean_magnitude = np.mean([np.abs(spectrum) for spectrum in spectra], axis=0)
|
|
phase = np.angle(mean_unit)
|
|
return coherence * np.exp(1j * phase), mean_magnitude
|
|
|
|
|
|
def _valid_half_plane(size: int, min_radius: float, max_radius_fraction: float) -> np.ndarray:
|
|
"""Return the nonredundant Fourier region allowed for carrier selection."""
|
|
center = size // 2
|
|
yy, xx = np.ogrid[:size, :size]
|
|
dy = yy - center
|
|
dx = xx - center
|
|
radius = np.sqrt(np.square(dy) + np.square(dx))
|
|
half_plane = (dy > 0) | ((dy == 0) & (dx > 0))
|
|
off_axis = (dy != 0) & (dx != 0)
|
|
return half_plane & off_axis & (radius >= min_radius) & (radius <= size * max_radius_fraction)
|
|
|
|
|
|
def discover_model(
|
|
groups: list[list[Path]],
|
|
*,
|
|
size: int = 512,
|
|
blur_radius: float = 2.0,
|
|
peak_count: int = 256,
|
|
min_radius: float = 8.0,
|
|
max_radius_fraction: float = 0.4,
|
|
) -> ConsensusModel:
|
|
"""Discover a polarity-invariant carrier from independent image GROUPS."""
|
|
if len(groups) < 2:
|
|
raise ValueError("at least two reference groups are required")
|
|
if any(len(group) < 2 for group in groups):
|
|
raise ValueError("each reference group requires at least two images")
|
|
|
|
group_units: list[np.ndarray] = []
|
|
group_magnitudes: list[np.ndarray] = []
|
|
for paths in groups:
|
|
unit, magnitude = _group_statistics(paths, size, blur_radius)
|
|
group_units.append(unit)
|
|
group_magnitudes.append(magnitude)
|
|
|
|
stacked = np.stack(group_units, axis=0)
|
|
within_coherence = np.abs(stacked)
|
|
group_phase = np.angle(stacked)
|
|
axial_mean = np.mean(np.exp(2j * group_phase) * within_coherence, axis=0)
|
|
axial_coherence = np.abs(axial_mean) / (np.mean(within_coherence, axis=0) + 1e-12)
|
|
mean_within = np.mean(within_coherence, axis=0)
|
|
expected_magnitude = np.mean(group_magnitudes, axis=0)
|
|
|
|
magnitude_scale = np.median(expected_magnitude, axis=(0, 1), keepdims=True) + 1e-12
|
|
magnitude_score = np.log1p(expected_magnitude / magnitude_scale)
|
|
selection_score = np.square(mean_within) * np.square(axial_coherence) * magnitude_score
|
|
selection_score *= _valid_half_plane(size, min_radius, max_radius_fraction)[:, :, None]
|
|
|
|
candidate_count = int(np.count_nonzero(selection_score))
|
|
if candidate_count < peak_count:
|
|
raise ValueError(f"only {candidate_count} valid carrier candidates for {peak_count} peaks")
|
|
flat = selection_score.ravel()
|
|
indices = np.argpartition(flat, -peak_count)[-peak_count:]
|
|
indices = indices[np.argsort(flat[indices])[::-1]]
|
|
rows, columns, channels = np.unravel_index(indices, selection_score.shape)
|
|
peaks = np.column_stack((rows - size // 2, columns - size // 2, channels)).astype(np.int32)
|
|
|
|
axial_phase = 0.5 * np.angle(axial_mean[rows, columns, channels])
|
|
weights = selection_score[rows, columns, channels]
|
|
weights /= np.sum(weights)
|
|
magnitudes = expected_magnitude[rows, columns, channels]
|
|
return ConsensusModel(
|
|
size=size,
|
|
peaks=peaks,
|
|
axial_phase=axial_phase.astype(np.float64),
|
|
weights=weights.astype(np.float64),
|
|
expected_magnitude=magnitudes.astype(np.float64),
|
|
)
|
|
|
|
|
|
def score_image(path: Path, model: ConsensusModel, *, blur_radius: float = 2.0) -> ConsensusScore:
|
|
"""Score PATH against a frozen polarity-invariant carrier model."""
|
|
spectrum = _spectrum(path, model.size, blur_radius)
|
|
center = model.size // 2
|
|
rows = center + model.peaks[:, 0]
|
|
columns = center + model.peaks[:, 1]
|
|
channels = model.peaks[:, 2]
|
|
values = spectrum[rows, columns, channels]
|
|
phase_alignment = np.cos(2.0 * (np.angle(values) - model.axial_phase))
|
|
magnitude_gate = np.minimum(np.abs(values) / (model.expected_magnitude + 1e-12), 1.0)
|
|
active_weights = model.weights * magnitude_gate
|
|
active_weight = float(np.sum(active_weights))
|
|
score = 0.0 if active_weight == 0.0 else float(np.sum(active_weights * phase_alignment) / active_weight)
|
|
return ConsensusScore(
|
|
path=str(path),
|
|
score=score,
|
|
active_weight_fraction=active_weight,
|
|
peak_count=len(model.peaks),
|
|
)
|
|
|
|
|
|
def save_model(path: Path, model: ConsensusModel) -> None:
|
|
"""Save MODEL as a pickle-free NPZ artifact."""
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
np.savez_compressed(
|
|
path,
|
|
size=np.asarray(model.size, dtype=np.int32),
|
|
peaks=model.peaks.astype(np.int32),
|
|
axial_phase=model.axial_phase.astype(np.float32),
|
|
weights=model.weights.astype(np.float32),
|
|
expected_magnitude=model.expected_magnitude.astype(np.float32),
|
|
)
|
|
|
|
|
|
def load_model(path: Path) -> ConsensusModel:
|
|
"""Load and validate a pickle-free carrier model."""
|
|
with np.load(path, allow_pickle=False) as artifact:
|
|
model = ConsensusModel(
|
|
size=int(artifact["size"]),
|
|
peaks=np.asarray(artifact["peaks"], dtype=np.int32),
|
|
axial_phase=np.asarray(artifact["axial_phase"], dtype=np.float64),
|
|
weights=np.asarray(artifact["weights"], dtype=np.float64),
|
|
expected_magnitude=np.asarray(artifact["expected_magnitude"], dtype=np.float64),
|
|
)
|
|
count = len(model.peaks)
|
|
if model.peaks.ndim != 2 or model.peaks.shape[1] != 3:
|
|
raise ValueError("invalid peak shape")
|
|
if any(array.shape != (count,) for array in (model.axial_phase, model.weights, model.expected_magnitude)):
|
|
raise ValueError("model arrays do not match peak count")
|
|
if model.size < 64 or np.any(np.abs(model.peaks[:, :2]) >= model.size // 2):
|
|
raise ValueError("invalid canonical size or peak coordinates")
|
|
if np.any(model.peaks[:, 2] < 0) or np.any(model.peaks[:, 2] > 2):
|
|
raise ValueError("invalid channel index")
|
|
if not np.isclose(np.sum(model.weights), 1.0, atol=1e-5):
|
|
raise ValueError("model weights must sum to one")
|
|
return model
|
|
|
|
|
|
@click.group()
|
|
def main() -> None:
|
|
"""Run low-texture carrier discovery and scoring experiments."""
|
|
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
|
|
|
|
|
@main.command()
|
|
@click.argument("group_dirs", nargs=-1, required=True, type=click.Path(exists=True, file_okay=False, path_type=Path))
|
|
@click.option("--limit", type=click.IntRange(min=2))
|
|
@click.option("--size", type=click.IntRange(min=64), default=512, show_default=True)
|
|
@click.option("--peak-count", type=click.IntRange(min=1), default=256, show_default=True)
|
|
@click.option("--model-out", type=click.Path(dir_okay=False, path_type=Path), required=True)
|
|
def discover(group_dirs: tuple[Path, ...], limit: int | None, size: int, peak_count: int, model_out: Path) -> None:
|
|
"""Discover a carrier from the images in each GROUP_DIRS directory."""
|
|
groups = [_image_paths(directory, limit) for directory in group_dirs]
|
|
model = discover_model(groups, size=size, peak_count=peak_count)
|
|
save_model(model_out, model)
|
|
log.info("Wrote consensus model: %s", model_out)
|
|
|
|
|
|
@main.command()
|
|
@click.argument("model_path", type=click.Path(exists=True, dir_okay=False, path_type=Path))
|
|
@click.argument("images", nargs=-1, required=True, type=click.Path(exists=True, dir_okay=False, path_type=Path))
|
|
@click.option("--report-out", type=click.Path(dir_okay=False, path_type=Path))
|
|
def score(model_path: Path, images: tuple[Path, ...], report_out: Path | None) -> None:
|
|
"""Score IMAGES against MODEL_PATH."""
|
|
model = load_model(model_path)
|
|
payload = {
|
|
"model": str(model_path),
|
|
"scores": [asdict(score_image(image, model)) for image in images],
|
|
}
|
|
rendered = json.dumps(payload, indent=2) + "\n"
|
|
if report_out is None:
|
|
log.info("%s", rendered.rstrip())
|
|
return
|
|
report_out.parent.mkdir(parents=True, exist_ok=True)
|
|
report_out.write_text(rendered, encoding="utf-8")
|
|
log.info("Wrote score report: %s", report_out)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|