video_model.py

"""
video_model.py — Spécialiste deepfake VIDÉO du Système Schrödinger (CPU).

Architecture : ResNext50_32x4d (backbone, tronqué avant avgpool/fc) -> features
2048 par frame -> LSTM(2048 -> 2048, 1 couche) -> Linear(2048 -> 2).
Checkpoint : Naman712/Deep-fake-detection / model_87_acc_20_frames_final_data.pt
(~87 % de précision, entraîné sur des séquences de 20 crops de visage 112×112).
Convention des classes : **classe 1 = FAKE** — constaté empiriquement sur ce
checkpoint (probs=[0.001, 0.999] sur une vidéo deepfake avérée) ; le projet
GitHub d'origine documentait l'inverse. Surchargeable via VIDEO_FAKE_INDEX.

TÉLÉCHARGEMENT + CHARGEMENT PARESSEUX, thread-safe : le checkpoint est récupéré
depuis HuggingFace ET chargé au tout premier appel à predict() (première analyse
vidéo), pas au démarrage du serveur. Un verrou (self._lock) sérialise l'opération
-> pas de téléchargements simultanés si plusieurs vidéos arrivent en même temps.
Tant que le modèle est indisponible (pas de HF_TOKEN, repo gated…), predict()
renvoie (0.5, False) et la fusion vidéo se rabat sur le pipeline forensique image.

CONTRAT D'ENTRÉE (fidélité au régime d'entraînement du checkpoint) : le projet
d'origine lisait les frames avec OpenCV et les passait à ToPILImage SANS
conversion BGR->RGB — le modèle a donc appris en ordre de canaux **BGR**, sur
des frames **consécutives** et des crops de visage **serrés**. predict() attend
donc des frames BGR (telles que décodées par OpenCV). Si un autre checkpoint
entraîné en RGB est branché, définir VIDEO_MODEL_CHANNELS=rgb.

Variables d'environnement :
  VIDEO_MODEL_PATH      chemin .pt local explicite (court-circuite models/)
  VIDEO_FAKE_INDEX      index de la classe « fake » (défaut 1 — checkpoint Naman712)
  VIDEO_MODEL_CHANNELS  ordre de canaux attendu par le checkpoint : bgr (défaut) | rgb
  TORCH_THREADS         threads intra-op PyTorch (défaut 2 — HF Spaces 2 vCPU)
"""

import os
import threading

import numpy as np
import cv2

_BASE_DIR = os.path.dirname(os.path.abspath(__file__))
_MODELS_DIR = os.path.join(_BASE_DIR, "models")
_DEBUG = os.getenv("UNKOR_DEBUG", "1").lower() not in ("0", "false", "no")

_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32)
_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32)
_IM_SIZE = 112          # taille d'entrée du projet d'origine


def _dbg(msg):
    if _DEBUG:
        try:
            print(f"[video-model] {msg}", flush=True)
        except UnicodeEncodeError:
            print(f"[video-model] {msg}".encode("ascii", "replace").decode("ascii"), flush=True)


def _build_model(num_classes):
    """Reconstruit l'architecture exacte du checkpoint (projet ResNext+LSTM)."""
    import torch.nn as nn
    from torchvision import models as tv_models

    class ResNextLSTM(nn.Module):
        def __init__(self, num_classes, latent_dim=2048, lstm_layers=1,
                     hidden_dim=2048, bidirectional=False):
            super().__init__()
            backbone = tv_models.resnext50_32x4d(weights=None)   # poids fournis par le checkpoint
            self.model = nn.Sequential(*list(backbone.children())[:-2])
            self.lstm = nn.LSTM(latent_dim, hidden_dim, lstm_layers, bidirectional)
            self.relu = nn.LeakyReLU()
            self.dp = nn.Dropout(0.4)
            self.linear1 = nn.Linear(2048, num_classes)
            self.avgpool = nn.AdaptiveAvgPool2d(1)

        def forward(self, x):
            b, t, c, h, w = x.shape
            x = x.view(b * t, c, h, w)
            fmap = self.model(x)
            x = self.avgpool(fmap)
            x = x.view(b, t, 2048)
            x_lstm, _ = self.lstm(x, None)
            return fmap, self.dp(self.linear1(x_lstm[:, -1, :]))

    return ResNextLSTM(num_classes)


