Harden scale-registered SynthID detection

src/remove_ai_watermarks/_synthid_registered.py

@@ -0,0 +1,304 @@
+"""Opt-in scale registration for the measured periodic 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 itertools
+import math
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import cv2
+import numpy as np
+
+from remove_ai_watermarks.synthid_detector import folded_template_score
+
+if TYPE_CHECKING:
+    from numpy.typing import NDArray
+
+_PYRAMID_SCALES = (0.75, 1.0, 1.25)
+_SEARCH_PERIODS = np.linspace(5.0, 32.0, 541, dtype=np.float64)
+_CANONICAL_PERIODS = np.linspace(7.5, 24.5, 1701, dtype=np.float64)
+_PERIOD_THRESHOLDS = (
+    (7.5, 8.5, 0.3770629524888979),
+    (8.5, 10.0, 0.25174716660523494),
+    (10.0, 12.0, 0.284692023502354),
+    (12.0, 14.0, 0.19794247706938645),
+    (14.0, 16.0, 0.33930082812296375),
+    (16.0, 18.0, 0.28915284982686323),
+    (18.0, 20.0, 0.22885510746595789),
+    (20.0, 22.0, 0.24570317032768269),
+    (22.0, 24.5, 0.3142958338390489),
+)
+REGISTERED_HIGH_BAND_THRESHOLD = 0.075
+
+
+@dataclass(frozen=True)
+class RegisteredComponents:
+    """Calibrated components of one scale-registered decision."""
+
+    raw_score: float
+    amplitude_threshold: float
+    selected_period: float
+    spectral_period: float
+    high_band_score: float
+
+    @property
+    def decision_score(self) -> float:
+        """Return a statistic that reaches one only when every gate passes."""
+        if self.selected_period != self.spectral_period:
+            return 0.0
+        return min(
+            self.raw_score / self.amplitude_threshold,
+            self.high_band_score / REGISTERED_HIGH_BAND_THRESHOLD,
+        )
+
+
+def _resize(pixels: NDArray[Any], width: int, height: int) -> NDArray[Any]:
+    interpolation = cv2.INTER_AREA if width < pixels.shape[1] else cv2.INTER_CUBIC
+    return np.asarray(cv2.resize(pixels, (width, height), interpolation=interpolation))
+
+
+def _template_frequency_features(
+    template: NDArray[Any],
+) -> tuple[NDArray[Any], NDArray[Any], NDArray[Any]]:
+    spectrum = np.fft.fft2(template, axes=(0, 1))
+    power = np.sum(np.abs(spectrum) ** 2, axis=2)
+    power[0, 0] = 0.0
+    indices = np.argsort(power.ravel())[::-1][:30]
+    rows, columns = np.unravel_index(indices, power.shape)
+    height, width = template.shape[:2]
+    signed_rows = np.where(rows <= height // 2, rows, rows - height)
+    signed_columns = np.where(columns <= width // 2, columns, columns - width)
+    harmonics = np.column_stack((signed_rows, signed_columns)).astype(np.float64)
+    return harmonics, spectrum[rows, columns], spectrum
+
+
+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 _spectral_curve(
+    pixels: NDArray[Any],
+    periods: NDArray[Any],
+    harmonics: NDArray[Any],
+    coefficients: NDArray[Any],
+) -> NDArray[Any]:
+    height, width = pixels.shape[:2]
+    y = (periods[:, None] ** -1) * harmonics[None, :, 0] * height
+    x = (periods[:, None] ** -1) * harmonics[None, :, 1] * width
+    sampled = np.empty((len(periods), len(harmonics), 3), dtype=np.complex128)
+    for channel in range(3):
+        residual = pixels[:, :, channel].astype(np.float32)
+        residual -= cv2.GaussianBlur(
+            residual,
+            (0, 0),
+            sigmaX=1.0,
+            sigmaY=1.0,
+            borderType=cv2.BORDER_REFLECT_101,
+        )
+        spectrum = np.fft.fft2(residual)
+        sampled[:, :, channel] = _bilinear_sample(spectrum, y % height, x % width)
+    numerator = np.real(np.sum(np.conj(coefficients)[None, :, :] * sampled, axis=(1, 2)))
+    denominator = np.linalg.norm(coefficients) * np.linalg.norm(sampled, axis=(1, 2))
+    return np.divide(
+        numerator,
+        denominator,
+        out=np.zeros_like(numerator),
+        where=denominator > 0.0,
+    )
+
+
+def _period_candidates(
+    periods: NDArray[Any],
+    scores: NDArray[Any],
+    count: int = 3,
+) -> list[float]:
+    candidates: list[float] = []
+    for index in np.argsort(scores)[::-1]:
+        period = float(periods[index])
+        if any(abs(period - existing_period) < 0.25 for existing_period in candidates):
+            continue
+        candidates.append(period)
+        if len(candidates) == count:
+            break
+    return candidates
+
+
+def _period_threshold(period: float) -> float:
+    for index, (lower, upper, threshold) in enumerate(_PERIOD_THRESHOLDS):
+        if lower <= period < upper or (index == len(_PERIOD_THRESHOLDS) - 1 and period == upper):
+            return threshold
+    raise ValueError(f"registered period {period} is outside the calibrated range")
+
+
+def _high_band_score(
+    folded: NDArray[Any],
+    template_spectrum: NDArray[Any],
+) -> float:
+    folded_spectrum = np.fft.fft2(folded, axes=(0, 1))
+    tile_height, tile_width = template_spectrum.shape[:2]
+    y_coordinates = np.minimum(np.arange(tile_height), tile_height - np.arange(tile_height))
+    x_coordinates = np.minimum(np.arange(tile_width), tile_width - np.arange(tile_width))
+    radius = np.sqrt(y_coordinates[:, None] ** 2 + x_coordinates[None, :] ** 2)
+    correlations = []
+    for lower, upper in ((4.5, 6.5), (6.5, 12.0)):
+        mask = (radius >= lower) & (radius < upper)
+        selected_folded = folded_spectrum[mask]
+        selected_template = template_spectrum[mask]
+        denominator = np.linalg.norm(selected_folded) * np.linalg.norm(selected_template)
+        correlations.append(
+            float(np.real(np.vdot(selected_template, selected_folded)) / denominator) if denominator > 0.0 else 0.0
+        )
+    return min(correlations)
+
+
+def _best_canonical(
+    pixels: NDArray[Any],
+    periods: list[float],
+    template: NDArray[Any],
+    sigma: float,
+) -> tuple[float, NDArray[Any], NDArray[Any], float]:
+    best_score = -math.inf
+    best_canonical: NDArray[Any] | None = None
+    best_folded: NDArray[Any] | None = None
+    best_period: float | None = None
+    for period in periods:
+        predicted_width = round(pixels.shape[1] * template.shape[1] / period)
+        for delta in range(-4, 5):
+            width = predicted_width + delta
+            height = round(pixels.shape[0] * width / pixels.shape[1])
+            canonical = _resize(pixels, width, height)
+            score, folded = folded_template_score(canonical, template, sigma)
+            if score > best_score:
+                best_score = score
+                best_canonical = canonical
+                best_folded = folded
+                best_period = period
+    if best_canonical is None or best_folded is None or best_period is None:
+        raise RuntimeError("scale registration produced no canonical view")
+    return float(best_score), best_canonical, best_folded, best_period
+
+
+def _quadrant_median(
+    canonical: NDArray[Any],
+    template: NDArray[Any],
+    sigma: float,
+) -> float:
+    tile_height, tile_width = template.shape[:2]
+    split_y = max(tile_height, (canonical.shape[0] // (2 * tile_height)) * tile_height)
+    split_x = max(tile_width, (canonical.shape[1] // (2 * tile_width)) * tile_width)
+    scores = []
+    for region in (
+        canonical[:split_y, :split_x],
+        canonical[:split_y, split_x:],
+        canonical[split_y:, :split_x],
+        canonical[split_y:, split_x:],
+    ):
+        score, _folded = folded_template_score(region, template, sigma)
+        scores.append(score)
+    return float(np.median(scores))
+
+
+def _pyramid_locked_mean(
+    pixels: NDArray[Any],
+    harmonics: NDArray[Any],
+    coefficients: NDArray[Any],
+    base_curve: NDArray[Any],
+) -> float:
+    curves = []
+    candidates = []
+    for scale in _PYRAMID_SCALES:
+        if scale == 1.0:
+            curve = base_curve
+        else:
+            level = _resize(
+                pixels,
+                max(16, round(pixels.shape[1] * scale)),
+                max(16, round(pixels.shape[0] * scale)),
+            )
+            curve = _spectral_curve(level, _SEARCH_PERIODS, harmonics, coefficients)
+        curves.append(curve)
+        candidates.append(_period_candidates(_SEARCH_PERIODS, curve))
+    combinations = itertools.product(*candidates)
+
+    def spread(combination: tuple[float, ...]) -> float:
+        normalized_periods = [
+            candidate / scale
+            for candidate, scale in zip(
+                combination,
+                _PYRAMID_SCALES,
+                strict=True,
+            )
+        ]
+        return float(np.std(np.log(normalized_periods)))
+
+    best = min(
+        combinations,
+        key=spread,
+    )
+    base_period = float(np.median([candidate / scale for candidate, scale in zip(best, _PYRAMID_SCALES, strict=True)]))
+    locked = [
+        float(np.interp(base_period * scale, _SEARCH_PERIODS, curve))
+        for curve, scale in zip(curves, _PYRAMID_SCALES, strict=True)
+    ]
+    return float(np.mean(locked))
+
+
+def registered_components(
+    pixels: NDArray[Any],
+    template: NDArray[Any],
+    sigma: float,
+) -> RegisteredComponents:
+    """Measure a carrier after bounded scale registration."""
+    harmonics, coefficients, template_spectrum = _template_frequency_features(template)
+    combined_periods = np.concatenate((_SEARCH_PERIODS, _CANONICAL_PERIODS))
+    combined_curve = _spectral_curve(pixels, combined_periods, harmonics, coefficients)
+    base_curve = combined_curve[: len(_SEARCH_PERIODS)]
+    canonical_curve = combined_curve[len(_SEARCH_PERIODS) :]
+    candidates = _period_candidates(_CANONICAL_PERIODS, canonical_curve)
+    baseline, canonical, folded, selected_period = _best_canonical(pixels, candidates, template, sigma)
+    quadrant = _quadrant_median(canonical, template, sigma)
+    pyramid = _pyramid_locked_mean(
+        pixels,
+        harmonics,
+        coefficients,
+        base_curve,
+    )
+    raw_score = float((baseline + quadrant + pyramid) / 3.0)
+    return RegisteredComponents(
+        raw_score=raw_score,
+        amplitude_threshold=_period_threshold(selected_period),
+        selected_period=selected_period,
+        spectral_period=candidates[0],
+        high_band_score=_high_band_score(folded, template_spectrum),
+    )
+
+
+def registered_score(
+    pixels: NDArray[Any],
+    template: NDArray[Any],
+    sigma: float,
+) -> float:
+    """Return the calibrated registered decision statistic."""
+    return registered_components(pixels, template, sigma).decision_score

src/remove_ai_watermarks/cli.py

@@ -1322,7 +1322,12 @@ def cmd_video_batch(
 @main.command("detect-synthid")
 @click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path))
 @click.option("--json", "as_json", is_flag=True, help="Emit the detector result as JSON.")
-def cmd_detect_synthid(source: Path, as_json: bool) -> None:
+@click.option(
+    "--register-scale",
+    is_flag=True,
+    help="Search the slower calibrated range of spatial carrier scales.",
+)
+def cmd_detect_synthid(source: Path, as_json: bool, register_scale: bool) -> None:
     """Detect the SynthID periodic pixel carrier at calibrated image sizes.
 
     A negative result means this detector did not find its supported carrier; it
