mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-31 09:40:38 +02:00
289 lines
11 KiB
Python
289 lines
11 KiB
Python
"""Independent split-patch confirmation for the registered SynthID carrier."""
|
|
|
|
# The optional numeric libraries do not provide complete types for this path.
|
|
# pyright: reportMissingTypeStubs=false, reportUnknownMemberType=false, reportUnknownVariableType=false, reportUnknownArgumentType=false
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
from dataclasses import dataclass
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
from synthid_runtime.synthid_detector import fold_residual_template, unit_tile
|
|
|
|
if TYPE_CHECKING:
|
|
from numpy.typing import NDArray
|
|
|
|
MIN_PERIOD = 10.0
|
|
MIN_COHERENCE = 0.30
|
|
MIN_AMPLITUDE = 0.0
|
|
H5_PERIOD = (18.0, 18.6)
|
|
H5_MIN = 0.13
|
|
STRONG_COHERENCE_PERIOD = (18.6, 20.0)
|
|
STRONG_COHERENCE_MIN = 0.40
|
|
WEAK_H5_PERIOD = (20.0, 22.0)
|
|
WEAK_H5_MIN = 0.02
|
|
PATCH_SIZE = 256
|
|
GRID_SIZE = 4
|
|
HARMONIC_COUNT = 16
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RegisteredConfirmationComponents:
|
|
"""Auditable split-patch confirmation components for one fixed period."""
|
|
|
|
period: float
|
|
joint_coherence: float
|
|
joint_amplitude: float
|
|
unknown_codeword_fixed_confirmation: float
|
|
selection_patches: int
|
|
confirmation_patches: int
|
|
|
|
@property
|
|
def passes(self) -> bool:
|
|
"""Whether every frozen period-aware confirmation gate passes."""
|
|
return registered_confirmation_passes(
|
|
self.period,
|
|
self.joint_coherence,
|
|
self.joint_amplitude,
|
|
self.unknown_codeword_fixed_confirmation,
|
|
)
|
|
|
|
|
|
def registered_confirmation_passes(
|
|
period: float,
|
|
joint_coherence: float,
|
|
joint_amplitude: float,
|
|
unknown_codeword_fixed_confirmation: float,
|
|
) -> bool:
|
|
"""Apply the single frozen registered-carrier confirmation rule."""
|
|
if period < MIN_PERIOD:
|
|
return False
|
|
if joint_coherence < MIN_COHERENCE or joint_amplitude < MIN_AMPLITUDE:
|
|
return False
|
|
if H5_PERIOD[0] <= period < H5_PERIOD[1]:
|
|
return unknown_codeword_fixed_confirmation >= H5_MIN
|
|
if STRONG_COHERENCE_PERIOD[0] <= period < STRONG_COHERENCE_PERIOD[1]:
|
|
return joint_coherence >= STRONG_COHERENCE_MIN
|
|
if WEAK_H5_PERIOD[0] <= period < WEAK_H5_PERIOD[1]:
|
|
return unknown_codeword_fixed_confirmation >= WEAK_H5_MIN
|
|
return True
|
|
|
|
|
|
def _opponent_channels(values: NDArray[Any]) -> NDArray[Any]:
|
|
red = values[:, :, 0]
|
|
green = values[:, :, 1]
|
|
blue = values[:, :, 2]
|
|
return np.stack((green, red - green, blue - 0.5 * (red + green)), axis=2)
|
|
|
|
|
|
def _template_harmonics(template: NDArray[Any]) -> tuple[NDArray[Any], NDArray[Any]]:
|
|
opponent = _opponent_channels(np.asarray(template, dtype=np.float64))
|
|
spectrum = np.fft.fft2(opponent, axes=(0, 1))
|
|
height, width = template.shape[:2]
|
|
candidates: list[tuple[float, int, int]] = []
|
|
for row in range(height):
|
|
signed_row = row if row <= height // 2 else row - height
|
|
for column in range(width):
|
|
signed_column = column if column <= width // 2 else column - width
|
|
if signed_row < 0 or (signed_row == 0 and signed_column <= 0):
|
|
continue
|
|
power = float(np.sum(np.abs(spectrum[row, column]) ** 2))
|
|
candidates.append((power, signed_row, signed_column))
|
|
candidates.sort(reverse=True)
|
|
selected = candidates[:HARMONIC_COUNT]
|
|
harmonics = np.asarray([(row, column) for _power, row, column in selected], dtype=np.float64)
|
|
coefficients = np.asarray([spectrum[int(row) % height, int(column) % width] for row, column in harmonics])
|
|
weights = np.abs(coefficients)
|
|
weight_sum = float(np.sum(weights))
|
|
if weight_sum <= 0.0:
|
|
raise ValueError("template has no nonzero periodic harmonics")
|
|
return harmonics, weights / weight_sum
|
|
|
|
|
|
def _patch_origins(height: int, width: int) -> list[tuple[int, int, int]]:
|
|
if height < PATCH_SIZE or width < PATCH_SIZE:
|
|
raise ValueError("registered confirmation needs both image sides to be at least 256 pixels")
|
|
y_values = np.linspace(0, height - PATCH_SIZE, min(GRID_SIZE, height // PATCH_SIZE), dtype=np.int64)
|
|
x_values = np.linspace(0, width - PATCH_SIZE, min(GRID_SIZE, width // PATCH_SIZE), dtype=np.int64)
|
|
origins = [
|
|
(int(y), int(x), (y_index + x_index) % 2)
|
|
for y_index, y in enumerate(np.unique(y_values))
|
|
for x_index, x in enumerate(np.unique(x_values))
|
|
]
|
|
if {group for _y, _x, group in origins} != {0, 1}:
|
|
raise ValueError("registered confirmation needs two independent patch groups")
|
|
return origins
|
|
|
|
|
|
def _bilinear_sample(spectrum: NDArray[Any], y: NDArray[Any], x: NDArray[Any]) -> NDArray[Any]:
|
|
height, width = spectrum.shape
|
|
y_floor = np.floor(y)
|
|
x_floor = np.floor(x)
|
|
y0 = y_floor.astype(np.int64) % height
|
|
x0 = x_floor.astype(np.int64) % width
|
|
y1 = (y0 + 1) % height
|
|
x1 = (x0 + 1) % width
|
|
dy = y - y_floor
|
|
dx = x - x_floor
|
|
return (
|
|
spectrum[y0, x0] * (1.0 - dy) * (1.0 - dx)
|
|
+ spectrum[y1, x0] * dy * (1.0 - dx)
|
|
+ spectrum[y0, x1] * (1.0 - dy) * dx
|
|
+ spectrum[y1, x1] * dy * dx
|
|
)
|
|
|
|
|
|
def _patch_unit_values(
|
|
pixels: NDArray[Any],
|
|
origin_y: int,
|
|
origin_x: int,
|
|
period: float,
|
|
harmonics: NDArray[Any],
|
|
denoise_sigma: float,
|
|
) -> NDArray[Any]:
|
|
patch = np.asarray(
|
|
pixels[origin_y : origin_y + PATCH_SIZE, origin_x : origin_x + PATCH_SIZE],
|
|
dtype=np.float32,
|
|
)
|
|
channels = _opponent_channels(patch)
|
|
window_1d = np.hanning(PATCH_SIZE).astype(np.float32)
|
|
window = window_1d[:, None] * window_1d[None, :]
|
|
frequencies_y = harmonics[:, 0] / period
|
|
frequencies_x = harmonics[:, 1] / period
|
|
sample_y = frequencies_y * PATCH_SIZE
|
|
sample_x = frequencies_x * PATCH_SIZE
|
|
sampled = np.empty((len(harmonics), 3), dtype=np.complex128)
|
|
for channel in range(3):
|
|
residual = channels[:, :, channel]
|
|
residual -= cv2.GaussianBlur(
|
|
residual,
|
|
(0, 0),
|
|
sigmaX=denoise_sigma,
|
|
sigmaY=denoise_sigma,
|
|
borderType=cv2.BORDER_REFLECT_101,
|
|
)
|
|
sampled[:, channel] = _bilinear_sample(np.fft.fft2(residual * window), sample_y, sample_x)
|
|
sampled *= np.exp(-2j * math.pi * (frequencies_y * origin_y + frequencies_x * origin_x))[:, None]
|
|
magnitudes = np.abs(sampled)
|
|
return np.divide(sampled, magnitudes, out=np.zeros_like(sampled), where=magnitudes > 1e-12)
|
|
|
|
|
|
def _coherence(values: list[NDArray[Any]], weights: NDArray[Any]) -> float:
|
|
coherence = np.abs(np.mean(np.stack(values), axis=0))
|
|
return float(np.sum(coherence * weights))
|
|
|
|
|
|
def _unknown_codeword_fixed_confirmation(
|
|
selection_values: list[NDArray[Any]],
|
|
confirmation_values: list[NDArray[Any]],
|
|
weights: NDArray[Any],
|
|
) -> float:
|
|
cross_codeword = np.mean(np.stack(confirmation_values), axis=0) * np.conj(
|
|
np.mean(np.stack(selection_values), axis=0)
|
|
)
|
|
confirmation_mask = np.arange(len(weights)) % 2 == 1
|
|
masked_weights = weights[confirmation_mask]
|
|
return float(np.abs(np.sum(cross_codeword[confirmation_mask] * masked_weights)) / np.sum(masked_weights))
|
|
|
|
|
|
def _canonical_pixels(pixels: NDArray[Any], template: NDArray[Any], period: float) -> NDArray[Any]:
|
|
width = max(template.shape[1], round(pixels.shape[1] * template.shape[1] / period))
|
|
height = max(template.shape[0], round(pixels.shape[0] * template.shape[0] / period))
|
|
if (height, width) == pixels.shape[:2]:
|
|
return pixels
|
|
interpolation = cv2.INTER_AREA if width < pixels.shape[1] else cv2.INTER_CUBIC
|
|
return np.asarray(cv2.resize(pixels, (width, height), interpolation=interpolation))
|
|
|
|
|
|
def _cyclic_correlations(template: NDArray[Any], tile: NDArray[Any]) -> NDArray[Any]:
|
|
template_spectrum = np.fft.fft2(template, axes=(0, 1))
|
|
tile_spectrum = np.fft.fft2(tile, axes=(0, 1))
|
|
return np.fft.ifft2(np.sum(template_spectrum * np.conj(tile_spectrum), axis=2)).real
|
|
|
|
|
|
def _joint_amplitude(
|
|
pixels: NDArray[Any],
|
|
template: NDArray[Any],
|
|
period: float,
|
|
denoise_sigma: float,
|
|
) -> tuple[float, int, int]:
|
|
canonical = _canonical_pixels(pixels, template, period)
|
|
tile_height, tile_width = template.shape[:2]
|
|
grouped_units: dict[int, list[NDArray[Any]]] = {0: [], 1: []}
|
|
origins = _patch_origins(*canonical.shape[:2])
|
|
for origin_y, origin_x, group in origins:
|
|
aligned_y = (origin_y // tile_height) * tile_height
|
|
aligned_x = (origin_x // tile_width) * tile_width
|
|
folded = fold_residual_template(
|
|
canonical[aligned_y : aligned_y + PATCH_SIZE, aligned_x : aligned_x + PATCH_SIZE],
|
|
tile_height=tile_height,
|
|
tile_width=tile_width,
|
|
denoise_sigma=denoise_sigma,
|
|
)
|
|
unit, _norm = unit_tile(folded)
|
|
grouped_units[group].append(unit)
|
|
selection_tile, _selection_norm = unit_tile(np.mean(grouped_units[0], axis=0))
|
|
confirmation_tile, _confirmation_norm = unit_tile(np.mean(grouped_units[1], axis=0))
|
|
selection_correlations = _cyclic_correlations(template, selection_tile)
|
|
confirmation_correlations = _cyclic_correlations(template, confirmation_tile)
|
|
shift_y, shift_x = np.unravel_index(int(np.argmax(selection_correlations)), selection_correlations.shape)
|
|
return (
|
|
min(
|
|
float(selection_correlations[shift_y, shift_x]),
|
|
float(confirmation_correlations[shift_y, shift_x]),
|
|
),
|
|
len(grouped_units[0]),
|
|
len(grouped_units[1]),
|
|
)
|
|
|
|
|
|
def registered_confirmation_components(
|
|
pixels: NDArray[Any],
|
|
template: NDArray[Any],
|
|
period: float,
|
|
denoise_sigma: float,
|
|
) -> RegisteredConfirmationComponents:
|
|
"""Measure the frozen split-patch gates at one registered carrier period."""
|
|
if pixels.ndim != 3 or pixels.shape[2] != 3:
|
|
raise ValueError("pixels must have shape (height, width, 3)")
|
|
if not math.isfinite(period) or period <= 0.0:
|
|
raise ValueError("registered period must be finite and positive")
|
|
harmonics, weights = _template_harmonics(template)
|
|
grouped_values: dict[int, list[NDArray[Any]]] = {0: [], 1: []}
|
|
for origin_y, origin_x, group in _patch_origins(*pixels.shape[:2]):
|
|
grouped_values[group].append(
|
|
_patch_unit_values(
|
|
pixels,
|
|
origin_y,
|
|
origin_x,
|
|
period,
|
|
harmonics,
|
|
denoise_sigma,
|
|
)
|
|
)
|
|
amplitude, selection_patches, confirmation_patches = _joint_amplitude(
|
|
pixels,
|
|
template,
|
|
period,
|
|
denoise_sigma,
|
|
)
|
|
return RegisteredConfirmationComponents(
|
|
period=period,
|
|
joint_coherence=min(
|
|
_coherence(grouped_values[0], weights),
|
|
_coherence(grouped_values[1], weights),
|
|
),
|
|
joint_amplitude=amplitude,
|
|
unknown_codeword_fixed_confirmation=_unknown_codeword_fixed_confirmation(
|
|
grouped_values[0],
|
|
grouped_values[1],
|
|
weights,
|
|
),
|
|
selection_patches=selection_patches,
|
|
confirmation_patches=confirmation_patches,
|
|
)
|