mirror of
https://github.com/wiltodelta/remove-ai-watermarks.git
synced 2026-08-10 08:00:32 +02:00
485 lines
19 KiB
Python
485 lines
19 KiB
Python
# pyright: reportUnknownMemberType=false, reportUnknownArgumentType=false, reportUnknownVariableType=false, reportMissingTypeStubs=false
|
|
"""The complete pixel-forensics layer for one image.
|
|
|
|
STATUS
|
|
|
|
Independent from provenance verdicts, removal, and the CLI. Consumers use the
|
|
versioned :meth:`PixelEvidence.to_dict` boundary; feature extraction failures are
|
|
reported per family without discarding successful measurements.
|
|
|
|
WHAT IS MEASURED
|
|
|
|
One decode, then six families of scale-robust statistics over it:
|
|
|
|
* ``dct`` -- AC coefficient histograms over the 8x8 block DCT, plus the deviation of
|
|
leading digits from Benford's law.
|
|
* ``fft`` -- radial band energies of the log-magnitude spectrum, plus the
|
|
color-filter-array periodicity peaks a demosaiced camera capture leaves.
|
|
* ``noise`` -- standard deviation and kurtosis of a high-pass residual.
|
|
* ``ela`` -- error level after a quality-90 JPEG re-save.
|
|
* ``gradient`` -- gradient-magnitude histogram and Laplacian variance.
|
|
* ``color`` -- 4x4x4 RGB histogram, mean saturation, mean value.
|
|
|
|
and, in ``artifacts``, the spatial layer those statistics are computed from: a
|
|
64-bit perceptual hash, a 128px JPEG thumbnail, and coarse ELA, noise-residual and
|
|
FFT-phase maps.
|
|
|
|
THE ARTIFACTS ARE NOT AGGREGATES
|
|
|
|
Everything above ``artifacts`` is a scalar or a fixed-length histogram, and an image
|
|
cannot be reconstructed from those. ``artifacts`` is different in kind: a thumbnail
|
|
is a picture, a perceptual hash identifies one, and the coarse maps carry layout.
|
|
Collecting them makes a record that identifies the source image, so a caller storing
|
|
or forwarding them is handling image content, not statistics about it. That is why
|
|
they are a separate field and not merged into the families.
|
|
|
|
REQUIREMENTS
|
|
|
|
Needs the ``pixels`` extra (numpy). Guard a call with :func:`is_available` when the
|
|
caller must not hard-depend on it.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import io
|
|
import logging
|
|
import time
|
|
from dataclasses import dataclass, field
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
from remove_ai_watermarks._internal.schema import require_schema_version
|
|
|
|
if TYPE_CHECKING:
|
|
from pathlib import Path
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Analysis resolution. Every statistic here is scale-robust, and a 2048px cap keeps
|
|
# the FFT and the sliding-window residual bounded on a 100 MP input.
|
|
MAX_SIDE = 2048
|
|
# The eight lowest-frequency AC positions of the 8x8 block DCT, zig-zag order.
|
|
AC_POSITIONS = ((0, 1), (1, 0), (1, 1), (0, 2), (2, 0), (2, 1), (1, 2), (0, 3))
|
|
FFT_BANDS = 8
|
|
# A Bayer CFA shows as symmetric peaks at half the Nyquist on the diagonals.
|
|
BAYER_OFFSETS = ((1, 1), (1, -1))
|
|
INSTALL_HINT = "install the pixel extra: uv add 'remove-ai-watermarks[pixels]'"
|
|
PIXEL_EVIDENCE_SCHEMA_VERSION = 1
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class PixelEvidence:
|
|
"""Pixel statistics for one image, and the spatial artifacts behind them.
|
|
|
|
``decode`` carries the source dimensions, or ``{"error": ...}`` when the image
|
|
could not be decoded -- in which case every other field is empty. A family is also
|
|
empty when the image is too small for it (the block DCT needs 8x8, the FFT 32x32,
|
|
the residual 3x3), so a caller must treat every field as optional rather than
|
|
assume a fixed feature width.
|
|
"""
|
|
|
|
path: Path
|
|
decode: dict[str, Any]
|
|
dct: dict[str, Any] = field(default_factory=dict[str, Any])
|
|
fft: dict[str, Any] = field(default_factory=dict[str, Any])
|
|
noise: dict[str, Any] = field(default_factory=dict[str, Any])
|
|
ela: dict[str, Any] = field(default_factory=dict[str, Any])
|
|
gradient: dict[str, Any] = field(default_factory=dict[str, Any])
|
|
color: dict[str, Any] = field(default_factory=dict[str, Any])
|
|
# Identifies the source image; see the module note. Empty unless asked for.
|
|
artifacts: dict[str, Any] = field(default_factory=dict[str, Any])
|
|
# Opt-in timings for callers measuring pipeline latency. Empty by default so
|
|
# repeated evidence collection remains value-deterministic.
|
|
timing_ms: dict[str, float] = field(default_factory=dict[str, float])
|
|
|
|
@property
|
|
def decoded(self) -> bool:
|
|
"""False when the source could not be decoded at all."""
|
|
return "error" not in self.decode
|
|
|
|
@property
|
|
def status(self) -> str:
|
|
"""``complete``, ``partial`` for a failed family, or ``error`` on decode."""
|
|
if not self.decoded:
|
|
return "error"
|
|
sections = (self.dct, self.fft, self.noise, self.ela, self.gradient, self.color, self.artifacts)
|
|
return "partial" if any("error" in section for section in sections) else "complete"
|
|
|
|
def to_dict(
|
|
self,
|
|
*,
|
|
schema_version: int = PIXEL_EVIDENCE_SCHEMA_VERSION,
|
|
) -> dict[str, Any]:
|
|
"""Return the selected JSON-safe transport schema without a local path."""
|
|
schema_version = require_schema_version(
|
|
schema_version,
|
|
contract="pixel evidence",
|
|
supported=(1,),
|
|
)
|
|
return {
|
|
"schema_version": schema_version,
|
|
"status": self.status,
|
|
"decode": dict(self.decode),
|
|
"dct": dict(self.dct),
|
|
"fft": dict(self.fft),
|
|
"noise": dict(self.noise),
|
|
"ela": dict(self.ela),
|
|
"gradient": dict(self.gradient),
|
|
"color": dict(self.color),
|
|
"artifacts": dict(self.artifacts),
|
|
"timing_ms": dict(self.timing_ms),
|
|
}
|
|
|
|
|
|
def is_available() -> bool:
|
|
"""True when the optional pixel dependencies are installed."""
|
|
from remove_ai_watermarks.optional_deps import module_available
|
|
|
|
return module_available("numpy")
|
|
|
|
|
|
def _numpy() -> Any:
|
|
from remove_ai_watermarks.optional_deps import module_available
|
|
|
|
if not module_available("numpy"):
|
|
raise RuntimeError(f"Pixel evidence needs numpy -- {INSTALL_HINT}")
|
|
import numpy as np
|
|
|
|
return np
|
|
|
|
|
|
def _dct_matrix(np: Any, n: int = 8) -> Any:
|
|
"""Orthonormal n x n DCT-II basis: M[i, j] = cos(pi (2j + 1) i / 2n)."""
|
|
i = np.arange(n)[:, None]
|
|
j = np.arange(n)[None, :]
|
|
m = np.cos(np.pi * (2 * j + 1) * i / (2 * n))
|
|
m[0, :] *= 1 / np.sqrt(2)
|
|
return m * np.sqrt(2 / n)
|
|
|
|
|
|
def read_gray(image_path: Path) -> tuple[Any, Any, dict[str, Any]]:
|
|
"""Decode to float32 grayscale (and RGB for color stats), downscaled.
|
|
|
|
Pillow, not cv2, and the source dimensions are recorded BEFORE the downscale.
|
|
"""
|
|
np = _numpy()
|
|
from PIL import Image
|
|
|
|
from remove_ai_watermarks import image_io
|
|
|
|
try:
|
|
image_io._register_heif() # pyright: ignore[reportPrivateUsage]
|
|
with Image.open(image_path) as img:
|
|
info: dict[str, Any] = {"width": img.width, "height": img.height}
|
|
if max(img.size) > MAX_SIDE:
|
|
img.thumbnail((MAX_SIDE, MAX_SIDE), Image.Resampling.LANCZOS)
|
|
rgb = np.asarray(img.convert("RGB"), dtype=np.float32)
|
|
gray = np.asarray(img.convert("L"), dtype=np.float32)
|
|
except Exception as exc:
|
|
logger.debug("pixel decode failed for %s: %s", image_path, exc)
|
|
# Exception text from Pillow commonly embeds the absolute source path.
|
|
# Keep that detail in the log, not in the pathless transport contract.
|
|
return None, None, {"error": type(exc).__name__}
|
|
return gray, rgb, info
|
|
|
|
|
|
def dct_features(gray: Any) -> dict[str, Any]:
|
|
"""AC coefficient histograms over the 8x8 block DCT + Benford deviation."""
|
|
np = _numpy()
|
|
height, width = gray.shape
|
|
h8, w8 = height // 8 * 8, width // 8 * 8
|
|
if h8 < 8 or w8 < 8:
|
|
return {}
|
|
basis = _dct_matrix(np)
|
|
bins = np.linspace(-20.5, 20.5, 22)
|
|
blocks = gray[:h8, :w8].reshape(h8 // 8, 8, w8 // 8, 8).swapaxes(1, 2)
|
|
rows = basis[[row for row, _ in AC_POSITIONS]]
|
|
columns = basis[[column for _, column in AC_POSITIONS]]
|
|
coeff = np.einsum("ki,abij,kj->abk", rows, blocks, columns)
|
|
hists = []
|
|
lead_vals: list[Any] = []
|
|
for index in range(len(AC_POSITIONS)):
|
|
values = coeff[:, :, index].ravel()
|
|
hists.append(np.histogram(values, bins=bins)[0].tolist())
|
|
lead_vals.append(np.abs(values))
|
|
out: dict[str, Any] = {"dct_ac_hist": hists}
|
|
flat = np.abs(np.concatenate(lead_vals))
|
|
flat = flat[flat >= 1]
|
|
if flat.size > 100:
|
|
leading = (flat / 10 ** np.floor(np.log10(flat))).astype(int)
|
|
leading = leading[(leading >= 1) & (leading <= 9)]
|
|
if leading.size > 100:
|
|
observed = np.bincount(leading, minlength=10)[1:10] / leading.size
|
|
benford = np.log10(1 + 1 / np.arange(1, 10))
|
|
out["benford_mad"] = float(np.abs(observed - benford).mean())
|
|
return out
|
|
|
|
|
|
def noise_residual_map(gray: Any) -> Any:
|
|
"""High-pass residual, the map the noise statistics are computed from."""
|
|
np = _numpy()
|
|
from numpy.lib.stride_tricks import sliding_window_view
|
|
|
|
if gray.shape[0] < 3 or gray.shape[1] < 3:
|
|
return None
|
|
kernel = np.array([[-1.0, -1.0, -1.0], [-1.0, 8.0, -1.0], [-1.0, -1.0, -1.0]])
|
|
height, width = gray.shape
|
|
# kernel is float64, so the residual is float64 like the unchunked form
|
|
out = np.empty((height - 2, width - 2), dtype=np.float64)
|
|
# Row-chunked: the (window * kernel) temporary is ~150 MB at 2048px if
|
|
# materialized whole. Per-element 9-tap sums are computed in the same order,
|
|
# so the result is bit-identical to the unchunked form.
|
|
for y0 in range(0, height - 2, 256):
|
|
y1 = min(y0 + 256, height - 2)
|
|
window = sliding_window_view(gray[y0 : y1 + 2], (3, 3))
|
|
out[y0:y1] = (window * kernel).sum(axis=(-1, -2))
|
|
return out
|
|
|
|
|
|
def noise_features(residual: Any) -> dict[str, Any]:
|
|
"""High-pass residual std and kurtosis."""
|
|
flat = residual.ravel()
|
|
std = float(flat.std())
|
|
if std < 1e-9:
|
|
return {"noise_std": 0.0, "noise_kurtosis": 0.0}
|
|
z = (flat - flat.mean()) / std
|
|
return {"noise_std": std, "noise_kurtosis": float((z**4).mean() - 3.0)}
|
|
|
|
|
|
def fft_decompose(gray: Any) -> tuple[Any, Any] | None:
|
|
"""Log-magnitude (fftshifted) and phase of the image spectrum."""
|
|
np = _numpy()
|
|
if min(gray.shape) < 32:
|
|
return None
|
|
spectrum = np.fft.fftshift(np.fft.fft2(gray - gray.mean()))
|
|
return np.log1p(np.abs(spectrum)), np.angle(spectrum)
|
|
|
|
|
|
def fft_features(mag: Any) -> dict[str, Any]:
|
|
"""Radial magnitude band energies (no phase) + CFA periodicity peaks."""
|
|
np = _numpy()
|
|
height, width = mag.shape
|
|
cy, cx = height // 2, width // 2
|
|
# 1D broadcast instead of an mgrid: saves ~160 MB of int64 temporaries at
|
|
# 2048px. The squares are exact in float64 (values < 2^53), so band means
|
|
# are identical to the mgrid form.
|
|
r2y = (np.arange(height, dtype=np.float64) - cy) ** 2
|
|
r2x = (np.arange(width, dtype=np.float64) - cx) ** 2
|
|
radius = np.sqrt(r2y[:, None] + r2x[None, :])
|
|
r_max = radius.max()
|
|
bands = []
|
|
for index in range(FFT_BANDS):
|
|
mask = (radius >= r_max * index / FFT_BANDS) & (radius < r_max * (index + 1) / FFT_BANDS)
|
|
bands.append(float(mag[mask].mean()) if mask.any() else 0.0)
|
|
peaks = []
|
|
for dy, dx in BAYER_OFFSETS:
|
|
y, x = cy + dy * (height // 4), cx + dx * (width // 4)
|
|
neighborhood = mag[y - 2 : y + 3, x - 2 : x + 3]
|
|
peaks.append(float(neighborhood.max() - mag.mean()))
|
|
return {"fft_band_energy": bands, "cfa_peaks": peaks, "cfa_peak": max(peaks)}
|
|
|
|
|
|
def ela_map(rgb: Any) -> Any:
|
|
"""Absolute per-pixel error after a quality-90 JPEG re-save."""
|
|
np = _numpy()
|
|
from PIL import Image
|
|
|
|
try:
|
|
buffer = io.BytesIO()
|
|
Image.fromarray(rgb.astype(np.uint8)).save(buffer, "JPEG", quality=90)
|
|
buffer.seek(0)
|
|
resaved = np.asarray(Image.open(buffer).convert("RGB"), dtype=np.float32)
|
|
except Exception as exc:
|
|
logger.debug("ELA re-save failed: %s", exc)
|
|
return None
|
|
if resaved.shape != rgb.shape:
|
|
return None
|
|
return np.abs(rgb - resaved).mean(axis=-1)
|
|
|
|
|
|
def ela_features(err: Any) -> dict[str, Any]:
|
|
"""Error-level stats after a quality-90 JPEG re-save."""
|
|
np = _numpy()
|
|
return {"ela_mean": float(err.mean()), "ela_p95": float(np.percentile(err, 95))}
|
|
|
|
|
|
def gradient_features(gray: Any) -> dict[str, Any]:
|
|
np = _numpy()
|
|
gy, gx = np.gradient(gray)
|
|
mag = np.sqrt(gx**2 + gy**2)
|
|
hist = np.histogram(mag, bins=10, range=(0, 255))[0].tolist()
|
|
laplacian = np.gradient(gy, axis=0) + np.gradient(gx, axis=1)
|
|
return {"gradient_hist": hist, "laplacian_var": float(laplacian.var())}
|
|
|
|
|
|
def color_features(rgb: Any) -> dict[str, Any]:
|
|
np = _numpy()
|
|
small = rgb[::4, ::4] # decimate; the histogram is position-blind anyway
|
|
bins = (small / 256 * 4).astype(int).clip(0, 3)
|
|
index = bins[..., 0] * 16 + bins[..., 1] * 4 + bins[..., 2]
|
|
hist = np.bincount(index.ravel(), minlength=64).tolist()
|
|
mx = small.max(axis=-1)
|
|
mn = small.min(axis=-1)
|
|
saturation = np.where(mx > 0, (mx - mn) / np.maximum(mx, 1e-6), 0)
|
|
return {
|
|
"color_hist_4x4x4": hist,
|
|
"saturation_mean": float(saturation.mean()),
|
|
"value_mean": float(mx.mean() / 255),
|
|
}
|
|
|
|
|
|
def perceptual_hash(gray: Any) -> str:
|
|
"""64-bit DCT perceptual hash. Identifies an image; see the module note."""
|
|
np = _numpy()
|
|
from PIL import Image
|
|
|
|
small = np.asarray(Image.fromarray(gray.astype(np.float32), mode="F").resize((32, 32), Image.Resampling.LANCZOS))
|
|
basis = _dct_matrix(np, 32)
|
|
low_basis = basis[:8]
|
|
low = (low_basis @ small @ low_basis.T).ravel()[1:] # drop DC
|
|
bits = low > np.median(low)
|
|
return f"{int(''.join('1' if bit else '0' for bit in bits), 2):016x}"
|
|
|
|
|
|
def _coarse(np: Any, arr: Any, side: int = 64) -> Any:
|
|
"""Downscale a 2D map to at most ``side`` on the long edge."""
|
|
from PIL import Image
|
|
|
|
height, width = arr.shape
|
|
if max(height, width) <= side:
|
|
return arr
|
|
img = Image.fromarray(arr.astype(np.float32), mode="F")
|
|
img.thumbnail((side, side), Image.Resampling.BILINEAR)
|
|
return np.asarray(img)
|
|
|
|
|
|
def _array_payload(arr: Any) -> dict[str, Any]:
|
|
return {
|
|
"shape": list(arr.shape),
|
|
"dtype": str(arr.dtype),
|
|
"base64": base64.b64encode(arr.tobytes()).decode("ascii"),
|
|
}
|
|
|
|
|
|
def spatial_artifacts(gray: Any, rgb: Any, *, ela: Any, residual: Any, phase: Any) -> dict[str, Any]:
|
|
"""Perceptual hash, thumbnail, and coarse ELA / residual / phase maps.
|
|
|
|
These identify the source image rather than describe it -- see the module note.
|
|
The maps are the ones the statistics were computed from, passed in rather than
|
|
recomputed.
|
|
"""
|
|
np = _numpy()
|
|
from PIL import Image
|
|
|
|
out: dict[str, Any] = {"phash": perceptual_hash(gray)}
|
|
thumbnail = Image.fromarray(rgb.astype(np.uint8))
|
|
thumbnail.thumbnail((128, 128), Image.Resampling.LANCZOS)
|
|
buffer = io.BytesIO()
|
|
thumbnail.save(buffer, "JPEG", quality=70)
|
|
out["thumbnail_jpeg_b64"] = base64.b64encode(buffer.getvalue()).decode("ascii")
|
|
|
|
if ela is not None:
|
|
out["ela_map"] = _array_payload(_coarse(np, ela))
|
|
if residual is not None:
|
|
clipped = np.clip(residual / 4.0, -1, 1)
|
|
out["noise_residual"] = _array_payload(_coarse(np, (clipped * 127).astype(np.int8)))
|
|
if phase is not None:
|
|
out["fft_phase"] = _array_payload(_coarse(np, phase.astype(np.float32), 32))
|
|
return out
|
|
|
|
|
|
def extract_pixel_evidence(image_path: Path, *, artifacts: bool = False, timings: bool = False) -> PixelEvidence:
|
|
"""Measure every pixel-statistic family for one image in a single decode.
|
|
|
|
The image is decoded ONCE and the intermediate maps (high-pass residual, ELA
|
|
error, FFT magnitude and phase) are computed once and shared, because the
|
|
residual's sliding window and the ELA re-save are the two expensive steps and
|
|
each family would otherwise redo them.
|
|
|
|
A family that fails or does not apply is left empty rather than raising: an
|
|
undecodable file, or one too small for the block DCT, still returns a
|
|
:class:`PixelEvidence` whose ``decoded`` / empty fields say so. Missing numpy is
|
|
the one hard error, since then nothing can be measured at all.
|
|
|
|
Args:
|
|
image_path: Path to the image. Any container Pillow can open.
|
|
artifacts: Also return the spatial layer -- perceptual hash, thumbnail and
|
|
coarse maps. Off by default: those identify the source image, so asking
|
|
for them is a decision the caller makes explicitly.
|
|
timings: Measure each stage and include rounded milliseconds in
|
|
:attr:`PixelEvidence.timing_ms`.
|
|
|
|
Returns:
|
|
A :class:`PixelEvidence`.
|
|
"""
|
|
started = time.perf_counter()
|
|
stage_started = started
|
|
measured: dict[str, float] = {}
|
|
|
|
gray, rgb, info = read_gray(image_path)
|
|
measured["decode"] = time.perf_counter() - stage_started
|
|
if gray is None or rgb is None:
|
|
measured["total"] = time.perf_counter() - started
|
|
timing_ms = {name: round(seconds * 1000, 1) for name, seconds in measured.items()} if timings else {}
|
|
return PixelEvidence(path=image_path, decode=info, timing_ms=timing_ms)
|
|
|
|
families: dict[str, dict[str, Any]] = {}
|
|
|
|
residual = None
|
|
stage_started = time.perf_counter()
|
|
try:
|
|
residual = noise_residual_map(gray)
|
|
families["noise"] = noise_features(residual) if residual is not None else {}
|
|
except Exception as exc:
|
|
logger.debug("pixel family noise failed for %s: %s", image_path, exc)
|
|
families["noise"] = {"error": type(exc).__name__}
|
|
measured["noise"] = time.perf_counter() - stage_started
|
|
|
|
spectrum = None
|
|
stage_started = time.perf_counter()
|
|
try:
|
|
spectrum = fft_decompose(gray)
|
|
families["fft"] = fft_features(spectrum[0]) if spectrum is not None else {}
|
|
except Exception as exc:
|
|
logger.debug("pixel family fft failed for %s: %s", image_path, exc)
|
|
families["fft"] = {"error": type(exc).__name__}
|
|
measured["fft"] = time.perf_counter() - stage_started
|
|
|
|
error = None
|
|
stage_started = time.perf_counter()
|
|
try:
|
|
error = ela_map(rgb)
|
|
families["ela"] = ela_features(error) if error is not None else {}
|
|
except Exception as exc:
|
|
logger.debug("pixel family ela failed for %s: %s", image_path, exc)
|
|
families["ela"] = {"error": type(exc).__name__}
|
|
measured["ela"] = time.perf_counter() - stage_started
|
|
|
|
for name, compute in (
|
|
("dct", lambda: dct_features(gray)),
|
|
("gradient", lambda: gradient_features(gray)),
|
|
("color", lambda: color_features(rgb)),
|
|
):
|
|
stage_started = time.perf_counter()
|
|
try:
|
|
families[name] = compute()
|
|
except Exception as exc: # one bad family must not lose the other five
|
|
logger.debug("pixel family %s failed for %s: %s", name, image_path, exc)
|
|
families[name] = {"error": type(exc).__name__}
|
|
measured[name] = time.perf_counter() - stage_started
|
|
|
|
if artifacts:
|
|
stage_started = time.perf_counter()
|
|
try:
|
|
families["artifacts"] = spatial_artifacts(
|
|
gray, rgb, ela=error, residual=residual, phase=spectrum[1] if spectrum is not None else None
|
|
)
|
|
except Exception as exc:
|
|
logger.debug("pixel artifacts failed for %s: %s", image_path, exc)
|
|
families["artifacts"] = {"error": type(exc).__name__}
|
|
measured["full_artifacts"] = time.perf_counter() - stage_started
|
|
|
|
measured["total"] = time.perf_counter() - started
|
|
timing_ms = {name: round(seconds * 1000, 1) for name, seconds in measured.items()} if timings else {}
|
|
return PixelEvidence(path=image_path, decode=info, timing_ms=timing_ms, **families)
|