@@ -1332,7 +1337,7 @@ def cmd_detect_synthid(source: Path, as_json: bool) -> None:
 
     source = _validate_image(source)
     try:
-        result = detect_synthid(source)
+        result = detect_synthid(source, register_scale=register_scale)
     except RuntimeError as exc:
         raise click.ClickException(str(exc)) from exc
 
@@ -1346,10 +1351,15 @@ def cmd_detect_synthid(source: Path, as_json: bool) -> None:
     if result.score is not None:
         console.print(f"  Score: {result.score:.6f}  (threshold: {result.threshold:.6f})")
     console.print(f"  Detector: {result.detector}")
+    scale_scope = (
+        "  Bounded spatial-scale registration was enabled. A negative or\n"
+        if register_scale
+        else "  Arbitrary spatial resampling was not registered. A negative or\n"
+    )
     console.print(
         "  Scope: one confirmed periodic carrier family in a calibrated image-size range.\n"
-        "  Arbitrary spatial resampling is not registered. A negative or\n"
-        "  unsupported result is not proof that SynthID is absent."
+        + scale_scope
+        + "  unsupported result is not proof that SynthID is absent."
     )
 
 

src/remove_ai_watermarks/synthid_detector.py

@@ -2,8 +2,8 @@
 
 This is a positive-only detector for one measured carrier epoch, not Google's
 private payload decoder. A positive result is strong local evidence for the
