src/deepfake_scanner/detectors/ensemble.py

"""
Detector ensemble.

Holds the registry of enabled detectors and aggregates their per-class
contributions into the final 4-class probability vector returned to the API.

Aggregation strategy (Stage 1)
------------------------------
Simple averaging of the per-class contributions across detectors that emit
non-empty contributions, then renormalisation. When only one detector is
active (Stage 1), the output is just its contributions (with the absent
classes filled with 0.0 and a tiny epsilon to keep softmax-style downstream
math safe).

Stage 2+ may switch to:
  • weighted averaging (per-detector confidence / calibration)
  • a small meta-classifier that takes detector scores as input
"""
from __future__ import annotations

from dataclasses import dataclass

from PIL import Image

from ..api.schemas import DetectorSignal, Probabilities, Verdict
from ..config import settings
from .base import Detector, DetectorResult
from .clip_classifier import ClipClassifier
from .face_swap import FaceSwapDetector
from .frequency import FrequencyDetector

# The fixed taxonomy. Order matters — used for deterministic argmax tie-break.
_CLASSES: tuple[str, ...] = ("authentic", "ai_generated", "deepfake", "edited")


@dataclass
class EnsembleOutput:
    probabilities: Probabilities
    signals: list[DetectorSignal]
    verdict: Verdict
    confidence: float


def _build_detectors() -> list[Detector]:
    """Instantiate the enabled detectors at process start."""
    detectors: list[Detector] = []
    if settings.enable_clip_detector:
        detectors.append(ClipClassifier())
    if settings.enable_frequency_detector:
        detectors.append(FrequencyDetector())
    if settings.enable_face_swap_detector:
        detectors.append(FaceSwapDetector())
    if not detectors:
        raise RuntimeError(
            "No detectors enabled. At minimum, ENABLE_CLIP_DETECTOR must be true."
        )
    return detectors


# Module-level singleton — built once at import.
_DETECTORS: list[Detector] = _build_detectors()


def _aggregate(results: list[DetectorResult]) -> Probabilities:
    """Average per-class contributions across detectors, renormalise."""
    sums = {c: 0.0 for c in _CLASSES}
    counts = {c: 0 for c in _CLASSES}

    for r in results:
        for cls, val in r.contributions.items():
            if cls in sums:
                sums[cls] += val
                counts[cls] += 1

    averaged = {
        c: (sums[c] / counts[c]) if counts[c] > 0 else 0.0 for c in _CLASSES
    }
    total = sum(averaged.values())
    if total <= 0.0:
        # No detector contributed — default to maximum uncertainty over the
        # two Stage 1 classes (graceful degradation).
        averaged = {"authentic": 0.5, "ai_generated": 0.5, "deepfake": 0.0, "edited": 0.0}
    else:
        averaged = {c: v / total for c, v in averaged.items()}

    return Probabilities(**averaged)


def _verdict_from(probs: Probabilities) -> tuple[Verdict, float]:
    """Pick the top class as the verdict; return it with its probability.

    If the top probability is below 0.55, return 'uncertain' to discourage
    callers from over-trusting low-confidence outputs.
    """
    items = [
        ("authentic", probs.authentic),
        ("ai_generated", probs.ai_generated),
        ("deepfake", probs.deepfake),
        ("edited", probs.edited),
    ]
    items.sort(key=lambda kv: kv[1], reverse=True)
    top_class, top_prob = items[0]
    if top_prob < 0.55:
        return "uncertain", top_prob
    return top_class, top_prob  # type: ignore[return-value]


def run_ensemble(image: Image.Image) -> EnsembleOutput:
    """Run all enabled detectors and aggregate their outputs."""
    results = [d.run(image) for d in _DETECTORS]
    probs = _aggregate(results)
    signals = [
        DetectorSignal(name=r.name, score=r.score, notes=r.notes) for r in results
    ]
    verdict, confidence = _verdict_from(probs)
    return EnsembleOutput(
        probabilities=probs,
        signals=signals,
        verdict=verdict,
        confidence=confidence,
    )