ai_detector.py

"""
ai_detector.py - Detection of AI-generated documents (Sora / Midjourney /
Stable Diffusion / DALL-E / DocGen outputs).

The 2026 threat vector is no longer Photoshop. Banks worry about applicants
submitting documents synthesized by generative models. These outputs leave
fingerprints in the frequency domain that classical forensics miss:

  1. Suppressed high-frequency energy ("smoothness")
     Generative models tend to produce images with less true high-frequency
     content than real scans, because their upsampling layers act as
     low-pass filters.

  2. Periodic spectral peaks
     Discrete-cosine-transform-based diffusion outputs and GAN outputs leave
     periodic peaks in the FFT magnitude spectrum (the "checkerboard
     artefact") - especially at frequencies corresponding to the upsampling
     stride.

  3. Missing JPEG quantization tables
     Real scanned documents are JPEGs with standard quantization tables.
     Generative model outputs are PNGs or re-saved JPEGs with non-standard
     quantization patterns.

This module implements all three as a single "AI-generated probability"
score in [0, 1].

Public API:
    detect_ai_generated(path)       - returns {'probability', 'confidence',
                                        'spectrum', 'flags', 'verdict'}

    radial_fft_profile(image)       - per-frequency log-magnitude profile
    spectral_peak_score(profile)    - peakiness of the high-frequency band
    high_freq_attenuation(profile)  - degree of high-frequency suppression
"""

import io
import numpy as np
from PIL import Image
import cv2


# ============================================================
# Helpers
# ============================================================

def _to_grayscale_float(path_or_img, max_dim=512):
    """Load + downscale + convert to grayscale float32 in [0, 1]."""
    if isinstance(path_or_img, (str, bytes)) or hasattr(path_or_img, "__fspath__"):
        img = Image.open(path_or_img).convert("L")
    else:
        img = path_or_img.convert("L")
    w, h = img.size
    scale = min(1.0, max_dim / max(w, h))
    if scale < 1.0:
        img = img.resize((int(w * scale), int(h * scale)), Image.LANCZOS)
    arr = np.asarray(img, dtype=np.float32) / 255.0
    return arr


# ============================================================
# Spectral analysis
# ============================================================

def radial_fft_profile(gray, n_bins=64):
    """
    Compute the radially-averaged log-magnitude FFT spectrum.

    Returns: numpy array of shape (n_bins,) with mean log-magnitude per
             radial frequency bin. Low index = low frequency.
    """
    H, W = gray.shape
    # 2D FFT, shifted so DC is at center
    F = np.fft.fftshift(np.fft.fft2(gray))
    mag = np.log1p(np.abs(F))

    # Radial coordinate map
    cy, cx = H // 2, W // 2
    yy, xx = np.indices((H, W))
    r = np.sqrt((xx - cx) ** 2 + (yy - cy) ** 2)
    r_max = min(H, W) // 2

    bins = np.linspace(0, r_max, n_bins + 1)
    profile = np.zeros(n_bins, dtype=np.float32)
    for i in range(n_bins):
        mask = (r >= bins[i]) & (r < bins[i + 1])
        if mask.any():
            profile[i] = mag[mask].mean()
    return profile