-carrier. A negative result means only that this exact detector did not find it;
-image sizes outside the calibrated pixel-count range are reported separately.
+carrier. A negative result means only that the selected detector did not find
+it; image sizes outside that mode's calibrated range are reported separately.
 
 The numeric runtime requires the ``pixels`` extra. Imports remain lazy so the
 package's metadata-only paths stay dependency-light.
@@ -25,6 +25,7 @@ if TYPE_CHECKING:
 SynthIDDetectionStatus = Literal["detected", "not_detected", "unsupported"]
 
 DETECTOR_ID = "synthid-periodic-tile-v2"
+REGISTERED_DETECTOR_ID = "synthid-periodic-tile-registered-v2"
 MODEL_FILENAME = "synthid_periodic_tile_2048_v1.npz"
 # The template remains frozen at this model geometry. Runtime images are never
 # resized. The supported pixel-count interval is the separately challenged domain:
@@ -35,6 +36,12 @@ MODEL_HEIGHT = 2048
 MIN_SUPPORTED_PIXELS = 1_000_000
 MAX_SUPPORTED_PIXELS = 18_000_000
 TILE_THRESHOLD = 0.17357069773071196
+REGISTERED_MIN_SUPPORTED_PIXELS = 250_000
+REGISTERED_MAX_SUPPORTED_PIXELS = 10_000_000
+REGISTERED_MIN_SIDE = 64
+# The registered score is the minimum normalized margin across its amplitude,
+# spectral-candidate, and high-frequency agreement gates.
+REGISTERED_THRESHOLD = 1.0
 INSTALL_HINT = "install the pixel extra: uv add 'remove-ai-watermarks[pixels]'"
 
 