class VideoDeepfakeAnalyzer:
    """ResNext50+LSTM sur une séquence de crops de visage — (proba_fake, ok)."""

    _MODEL = os.path.join(_MODELS_DIR, "video_resnext_lstm.pt")

    def __init__(self):
        self._model = None
        self._tried = False          # tentative (téléchargement + chargement) déjà faite ?
        self._lock = threading.Lock()
        self._fake_idx = int(os.getenv("VIDEO_FAKE_INDEX", "1"))

    def _resolve_or_download(self):
        """Chemin du .pt : explicite -> local -> téléchargement paresseux HF.
        Renvoie None si indisponible (l'analyse vidéo se rabat alors sur le
        pipeline image). Appelé UNIQUEMENT sous self._lock (voir _ensure)."""
        p = os.getenv("VIDEO_MODEL_PATH")
        if p and os.path.exists(p):
            return p
        if os.path.exists(self._MODEL):
            return self._MODEL
        # Téléchargement paresseux (1er predict) — repo gated : HF_TOKEN requis.
        token = os.getenv("HF_TOKEN") or os.getenv("HUGGINGFACE_TOKEN") or None
        if not token:
            print("[VIDEO] HF_TOKEN absent — modèle vidéo non téléchargé (repo gated). "
                  "Analyse vidéo = pipeline image seul.", flush=True)
            return None
        repo = os.getenv("VIDEO_MODEL_REPO", "Naman712/Deep-fake-detection")
        fname = os.getenv("VIDEO_MODEL_FILE", "model_87_acc_20_frames_final_data.pt")
        try:
            import shutil
            from huggingface_hub import hf_hub_download
            os.makedirs(_MODELS_DIR, exist_ok=True)
            print(f"[VIDEO] téléchargement paresseux {repo}/{fname} "
                  "(~220 Mo, 1re analyse vidéo)…", flush=True)
            f = hf_hub_download(repo_id=repo, filename=fname, token=token)
            if os.path.getsize(f) < 10_000_000:      # pointeur LFS/HTML = quelques Ko
                print("[VIDEO] téléchargement invalide (fichier trop petit).", flush=True)
                return None
            shutil.copy(f, self._MODEL)
            print(f"[VIDEO] modèle téléchargé -> {os.path.basename(self._MODEL)} "
                  f"({os.path.getsize(self._MODEL) // (1024 * 1024)} Mo)", flush=True)
            return self._MODEL
        except Exception as e:
            print(f"[VIDEO] téléchargement échoué : {e}", flush=True)
            return None

    def available(self):
        self._ensure()
        return self._model is not None

    def _ensure(self):
        if self._model is not None:
            return
        # Verrou : sérialise téléchargement + chargement -> pas de DL simultanés.
        with self._lock:
            if self._model is not None or self._tried:
                return
            self._tried = True          # une seule tentative (DL inclus) par process
            path = self._resolve_or_download()
            if not path:
                return
            try:
                import torch
                torch.set_num_threads(max(1, int(os.getenv("TORCH_THREADS", "2"))))
                # weights_only=True : refuse tout pickle arbitraire (state_dict pur).
                sd = torch.load(path, map_location="cpu", weights_only=True)
                if isinstance(sd, dict) and "state_dict" in sd and isinstance(sd["state_dict"], dict):
                    sd = sd["state_dict"]
                sd = {k[7:] if k.startswith("module.") else k: v for k, v in sd.items()}
                if "linear1.weight" not in sd:
                    raise RuntimeError("clé linear1.weight absente — architecture inattendue")
                num_classes = int(sd["linear1.weight"].shape[0])
                model = _build_model(num_classes)
                missing, unexpected = model.load_state_dict(sd, strict=False)
                critical = [k for k in missing if k.split(".")[0] in ("model", "lstm", "linear1")]
                if critical:
                    raise RuntimeError(f"poids manquants: {critical[:5]} (+{max(0, len(critical)-5)})")
                model.eval()
                self._model = model
                print(f"[VIDEO] ResNext50+LSTM chargé : {os.path.basename(path)} "
                      f"(classes={num_classes}, fake_idx={self._fake_idx}, "
                      f"inattendues={len(unexpected)})", flush=True)
            except Exception as e:
                print(f"[VIDEO] chargement échoué : {e}", flush=True)
                self._model = None

    def _preprocess(self, frames):
        # L'appelant fournit du BGR (décodage OpenCV) ; on ne convertit en RGB
        # QUE si le checkpoint branché a été entraîné en RGB.
        to_rgb = os.getenv("VIDEO_MODEL_CHANNELS", "bgr").lower() == "rgb"
        seq = []
        for f in frames:
            img = f
            if img is None or not isinstance(img, np.ndarray) or img.ndim < 2:
                continue
            if img.ndim == 2:
                img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
            elif img.shape[2] == 4:
                img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
            if img.shape[0] < 8 or img.shape[1] < 8:
                continue
            if to_rgb:
                img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
            img = cv2.resize(img, (_IM_SIZE, _IM_SIZE), interpolation=cv2.INTER_AREA)
            img = img.astype(np.float32) / 255.0
            img = (img - _MEAN) / _STD
            seq.append(np.transpose(img, (2, 0, 1)))
        return seq

    def predict(self, frames):
        """frames : liste de np.ndarray BGR (crops visage serrés, consécutifs de
        préférence — voir le contrat d'entrée en tête de module).
        -> (proba_fake dans [0.02, 0.98], ok)."""
        self._ensure()
        if self._model is None:
            return 0.5, False
        try:
            seq = self._preprocess(frames)
            if len(seq) < 4:
                _dbg(f"séquence trop courte ({len(seq)} frames exploitables) — modèle ignoré")
                return 0.5, False
            import torch
            x = torch.from_numpy(np.ascontiguousarray(
                np.stack(seq)[None, ...].astype(np.float32)))     # (1, T, 3, 112, 112)
            with torch.no_grad():
                _, logits = self._model(x)
                probs = torch.softmax(logits, dim=1)[0].tolist()
            p = float(probs[min(self._fake_idx, len(probs) - 1)])
            _dbg(f"T={len(seq)} probs={[round(v, 3) for v in probs]} "
                 f"fake_idx={self._fake_idx} -> p_fake={p:.3f}")
            return float(np.clip(p, 0.02, 0.98)), True
        except Exception as e:
            print(f"[VIDEO] inférence échouée : {e}", flush=True)
            return 0.5, False