mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-19 20:17:12 +02:00
313 lines
12 KiB
Python
313 lines
12 KiB
Python
"""Discover and score a shared spectral carrier from exact image pairs.
|
|
|
|
This is a research harness, not a production SynthID detector. Pair provenance
|
|
and oracle labels remain external evidence. The harness deliberately separates
|
|
template discovery from single-image scoring and stores arrays in NPZ without
|
|
pickle.
|
|
|
|
Usage:
|
|
uv run python scripts/synthid_spectral_probe.py discover \
|
|
--pair clean.png marked.png --pair clean2.png marked2.png \
|
|
--template-out .local-eval/synthid/template.npz \
|
|
--report-out .local-eval/synthid/pair-report.json
|
|
|
|
uv run python scripts/synthid_spectral_probe.py score \
|
|
.local-eval/synthid/template.npz image.png other.png
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
from dataclasses import asdict, dataclass
|
|
from itertools import combinations
|
|
from pathlib import Path
|
|
|
|
import click
|
|
import numpy as np
|
|
from PIL import Image, ImageFilter
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PairMeasurement:
|
|
"""Pixel-domain measurements for one exact clean/marked pair."""
|
|
|
|
clean: str
|
|
marked: str
|
|
width: int
|
|
height: int
|
|
psnr_db: float
|
|
changed_pixel_fraction: float
|
|
difference_min: float
|
|
difference_max: float
|
|
channel_mean: tuple[float, float, float]
|
|
channel_std: tuple[float, float, float]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ImageScore:
|
|
"""Single-image phase alignment against a discovered template."""
|
|
|
|
path: str
|
|
phase_mean: float
|
|
phase_weighted: float
|
|
channel_phase_weighted: tuple[float, float, float]
|
|
top_two_channel_phase_12: float
|
|
peak_count: int
|
|
|
|
|
|
def load_rgb(path: Path) -> np.ndarray:
|
|
"""Load PATH as float64 RGB pixels."""
|
|
with Image.open(path) as image:
|
|
return np.asarray(image.convert("RGB"), dtype=np.float64)
|
|
|
|
|
|
def _resize_float(channel: np.ndarray, size: int) -> np.ndarray:
|
|
"""Resize one floating-point channel without quantizing the residual."""
|
|
image = Image.fromarray(channel.astype(np.float32), mode="F")
|
|
return np.asarray(image.resize((size, size), Image.Resampling.BILINEAR), dtype=np.float64)
|
|
|
|
|
|
def pair_residual(clean: Path, marked: Path, *, size: int = 512) -> tuple[np.ndarray, PairMeasurement]:
|
|
"""Return a canonical RGB residual and measurements for an exact pair."""
|
|
clean_rgb = load_rgb(clean)
|
|
marked_rgb = load_rgb(marked)
|
|
if clean_rgb.shape != marked_rgb.shape:
|
|
raise ValueError(f"pair shapes differ: {clean_rgb.shape} != {marked_rgb.shape}")
|
|
|
|
difference = marked_rgb - clean_rgb
|
|
mse = float(np.mean(np.square(difference)))
|
|
psnr = float("inf") if mse == 0.0 else float(20.0 * np.log10(255.0 / np.sqrt(mse)))
|
|
residual = np.stack([_resize_float(difference[:, :, channel], size) for channel in range(3)], axis=2)
|
|
measurement = PairMeasurement(
|
|
clean=str(clean),
|
|
marked=str(marked),
|
|
width=int(clean_rgb.shape[1]),
|
|
height=int(clean_rgb.shape[0]),
|
|
psnr_db=psnr,
|
|
changed_pixel_fraction=float(np.mean(np.any(difference != 0.0, axis=2))),
|
|
difference_min=float(np.min(difference)),
|
|
difference_max=float(np.max(difference)),
|
|
channel_mean=tuple(float(value) for value in np.mean(difference, axis=(0, 1))),
|
|
channel_std=tuple(float(value) for value in np.std(difference, axis=(0, 1))),
|
|
)
|
|
return residual, measurement
|
|
|
|
|
|
def normalized_channels(residual: np.ndarray) -> np.ndarray:
|
|
"""Zero-center and unit-normalize each residual channel."""
|
|
centered = residual - np.mean(residual, axis=(0, 1), keepdims=True)
|
|
norms = np.linalg.norm(centered, axis=(0, 1), keepdims=True)
|
|
return np.divide(centered, norms, out=np.zeros_like(centered), where=norms != 0.0)
|
|
|
|
|
|
def channel_ncc(first: np.ndarray, second: np.ndarray) -> tuple[float, float, float]:
|
|
"""Return per-channel normalized cross-correlation."""
|
|
first_norm = normalized_channels(first)
|
|
second_norm = normalized_channels(second)
|
|
values = np.sum(first_norm * second_norm, axis=(0, 1))
|
|
return tuple(float(value) for value in values)
|
|
|
|
|
|
def build_template(residuals: list[np.ndarray]) -> np.ndarray:
|
|
"""Average canonical residuals after per-channel normalization."""
|
|
if not residuals:
|
|
raise ValueError("at least one residual is required")
|
|
shape = residuals[0].shape
|
|
if any(residual.shape != shape for residual in residuals):
|
|
raise ValueError("all canonical residuals must have the same shape")
|
|
return np.mean([normalized_channels(residual) for residual in residuals], axis=0)
|
|
|
|
|
|
def _template_fft(template: np.ndarray) -> np.ndarray:
|
|
"""Return a centered two-dimensional FFT for each RGB channel."""
|
|
return np.fft.fftshift(np.fft.fft2(template, axes=(0, 1)), axes=(0, 1))
|
|
|
|
|
|
def select_peaks(
|
|
template: np.ndarray,
|
|
*,
|
|
count: int = 64,
|
|
min_radius: float = 8.0,
|
|
max_radius_fraction: float = 0.35,
|
|
min_distance: float = 3.0,
|
|
) -> np.ndarray:
|
|
"""Select separated high-energy carrier bins from one Fourier half-plane."""
|
|
if count <= 0:
|
|
raise ValueError("count must be positive")
|
|
height, width, channels = template.shape
|
|
if height != width or channels != 3:
|
|
raise ValueError("template must be a square RGB array")
|
|
center = height // 2
|
|
spectrum = _template_fft(template)
|
|
magnitude = np.linalg.norm(spectrum, axis=2)
|
|
yy, xx = np.ogrid[:height, :width]
|
|
radius = np.sqrt(np.square(yy - center) + np.square(xx - center))
|
|
valid = (radius >= min_radius) & (radius <= height * max_radius_fraction)
|
|
candidates = np.flatnonzero(valid)
|
|
order = candidates[np.argsort(magnitude.ravel()[candidates])[::-1]]
|
|
|
|
selected: list[tuple[int, int]] = []
|
|
for flat_index in order:
|
|
row, column = np.unravel_index(flat_index, magnitude.shape)
|
|
dy, dx = int(row - center), int(column - center)
|
|
if dy < 0 or (dy == 0 and dx < 0):
|
|
continue
|
|
if any((dy - old_dy) ** 2 + (dx - old_dx) ** 2 < min_distance**2 for old_dy, old_dx in selected):
|
|
continue
|
|
selected.append((dy, dx))
|
|
if len(selected) == count:
|
|
break
|
|
if len(selected) != count:
|
|
raise ValueError(f"could select only {len(selected)} of {count} peaks")
|
|
return np.asarray(selected, dtype=np.int32)
|
|
|
|
|
|
def _high_pass_rgb(path: Path, size: int, blur_radius: float) -> np.ndarray:
|
|
"""Decode, resize, and subtract a small Gaussian blur from RGB pixels."""
|
|
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 score_image(path: Path, template: np.ndarray, peaks: np.ndarray, *, blur_radius: float = 2.0) -> ImageScore:
|
|
"""Score one image by phase alignment at discovered carrier bins."""
|
|
size = int(template.shape[0])
|
|
image_fft = np.fft.fftshift(np.fft.fft2(_high_pass_rgb(path, size, blur_radius), axes=(0, 1)), axes=(0, 1))
|
|
template_fft = _template_fft(template)
|
|
center = size // 2
|
|
phase_values: list[np.ndarray] = []
|
|
weights: list[np.ndarray] = []
|
|
for dy, dx in peaks:
|
|
image_value = image_fft[center + int(dy), center + int(dx)]
|
|
template_value = template_fft[center + int(dy), center + int(dx)]
|
|
phase = np.real(image_value * np.conj(template_value)) / (np.abs(image_value) * np.abs(template_value) + 1e-12)
|
|
phase_values.append(phase)
|
|
weights.append(np.abs(template_value))
|
|
phases = np.asarray(phase_values)
|
|
carrier_weights = np.asarray(weights)
|
|
channel_weighted = np.sum(phases * carrier_weights, axis=0) / np.sum(carrier_weights, axis=0)
|
|
consensus_count = min(12, len(peaks))
|
|
consensus_channels = np.sum(phases[:consensus_count] * carrier_weights[:consensus_count], axis=0) / np.sum(
|
|
carrier_weights[:consensus_count], axis=0
|
|
)
|
|
top_two_channel_phase_12 = float(np.mean(np.sort(consensus_channels)[-2:]))
|
|
return ImageScore(
|
|
path=str(path),
|
|
phase_mean=float(np.mean(phases)),
|
|
phase_weighted=float(np.sum(phases * carrier_weights) / np.sum(carrier_weights)),
|
|
channel_phase_weighted=tuple(float(value) for value in channel_weighted),
|
|
top_two_channel_phase_12=top_two_channel_phase_12,
|
|
peak_count=len(peaks),
|
|
)
|
|
|
|
|
|
def save_template(path: Path, template: np.ndarray, peaks: np.ndarray) -> None:
|
|
"""Store a template in a pickle-free compressed NPZ artifact."""
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
np.savez_compressed(path, template=template.astype(np.float32), peaks=peaks.astype(np.int32))
|
|
|
|
|
|
def load_template(path: Path) -> tuple[np.ndarray, np.ndarray]:
|
|
"""Load a template artifact without enabling pickle."""
|
|
with np.load(path, allow_pickle=False) as artifact:
|
|
template = np.asarray(artifact["template"], dtype=np.float64)
|
|
peaks = np.asarray(artifact["peaks"], dtype=np.int32)
|
|
if template.ndim != 3 or template.shape[2] != 3 or template.shape[0] != template.shape[1]:
|
|
raise ValueError("invalid template shape")
|
|
if peaks.ndim != 2 or peaks.shape[1] != 2:
|
|
raise ValueError("invalid peak shape")
|
|
return template, peaks
|
|
|
|
|
|
def discovery_report(
|
|
residuals: list[np.ndarray], measurements: list[PairMeasurement], peaks: np.ndarray
|
|
) -> dict[str, object]:
|
|
"""Build a JSON-safe report with pair statistics and cross-pair NCC."""
|
|
pairwise = [
|
|
{
|
|
"first": measurements[first].marked,
|
|
"second": measurements[second].marked,
|
|
"channel_ncc": channel_ncc(residuals[first], residuals[second]),
|
|
}
|
|
for first, second in combinations(range(len(residuals)), 2)
|
|
]
|
|
return {
|
|
"pair_count": len(measurements),
|
|
"pairs": [asdict(measurement) for measurement in measurements],
|
|
"pairwise": pairwise,
|
|
"peaks": peaks.tolist(),
|
|
}
|
|
|
|
|
|
@click.group()
|
|
def main() -> None:
|
|
"""Discover and score an experimental shared spectral carrier."""
|
|
logging.basicConfig(level=logging.INFO, format="%(message)s")
|
|
|
|
|
|
@main.command()
|
|
@click.option(
|
|
"--pair",
|
|
"pairs",
|
|
type=(
|
|
click.Path(exists=True, dir_okay=False, path_type=Path),
|
|
click.Path(exists=True, dir_okay=False, path_type=Path),
|
|
),
|
|
multiple=True,
|
|
required=True,
|
|
help="Exact CLEAN MARKED pair; repeat for multiple pairs.",
|
|
)
|
|
@click.option("--size", type=click.IntRange(min=64), default=512, show_default=True)
|
|
@click.option("--peak-count", type=click.IntRange(min=1), default=64, show_default=True)
|
|
@click.option("--template-out", type=click.Path(dir_okay=False, path_type=Path), required=True)
|
|
@click.option("--report-out", type=click.Path(dir_okay=False, path_type=Path), required=True)
|
|
def discover(
|
|
pairs: tuple[tuple[Path, Path], ...],
|
|
size: int,
|
|
peak_count: int,
|
|
template_out: Path,
|
|
report_out: Path,
|
|
) -> None:
|
|
"""Build a template and report from exact CLEAN MARKED pairs."""
|
|
residuals: list[np.ndarray] = []
|
|
measurements: list[PairMeasurement] = []
|
|
for clean, marked in pairs:
|
|
residual, measurement = pair_residual(clean, marked, size=size)
|
|
residuals.append(residual)
|
|
measurements.append(measurement)
|
|
template = build_template(residuals)
|
|
peaks = select_peaks(template, count=peak_count)
|
|
save_template(template_out, template, peaks)
|
|
report = discovery_report(residuals, measurements, peaks)
|
|
report_out.parent.mkdir(parents=True, exist_ok=True)
|
|
report_out.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
|
|
log.info("Wrote template: %s", template_out)
|
|
log.info("Wrote report: %s", report_out)
|
|
|
|
|
|
@main.command()
|
|
@click.argument("template_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(template_path: Path, images: tuple[Path, ...], report_out: Path | None) -> None:
|
|
"""Score IMAGES against TEMPLATE_PATH."""
|
|
template, peaks = load_template(template_path)
|
|
scores = [asdict(score_image(image, template, peaks)) for image in images]
|
|
payload = {"template": str(template_path), "scores": scores}
|
|
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()
|