def high_freq_attenuation(profile):
    """
    Measure how much energy is missing from the top half of the spectrum.
    Real scans have roughly 1/f decay; AI outputs are much smoother
    (sharper drop-off).

    Returns: a score in [0, 1] where 1 = very suppressed (suspicious).
    """
    n = len(profile)
    low = profile[: n // 4].mean()        # low frequencies
    high = profile[3 * n // 4:].mean()    # high frequencies
    if low < 1e-6:
        return 0.0
    ratio = high / low
    # In real scans this ratio is typically 0.45-0.65.
    # In AI-generated images it drops below 0.30.
    if ratio >= 0.45: return 0.0
    if ratio <= 0.20: return 1.0
    return float((0.45 - ratio) / 0.25)


def spectral_peak_score(profile):
    """
    Detect periodic peaks in the high-frequency band - a signature of
    checkerboard upsampling in GANs/diffusion.

    Returns: a score in [0, 1] where 1 = lots of peaks (suspicious).
    """
    n = len(profile)
    high_band = profile[n // 2:]
    if len(high_band) < 4:
        return 0.0
    # Differences from local trend
    smooth = np.convolve(high_band, np.ones(3) / 3, mode="same")
    deviations = high_band - smooth
    # Count peaks (positive spikes > 1 std above local mean)
    spikes = (deviations > deviations.std()).sum()
    # Normalise: ~6 spikes in 32 bins is typical for AI output
    return float(min(1.0, spikes / 6.0))


def jpeg_quantization_check(path):
    """
    For images saved as JPEG, examine the quantization tables.
    Real scanners write standard tables; AI outputs often have all-1s
    (no real compression history) or unusual table structures.

    Returns: dict with 'has_jpeg_qtable', 'is_standard'.
    """
    try:
        img = Image.open(path)
        if img.format != "JPEG":
            return {"has_jpeg_qtable": False, "is_standard": None,
                    "note": f"Not a JPEG (format: {img.format})"}
        qtables = getattr(img, "quantization", {})
        if not qtables:
            return {"has_jpeg_qtable": False, "is_standard": False,
                    "note": "JPEG without quantization tables - suspicious"}
        # Standard tables have mean value > 5 (real compression)
        means = [np.array(q).mean() for q in qtables.values()]
        is_standard = all(m > 5 for m in means)
        return {"has_jpeg_qtable": True, "is_standard": bool(is_standard),
                "note": f"qtable means: {[round(m, 1) for m in means]}"}
    except Exception as e:
        return {"has_jpeg_qtable": False, "is_standard": None,
                "note": f"qtable read failed: {e}"}


# ============================================================
# Main entry point
# ============================================================

def detect_ai_generated(path):
    """
    Run the full AI-generated-content detection pipeline.

    Returns:
        {
          'probability': float in [0, 1],   - overall AI-generated probability
          'confidence': str,                 - 'low' | 'medium' | 'high'
          'verdict': str,                    - 'likely_real' | 'suspicious' | 'likely_ai_generated'
          'profile': list[float],            - radial FFT profile
          'sub': dict,                       - individual detector scores
          'flags': list[str],                - human-readable evidence
        }
    """
    gray = _to_grayscale_float(path)
    profile = radial_fft_profile(gray)

    s_hf  = high_freq_attenuation(profile)
    s_pk  = spectral_peak_score(profile)
    q     = jpeg_quantization_check(path)
    s_jpg = 0.6 if (q["has_jpeg_qtable"] is False and q.get("note", "").endswith("not a JPEG")) else \
            (0.4 if q.get("is_standard") is False else 0.0)

    # Weighted blend
    prob = 0.50 * s_hf + 0.30 * s_pk + 0.20 * s_jpg

    flags = []
    if s_hf > 0.5:
        flags.append(f"High-frequency content suppressed (smoothness score {s_hf:.2f}) - typical of generative-model output.")
    if s_pk > 0.5:
        flags.append(f"Periodic peaks detected in high-frequency band (score {s_pk:.2f}) - GAN/diffusion checkerboard signature.")
    if q.get("is_standard") is False:
        flags.append(f"Non-standard JPEG quantization tables ({q.get('note','')}).")
    if not flags:
        flags.append("No AI-generated indicators above threshold.")

    if prob >= 0.65:
        verdict = "likely_ai_generated"; confidence = "high"
    elif prob >= 0.40:
        verdict = "suspicious"; confidence = "medium"
    else:
        verdict = "likely_real"; confidence = "low"

    return {
        "probability": round(float(prob), 3),
        "confidence": confidence,
        "verdict": verdict,
        "profile": profile.tolist(),
        "sub": {
            "high_freq_suppression": round(float(s_hf), 3),
            "spectral_peakiness": round(float(s_pk), 3),
            "jpeg_artefact_score": round(float(s_jpg), 3),
        },
        "jpeg_qtable": q,
        "flags": flags,
    }


if __name__ == "__main__":
    import sys
    if len(sys.argv) < 2:
        # Smoke test on synthetic inputs
        from PIL import ImageDraw
        # 1. Smooth/blurry image (AI-like)
        smooth = Image.new("RGB", (400, 300), "white")
        d = ImageDraw.Draw(smooth)
        d.text((40, 120), "AI-LIKE smooth output", fill="black")
        smooth.save("/tmp/_smooth.png")
        # 2. Noisy/real-like image
        noisy_arr = (np.random.RandomState(0).randint(0, 256, (300, 400, 3))
                     .astype(np.uint8))
        noisy = Image.fromarray(noisy_arr)
        noisy.save("/tmp/_noisy.png")

        for tag, p in [("AI-like smooth", "/tmp/_smooth.png"),
                        ("Noisy real-like", "/tmp/_noisy.png")]:
            r = detect_ai_generated(p)
            print(f"  {tag:18s} -> verdict={r['verdict']:20s} "
                   f"prob={r['probability']:.2f} sub={r['sub']}")
    else:
        import json
        r = detect_ai_generated(sys.argv[1])
        print(json.dumps({k: v for k, v in r.items() if k != "profile"}, indent=2))