app/deepfake/classifier.py

import os
import logging
from pathlib import Path
from typing import Union, Dict, Any, Optional
import numpy as np
from PIL import Image

logger = logging.getLogger(__name__)

DEFAULT_MODEL_ID = "prithivMLmods/Deep-Fake-Detector-v2-Model"

class DeepfakeClassifier:
    """
    Vision Transformer (ViT) based deepfake and synthetic face detector.
    Evaluates candidate media to alert investigators to manipulated, swapped,
    or synthetically generated facial content.
    """

    def __init__(self, model_name: str = DEFAULT_MODEL_ID):
        self.model_name = model_name
        self._processor = None
        self._model = None
        self._load_attempted = False
        self._available = False

    def _lazy_load(self):
        if self._load_attempted:
            return
        self._load_attempted = True

        # Memory conservation safeguard for cloud free tiers (e.g. Render 512MB RAM)
        enable_vit = os.getenv("ENABLE_VIT_MODEL", "").lower()
        if enable_vit == "false":
            logger.info("ViT model disabled via ENABLE_VIT_MODEL=false. Using lightweight spatial-FFT detector.")
            self._available = False
            return
        elif enable_vit != "true":
            try:
                import psutil
                total_mb = psutil.virtual_memory().total / (1024 * 1024)
                if total_mb < 1200:
                    logger.info("Host memory is %.0f MB (< 1200 MB). Using lightweight spatial-FFT detector to protect container stability.", total_mb)
                    self._available = False
                    return
            except Exception:
                pass

        try:
            from transformers import AutoImageProcessor, AutoModelForImageClassification
            import torch

            logger.info("Initializing ViT deepfake detector: %s", self.model_name)
            self._processor = AutoImageProcessor.from_pretrained(
                self.model_name
            )
            self._model = AutoModelForImageClassification.from_pretrained(
                self.model_name
            )
            self._model.eval()
            self._available = True
            logger.info("ViT Deepfake detector loaded successfully.")
        except Exception as exc:
            logger.warning(
                "Could not load HuggingFace ViT weights (%s): %s. "
                "Engaging resilient heuristic synthetic media analysis.",
                self.model_name, exc
            )
            self._available = False

    def _heuristic_analysis(self, image: Image.Image) -> Dict[str, Any]:
        """
        Fallback spatial-frequency analysis when offline or weights unavailable.
        Checks for spectral artifacts and boundary anomalies characteristic of synthetic generation.
        """
        try:
            gray = image.convert("L").resize((256, 256))
            arr = np.asarray(gray, dtype=np.float32)
            # 2D FFT spectral high-frequency decay
            f = np.fft.fft2(arr)
            fshift = np.fft.fftshift(f)
            magnitude = 20 * np.log(np.abs(fshift) + 1e-5)
            h, w = magnitude.shape
            center_h, center_w = h // 2, w // 2
            # High frequency energy ratio
            radius = 30
            y, x = np.ogrid[:h, :w]
            dist_from_center = np.sqrt((x - center_w)**2 + (y - center_h)**2)
            high_freq_mask = dist_from_center > radius
            high_freq_energy = float(np.mean(magnitude[high_freq_mask]))
            total_energy = float(np.mean(magnitude))
            ratio = (high_freq_energy / (total_energy + 1e-5)) if total_energy > 0 else 0.5

            # Typical synthetic GAN/diffusion faces have muted or irregular high-frequency spectra
            synthetic_score = float(np.clip(1.0 - (ratio * 1.2), 0.05, 0.95))
            is_synthetic = synthetic_score > 0.65

            return {
                "label": "Synthetic / Deepfake" if is_synthetic else "Authentic / Real",
                "risk_score": round(synthetic_score, 4),
                "confidence": round(abs(synthetic_score - 0.5) * 2, 4),
                "is_synthetic": is_synthetic,
                "model": "Spatial-FFT-Heuristic (Offline Fallback)",
                "status": "ANALYZED",
            }
        except Exception as exc:
            return {
                "label": "Undetermined",
                "risk_score": 0.5,
                "confidence": 0.0,
                "is_synthetic": False,
                "model": "None",
                "status": f"SKIPPED ({exc})",
            }

    def analyze(self, image_input: Union[Path, str, Image.Image, np.ndarray]) -> Dict[str, Any]:
        """
        Analyze an input image for deepfake or synthetic facial generation signals.
        """
        if isinstance(image_input, (str, Path)):
            img_path = Path(image_input)
            if not img_path.exists():
                return {
                    "label": "Image Not Found",
                    "risk_score": 0.0,
                    "confidence": 0.0,
                    "is_synthetic": False,
                    "model": "None",
                    "status": "ERROR",
                }
            image = Image.open(img_path).convert("RGB")
        elif isinstance(image_input, Image.Image):
            image = image_input.convert("RGB")
        elif isinstance(image_input, np.ndarray):
            image = Image.fromarray(image_input).convert("RGB")
        else:
            return {
                "label": "Invalid Format",
                "risk_score": 0.0,
                "confidence": 0.0,
                "is_synthetic": False,
                "model": "None",
                "status": "ERROR",
            }

        self._lazy_load()

        if not self._available or self._model is None or self._processor is None:
            return self._heuristic_analysis(image)

        try:
            import torch
            inputs = self._processor(images=image, return_tensors="pt")
            with torch.no_grad():
                outputs = self._model(**inputs)
                logits = outputs.logits
                probs = torch.nn.functional.softmax(logits, dim=-1).squeeze().tolist()

            id2label = self._model.config.id2label or {0: "Real", 1: "Fake"}
            # Determine fake/synthetic index
            fake_idx = 1
            for idx, label_str in id2label.items():
                if any(w in label_str.lower() for w in ["fake", "synth", "deepfake", "manipulated"]):
                    fake_idx = int(idx)
                    break

            risk_score = float(probs[fake_idx]) if isinstance(probs, list) and len(probs) > fake_idx else float(probs)
            pred_idx = int(np.argmax(probs)) if isinstance(probs, list) else (1 if probs >= 0.5 else 0)
            pred_label = id2label.get(pred_idx, "Unknown")
            is_synthetic = risk_score >= 0.50

            return {
                "label": pred_label,
                "risk_score": round(risk_score, 4),
                "confidence": round(float(np.max(probs)), 4) if isinstance(probs, list) else round(risk_score, 4),
                "is_synthetic": is_synthetic,
                "model": self.model_name,
                "status": "ANALYZED",
            }
        except Exception as exc:
            logger.warning("ViT inference failed: %s. Reverting to heuristic analysis.", exc)
            return self._heuristic_analysis(image)