@@ -206,11 +213,43 @@ def _geometry_supported(width: int, height: int) -> bool:
     return MIN_SUPPORTED_PIXELS <= pixels <= MAX_SUPPORTED_PIXELS
 
 
-def detect_synthid(image_path: str | Path, *, image: NDArray[Any] | None = None) -> SynthIDDetection:
+def _registered_geometry_supported(width: int, height: int) -> bool:
+    """Whether scale registration was challenged at this decoded size."""
+    pixels = width * height
+    return (
+        min(width, height) >= REGISTERED_MIN_SIDE
+        and REGISTERED_MIN_SUPPORTED_PIXELS <= pixels <= REGISTERED_MAX_SUPPORTED_PIXELS
+    )
+
+
+def folded_template_score(
+    pixels: NDArray[Any],
+    template: NDArray[Any],
+    denoise_sigma: float,
+) -> tuple[float, NDArray[Any]]:
+    """Fold PIXELS at the model geometry and score the normalized tile."""
+    tile_height, tile_width = template.shape[:2]
+    folded = fold_residual_template(
+        pixels,
+        tile_height=tile_height,
+        tile_width=tile_width,
+        denoise_sigma=denoise_sigma,
+    )
+    normalized, _norm = unit_tile(folded)
+    return float((template * normalized).sum()), folded
+
+
+def detect_synthid(
+    image_path: str | Path,
+    *,
+    image: NDArray[Any] | None = None,
+    register_scale: bool = False,
+) -> SynthIDDetection:
     """Detect the supported periodic carrier in IMAGE_PATH.
 
     ``not_detected`` is not a clean-image guarantee. It means only that the
-    frozen periodic carrier did not cross its calibrated threshold.
+    frozen periodic carrier did not cross its calibrated threshold. Set
+    ``register_scale`` for the slower, separately calibrated resize search.
     """
     path = Path(image_path)
     if image is None:
