detectors/registry.py

"""
Detector registry - maps modes to detector instances with weights.
Central place to configure ensemble composition.
"""

import torch
from typing import Dict, List, Tuple
from .base import BaseDetector
from .community_forensics import CommunityForensicsDetector
from .itsnotai import ItsNotAIDetector
from .face_swap_clip import FaceSwapCLIPDetector, FaceSwapBackupDetector
from .frequency import FrequencyDetector


# Global detector cache to avoid reloading
_detector_cache: Dict[str, BaseDetector] = {}


def get_detector(detector_class: type, device: torch.device) -> BaseDetector:
    """Get or create a detector instance (singleton per class)."""
    key = detector_class.__name__
    if key not in _detector_cache:
        _detector_cache[key] = detector_class(device)
    return _detector_cache[key]


def get_ensemble_for_mode(mode: str, device: torch.device) -> Tuple[List[BaseDetector], List[float]]:
    """
    Returns (detectors, weights) for a given mode.
    Weights should sum to 1.0.
    
    AI Image Detection: Two-Pillar Ensemble (Community-Forensics ViT + ItsNotAI v2 BEiT + Frequency FFT)
    Deepfake Detection: Face-Swap CLIP + Backup + Frequency
    """
    if mode == "ai_image":
        # AI-generated images: High-accuracy two-pillar neural ensemble + spectral frequency analysis
        # Community-Forensics: ViT-Small trained on 2.7M samples from 4803 generators (high sensitivity)
        # ItsNotAI v2: BEiT-Large fine-tuned on modern generators (FLUX, Midjourney v6, DALL-E 3, SDXL)
        # Frequency: handcrafted FFT analysis for spectral checkerboard/upsampling artifacts
        detectors = [
            get_detector(CommunityForensicsDetector, device),
            get_detector(ItsNotAIDetector, device),
            get_detector(FrequencyDetector, device),
        ]
        weights = [0.48, 0.47, 0.05]
        
    elif mode == "deepfake":
        # Face-swap: CLIP-based primary, backup for diversity, frequency as minor signal
        detectors = [
            get_detector(FaceSwapCLIPDetector, device),
            get_detector(FaceSwapBackupDetector, device),
            get_detector(FrequencyDetector, device),
        ]
        # Weights: CLIP primary (SOTA cross-dataset), backup for architecture diversity, frequency minor
        weights = [0.65, 0.25, 0.10]
        
    else:
        raise ValueError(f"Unknown mode: {mode}")

    # Normalize weights
    total = sum(weights)
    weights = [w / total for w in weights]
    
    return detectors, weights


def list_available_detectors() -> Dict[str, List[str]]:
    """Return available detectors per mode for health check."""
    return {
        "ai_image": [
            "Community-Forensics (ViT-Small)",
            "ItsNotAI v2 (BEiT-Large)",
            "Frequency (FFT+LR)",
        ],
        "deepfake": [
            "Face-Swap CLIP (ViT-L/14)",
            "Face-Swap Backup (ViT-Base)",
            "Frequency (FFT+LR)",
        ],
    }