@@ -219,13 +258,19 @@ def detect_synthid(image_path: str | Path, *, image: NDArray[Any] | None = None)
         if image.ndim != 3 or image.shape[2] != 3:
             raise ValueError("image must be a three-channel BGR array")
         height, width = image.shape[:2]
-    if not _geometry_supported(width, height):
+    geometry_supported = (
+        _registered_geometry_supported(width, height) if register_scale else _geometry_supported(width, height)
+    )
+    threshold = REGISTERED_THRESHOLD if register_scale else TILE_THRESHOLD
+    detector_id = REGISTERED_DETECTOR_ID if register_scale else DETECTOR_ID
+    if not geometry_supported:
         return SynthIDDetection(
             status="unsupported",
             width=width,
             height=height,
             score=None,
-            threshold=TILE_THRESHOLD,
+            threshold=threshold,
+            detector=detector_id,
         )
     if not is_available():
         raise RuntimeError(f"SynthID pixel detection needs numpy and OpenCV; {INSTALL_HINT}")
@@ -233,7 +278,7 @@ def detect_synthid(image_path: str | Path, *, image: NDArray[Any] | None = None)
     import numpy as np
     from PIL import Image
 
-    template, sigma, _model_height, _model_width, tile_height, tile_width = _load_template()
+    template, sigma, *_model = _load_template()
     if image is None:
         with Image.open(path) as source:
             pixels = np.asarray(source.convert("RGB"), dtype=np.uint8)
@@ -241,18 +286,17 @@ def detect_synthid(image_path: str | Path, *, image: NDArray[Any] | None = None)
         pixels = np.asarray(image[:, :, ::-1], dtype=np.uint8)
     if pixels.shape != (height, width, 3):
         raise RuntimeError("decoded image geometry does not match its header")
-    folded = fold_residual_template(
-        pixels,
-        tile_height=tile_height,
-        tile_width=tile_width,
-        denoise_sigma=sigma,
-    )
-    normalized, _norm = unit_tile(folded)
-    score = float(np.sum(template * normalized))
+    if register_scale:
+        from remove_ai_watermarks._synthid_registered import registered_score
+
+        score = registered_score(pixels, template, sigma)
+    else:
+        score, _folded = folded_template_score(pixels, template, sigma)
     return SynthIDDetection(
-        status="detected" if score >= TILE_THRESHOLD else "not_detected",
+        status="detected" if score >= threshold else "not_detected",
         width=width,
         height=height,
         score=score,
-        threshold=TILE_THRESHOLD,
+        threshold=threshold,
+        detector=detector_id,
     )