{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.13.15","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","isInternetEnabled":false,"language":"python","sourceType":"notebook"}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"daf072d7","cell_type":"markdown","source":"# ForensicX: Multi-Domain Deepfake Detection (Kaggle Edition)\nHybrid ConvNeXt-V2 + Swin-V2 detector with a frequency encoder, rPPG biological cues, evidential uncertainty and explainability. Enable a GPU, attach FF++, Celeb-DF and/or DFDC, then run all cells. If no dataset is found, a small synthetic set is generated for a smoke test.","metadata":{}},{"id":"43bc70d0","cell_type":"markdown","source":"## Section 1: Environment Setup and Kaggle Configuration\nThis section imports the pre-installed libraries, fixes random seeds, selects the GPU with mixed precision and defines paths and hyperparameters. Datasets under `/kaggle/input/` are detected automatically; use `DATASET_OVERRIDES` if your folder names differ.","metadata":{}},{"id":"4d8fd57e","cell_type":"code","source":"import gc\nimport json\nimport math\nimport multiprocessing as mp\nimport os\nimport random\nimport re\nimport time\nimport traceback\nimport warnings\nfrom pathlib import Path\nfrom typing import Any, Dict, List, Optional, Sequence, Tuple\n\nimport albumentations as A\nimport cv2\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport pandas as pd\nimport seaborn as sns\nimport timm\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom einops import rearrange\nfrom scipy.signal import detrend\nfrom sklearn.metrics import accuracy_score, f1_score, roc_auc_score\nfrom sklearn.model_selection import GroupShuffleSplit\nfrom torch.utils.checkpoint import checkpoint\nfrom torch.utils.data import DataLoader, Dataset\nfrom tqdm.auto import tqdm\n\nwarnings.filterwarnings(\"ignore\")\nsns.set_theme(style=\"whitegrid\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-05T03:44:28.040139Z","iopub.execute_input":"2026-10-05T03:44:28.040484Z","iopub.status.idle":"2026-10-05T03:45:17.886857Z","shell.execute_reply.started":"2026-10-05T03:44:28.040458Z","shell.execute_reply":"2026-10-05T03:45:17.886245Z"}},"outputs":[],"execution_count":null},{"id":"87927671","cell_type":"code","source":"SEED: int = 42\nOFFLINE_MODE: bool = True  # Kaggle has no internet: fail fast on hub downloads.\nif OFFLINE_MODE:\n    os.environ.setdefault(\"HF_HUB_OFFLINE\", \"1\")\n\n\ndef set_seed(seed: int = SEED) -> None:\n    \"\"\"Seeds all random number generators.\n\n    Args:\n        seed: Seed value shared by python, numpy and torch.\n    \"\"\"\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n\n\ndef print_gpu_memory(tag: str = \"\") -> None:\n    \"\"\"Prints allocated and reserved GPU memory.\n\n    Args:\n        tag: Label describing the point of measurement.\n    \"\"\"\n    if not torch.cuda.is_available():\n        return\n    allocated = torch.cuda.memory_allocated() / 1024 ** 3\n    reserved = torch.cuda.memory_reserved() / 1024 ** 3\n    print(f\"[GPU {tag}] allocated={allocated:.2f} GB, reserved={reserved:.2f} GB\")\n\n\nset_seed()\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nUSE_AMP: bool = DEVICE.type == \"cuda\"\ntorch.backends.cudnn.benchmark = True\nprint(f\"Device: {DEVICE} | AMP: {USE_AMP} | torch {torch.__version__} | timm {timm.__version__}\")\nif DEVICE.type == \"cuda\":\n    print(f\"GPU: {torch.cuda.get_device_name(0)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-05T03:45:17.88813Z","iopub.execute_input":"2026-10-05T03:45:17.889142Z","iopub.status.idle":"2026-10-05T03:45:18.235982Z","shell.execute_reply.started":"2026-10-05T03:45:17.889095Z","shell.execute_reply":"2026-10-05T03:45:18.235351Z"}},"outputs":[],"execution_count":null},{"id":"c128bb95","cell_type":"code","source":"IMAGE_SIZE: int = 224\nNUM_FRAMES: int = 8\nBATCH_SIZE: int = 16\nNUM_EPOCHS: int = 30\nLEARNING_RATE: float = 1e-4\nWEIGHT_DECAY: float = 1e-2\nEARLY_STOP_PATIENCE: int = 5\nNUM_WORKERS: int = 2\nFREEZE_STAGES: int = 1            # Early backbone stages kept frozen to save memory.\nUSE_GRAD_CHECKPOINT: bool = True  # Recompute backbone activations in backward.\nPRETRAINED: bool = True           # Falls back to random init when weights are unavailable.\nMAX_VIDEOS_PER_CLASS: Optional[int] = 400  # Keeps an epoch short enough for a 9-hour session.\nTIME_BUDGET_HOURS: float = 8.5    # Training stops cleanly before Kaggle's 9-hour limit.","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-05T03:45:18.236828Z","iopub.execute_input":"2026-10-05T03:45:18.237122Z","iopub.status.idle":"2026-10-05T03:45:18.241649Z","shell.execute_reply.started":"2026-10-05T03:45:18.237103Z","shell.execute_reply":"2026-10-05T03:45:18.240993Z"}},"outputs":[],"execution_count":null},{"id":"5755f59c","cell_type":"code","source":"IS_KAGGLE: bool = Path(\"/kaggle\").exists()\nINPUT_ROOT: Path = Path(\"/kaggle/input\") if IS_KAGGLE else Path(\"./input\")\nWORKING_ROOT: Path = Path(\"/kaggle/working\") if IS_KAGGLE else Path(\"./working\")\n\nKAGGLE_PATHS: Dict[str, Path] = {\n    \"input\": INPUT_ROOT,\n    \"working\": WORKING_ROOT,\n    \"checkpoints\": WORKING_ROOT / \"checkpoints\",\n    \"xai\": WORKING_ROOT / \"xai_outputs\",\n    \"results\": WORKING_ROOT / \"results\",\n    \"synthetic\": WORKING_ROOT / \"synthetic_data\",\n    \"log_csv\": WORKING_ROOT / \"training_log.csv\",\n}\n\nDATASET_KEYWORDS: Dict[str, Tuple[str, ...]] = {\n    \"ffpp\": (\"faceforensics\", \"ff++\", \"ffpp\", \"ff-c23\", \"ff_c23\", \"c23\"),\n    \"celeb_df\": (\"celeb\",),\n    \"dfdc\": (\"dfdc\", \"deepfake-detection-challenge\"),\n}\n# Example: DATASET_OVERRIDES = {\"ffpp\": \"/kaggle/input/my-ff-dataset\"}\nDATASET_OVERRIDES: Dict[str, str] = {\n    \"ffpp\": \"/kaggle/input/datasets/xdxd003/ff-c23\",\n    \"celeb_df\": \"/kaggle/input/datasets/reubensuju/celeb-df-v2\",\n    \"dfdc\": \"/kaggle/input/competitions/deepfake-detection-challenge\",\n}\n\n\ndef create_output_dirs() -> None:\n    \"\"\"Creates all output directories under the working directory.\"\"\"\n    for key, path in KAGGLE_PATHS.items():\n        if key not in (\"input\", \"log_csv\"):\n            path.mkdir(parents=True, exist_ok=True)\n\n\ndef detect_datasets(input_root: Path = INPUT_ROOT, max_depth: int = 2) -> Dict[str, Path]:\n    \"\"\"Detects known deepfake datasets by folder name.\n\n    Args:\n        input_root: Root directory containing attached Kaggle datasets.\n        max_depth: Maximum directory depth to inspect.\n\n    Returns:\n        Mapping from dataset key (ffpp, celeb_df, dfdc) to its root path.\n    \"\"\"\n    found: Dict[str, Path] = {}\n    if not input_root.exists():\n        return found\n    base_depth = len(input_root.parts)\n    for current, dirs, _ in os.walk(input_root):\n        for name in list(dirs):\n            lowered = name.lower()\n            for key, keywords in DATASET_KEYWORDS.items():\n                if key not in found and any(k in lowered for k in keywords):\n                    found[key] = Path(current) / name\n        if len(Path(current).parts) - base_depth >= max_depth:\n            dirs[:] = []\n    return found\n\n\ncreate_output_dirs()\nDATASET_PATHS: Dict[str, Path] = {**detect_datasets(), **{k: Path(v) for k, v in DATASET_OVERRIDES.items()}}\nprint(\"Detected datasets:\", {k: str(v) for k, v in DATASET_PATHS.items()} or \"none (synthetic fallback will be used)\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-05T03:45:18.24347Z","iopub.execute_input":"2026-10-05T03:45:18.243726Z","iopub.status.idle":"2026-10-05T03:45:18.263321Z","shell.execute_reply.started":"2026-10-05T03:45:18.243709Z","shell.execute_reply":"2026-10-05T03:45:18.262696Z"}},"outputs":[],"execution_count":null},{"id":"a1c49c02","cell_type":"markdown","source":"## Section 2: Dataset Class (Kaggle-Optimized)\n`ForensicXDataset` decodes videos on demand, samples frames uniformly, crops faces (MTCNN or Haar, with center-crop fallback) and optionally applies JPEG compression. It also extracts a green-channel rPPG trace per video and returns `(frames, rppg, label)`.","metadata":{}},{"id":"e85d6153","cell_type":"code","source":"VIDEO_EXTS: Tuple[str, ...] = (\".mp4\", \".avi\", \".mov\", \".mkv\", \".webm\")\nIMAGENET_MEAN: Tuple[float, float, float] = (0.485, 0.456, 0.406)\nIMAGENET_STD: Tuple[float, float, float] = (0.229, 0.224, 0.225)\nBox = Tuple[int, int, int, int]\n\n\nclass FaceDetector:\n    \"\"\"Face detector using MTCNN when requested and available, else Haar cascades.\"\"\"\n\n    def __init__(self, backend: str = \"haar\", margin: float = 0.25, detect_size: int = 480) -> None:\n        \"\"\"Initializes the detector.\n\n        Args:\n            backend: Either \"haar\" or \"mtcnn\".\n            margin: Fractional box expansion on each side.\n            detect_size: Longest side used for detection (speed/accuracy trade-off).\n        \"\"\"\n        self.margin = margin\n        self.detect_size = detect_size\n        self.backend = \"none\"\n        self._haar: Optional[cv2.CascadeClassifier] = None\n        self._mtcnn: Any = None\n        if backend == \"mtcnn\":\n            try:\n                from facenet_pytorch import MTCNN\n\n                self._mtcnn = MTCNN(keep_all=False, device=\"cpu\", post_process=False)\n                self.backend = \"mtcnn\"\n            except Exception:\n                backend = \"haar\"\n        if backend == \"haar\":\n            cascade = cv2.CascadeClassifier(cv2.data.haarcascades + \"haarcascade_frontalface_default.xml\")\n            if not cascade.empty():\n                self._haar = cascade\n                self.backend = \"haar\"\n\n    def _raw_detect(self, frame_bgr: np.ndarray) -> Optional[Box]:\n        \"\"\"Detects the primary face without margin.\n\n        Args:\n            frame_bgr: Full frame in BGR order.\n\n        Returns:\n            Box (x1, y1, x2, y2) in original resolution, or None.\n        \"\"\"\n        height, width = frame_bgr.shape[:2]\n        scale = min(1.0, self.detect_size / max(height, width))\n        small = cv2.resize(frame_bgr, None, fx=scale, fy=scale, interpolation=cv2.INTER_AREA) if scale < 1.0 else frame_bgr\n        if self._mtcnn is not None:\n            boxes, _ = self._mtcnn.detect(cv2.cvtColor(small, cv2.COLOR_BGR2RGB))\n            if boxes is None or len(boxes) == 0:\n                return None\n            x1, y1, x2, y2 = boxes[0]\n        elif self._haar is not None:\n            gray = cv2.cvtColor(small, cv2.COLOR_BGR2GRAY)\n            faces = self._haar.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(32, 32))\n            if len(faces) == 0:\n                return None\n            x, y, fw, fh = max(faces, key=lambda f: f[2] * f[3])\n            x1, y1, x2, y2 = x, y, x + fw, y + fh\n        else:\n            return None\n        return tuple(int(v / scale) for v in (x1, y1, x2, y2))  # type: ignore[return-value]\n\n    def detect(self, frame_bgr: np.ndarray) -> Optional[Box]:\n        \"\"\"Detects the primary face and expands the box by the margin.\n\n        Args:\n            frame_bgr: Full frame in BGR order.\n\n        Returns:\n            Clipped box (x1, y1, x2, y2) or None when no face is found.\n        \"\"\"\n        box = self._raw_detect(frame_bgr)\n        if box is None:\n            return None\n        x1, y1, x2, y2 = box\n        height, width = frame_bgr.shape[:2]\n        mx, my = int((x2 - x1) * self.margin), int((y2 - y1) * self.margin)\n        return max(0, x1 - mx), max(0, y1 - my), min(width, x2 + mx), min(height, y2 + my)\n\n\ndef center_crop(frame: np.ndarray, ratio: float = 0.8) -> np.ndarray:\n    \"\"\"Returns a centered square crop used when face detection fails.\n\n    Args:\n        frame: Image array of shape [H, W, C].\n        ratio: Crop side as a fraction of the shorter image side.\n\n    Returns:\n        Cropped image.\n    \"\"\"\n    height, width = frame.shape[:2]\n    side = int(min(height, width) * ratio)\n    y0, x0 = (height - side) // 2, (width - side) // 2\n    return frame[y0:y0 + side, x0:x0 + side]\n\n\ndef jpeg_compress(image: np.ndarray, quality: int = 50) -> np.ndarray:\n    \"\"\"Applies JPEG compression to simulate low-quality (LQ) video.\n\n    Args:\n        image: BGR uint8 image.\n        quality: JPEG quality in [1, 100].\n\n    Returns:\n        Compressed and decoded image.\n    \"\"\"\n    ok, encoded = cv2.imencode(\".jpg\", image, [int(cv2.IMWRITE_JPEG_QUALITY), quality])\n    return cv2.imdecode(encoded, cv2.IMREAD_COLOR) if ok else image\n\n\ndef build_transforms(train: bool) -> A.ReplayCompose:\n    \"\"\"Builds albumentations transforms replayed identically across frames.\n\n    Args:\n        train: Whether to include random augmentations.\n\n    Returns:\n        A ReplayCompose pipeline ending with ImageNet normalization.\n    \"\"\"\n    ops: List[Any] = []\n    if train:\n        ops += [A.HorizontalFlip(p=0.5), A.RandomBrightnessContrast(p=0.3), A.HueSaturationValue(p=0.2)]\n    ops.append(A.Normalize(mean=IMAGENET_MEAN, std=IMAGENET_STD))\n    return A.ReplayCompose(ops)\n\n\ndef worker_init_fn(worker_id: int) -> None:\n    \"\"\"Seeds each DataLoader worker and limits OpenCV threads.\n\n    Args:\n        worker_id: Index of the worker process.\n    \"\"\"\n    seed = (torch.initial_seed() + worker_id) % (2 ** 32)\n    np.random.seed(seed)\n    random.seed(seed)\n    cv2.setNumThreads(0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-05T03:45:18.264364Z","iopub.execute_input":"2026-10-05T03:45:18.264609Z","iopub.status.idle":"2026-10-05T03:45:18.286244Z","shell.execute_reply.started":"2026-10-05T03:45:18.264592Z","shell.execute_reply":"2026-10-05T03:45:18.285686Z"}},"outputs":[],"execution_count":null},{"id":"3614e8fd","cell_type":"code","source":"class ForensicXDataset(Dataset):\n    \"\"\"Video dataset that decodes, face-crops and normalizes frames on demand.\"\"\"\n\n    _MAX_WARNINGS: int = 3\n\n    def __init__(\n        self,\n        video_paths: Sequence[str],\n        labels: Sequence[int],\n        transform: Optional[A.ReplayCompose] = None,\n        num_frames: int = NUM_FRAMES,\n        quality_mode: str = \"HQ\",\n        face_detector: Optional[FaceDetector] = None,\n    ) -> None:\n        \"\"\"Initializes the dataset.\n\n        Args:\n            video_paths: Paths to video files.\n            labels: Labels per video (0 = real, 1 = fake).\n            transform: Optional albumentations ReplayCompose pipeline.\n            num_frames: Number of uniformly sampled frames per video.\n            quality_mode: \"HQ\", \"LQ\" (JPEG q=50) or \"MIXED\" (LQ with p=0.5).\n            face_detector: Detector exposing `detect`; None means center crop.\n        \"\"\"\n        if len(video_paths) != len(labels):\n            raise ValueError(\"video_paths and labels must have the same length.\")\n        if quality_mode not in (\"HQ\", \"LQ\", \"MIXED\"):\n            raise ValueError(f\"Unknown quality_mode: {quality_mode}\")\n        self.video_paths = [str(p) for p in video_paths]\n        self.labels = [int(label) for label in labels]\n        self.transform = transform\n        self.num_frames = num_frames\n        self.quality_mode = quality_mode\n        self.face_detector = face_detector\n        self._warnings = 0\n\n    def __len__(self) -> int:\n        \"\"\"Returns the number of videos.\"\"\"\n        return len(self.video_paths)\n\n    def _read_frames(self, path: str) -> List[np.ndarray]:\n        \"\"\"Reads `num_frames` uniformly spaced frames from a video.\n\n        Args:\n            path: Video file path.\n\n        Returns:\n            List of BGR frames of length `num_frames`.\n\n        Raises:\n            IOError: If the video cannot be opened or decoded.\n        \"\"\"\n        cap = cv2.VideoCapture(path)\n        try:\n            if not cap.isOpened():\n                raise IOError(f\"Cannot open video: {path}\")\n            total = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))\n            if total <= 0:\n                raise IOError(f\"No frames in video: {path}\")\n            indices = np.linspace(0, total - 1, self.num_frames).round().astype(int)\n            frames: List[np.ndarray] = []\n            last: Optional[np.ndarray] = None\n            for index in indices:\n                cap.set(cv2.CAP_PROP_POS_FRAMES, int(index))\n                ok, frame = cap.read()\n                frame = frame if ok else last\n                if frame is not None:\n                    frames.append(frame)\n                    last = frame\n        finally:\n            cap.release()\n        if not frames:\n            raise IOError(f\"Could not decode any frame: {path}\")\n        while len(frames) < self.num_frames:\n            frames.append(frames[-1])\n        return frames\n\n    def _extract_faces(self, frames: List[np.ndarray], low_quality: bool) -> List[np.ndarray]:\n        \"\"\"Crops faces, resizes to IMAGE_SIZE and optionally compresses them.\n\n        Args:\n            frames: BGR frames.\n            low_quality: Whether to apply JPEG compression (quality=50).\n\n        Returns:\n            List of BGR crops of shape [IMAGE_SIZE, IMAGE_SIZE, 3].\n        \"\"\"\n        crops: List[np.ndarray] = []\n        last_box: Optional[Box] = None\n        for frame in frames:\n            box = self.face_detector.detect(frame) if self.face_detector is not None else None\n            box = box if box is not None else last_box  # Reuse last box to keep crops temporally stable.\n            crop = center_crop(frame)\n            if box is not None:\n                x1, y1, x2, y2 = box\n                candidate = frame[y1:y2, x1:x2]\n                if candidate.size > 0:\n                    crop, last_box = candidate, box\n            crop = cv2.resize(crop, (IMAGE_SIZE, IMAGE_SIZE), interpolation=cv2.INTER_AREA)\n            crops.append(jpeg_compress(crop, 50) if low_quality else crop)\n        return crops\n\n    @staticmethod\n    def extract_rppg(crops: List[np.ndarray]) -> np.ndarray:\n        \"\"\"Extracts a normalized green-channel rPPG trace from face crops.\n\n        Args:\n            crops: BGR face crops.\n\n        Returns:\n            Zero-mean, unit-variance array of shape [num_frames].\n        \"\"\"\n        signal = np.array([crop[:, :, 1].mean() for crop in crops], dtype=np.float64)\n        if len(signal) > 2:\n            signal = detrend(signal)\n        return ((signal - signal.mean()) / (signal.std() + 1e-6)).astype(np.float32)\n\n    def _apply_transform(self, rgb_frames: List[np.ndarray]) -> torch.Tensor:\n        \"\"\"Applies the same augmentation to all frames and stacks them.\n\n        Args:\n            rgb_frames: RGB uint8 frames.\n\n        Returns:\n            Float tensor of shape [T, 3, IMAGE_SIZE, IMAGE_SIZE].\n        \"\"\"\n        if self.transform is None:\n            out = [((f.astype(np.float32) / 255.0) - IMAGENET_MEAN) / IMAGENET_STD for f in rgb_frames]\n        else:\n            first = self.transform(image=rgb_frames[0])\n            out = [first[\"image\"]] + [\n                A.ReplayCompose.replay(first[\"replay\"], image=f)[\"image\"] for f in rgb_frames[1:]\n            ]\n        return torch.from_numpy(np.stack(out).astype(np.float32)).permute(0, 3, 1, 2).contiguous()\n\n    def __getitem__(self, index: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n        \"\"\"Loads one video.\n\n        Args:\n            index: Sample index.\n\n        Returns:\n            Tuple of frames [T, 3, H, W], rPPG [T] and label (scalar long tensor).\n        \"\"\"\n        label = torch.tensor(self.labels[index], dtype=torch.long)\n        try:\n            low_quality = self.quality_mode == \"LQ\" or (self.quality_mode == \"MIXED\" and random.random() < 0.5)\n            crops = self._extract_faces(self._read_frames(self.video_paths[index]), low_quality)\n            rppg = self.extract_rppg(crops)\n            rgb = [cv2.cvtColor(c, cv2.COLOR_BGR2RGB) for c in crops]\n            return self._apply_transform(rgb), torch.from_numpy(rppg), label\n        except Exception as exc:  # A single corrupt video must not abort an epoch.\n            if self._warnings < self._MAX_WARNINGS:\n                print(f\"[WARN] failed to load {self.video_paths[index]}: {exc}\")\n                self._warnings += 1\n            frames = torch.zeros(self.num_frames, 3, IMAGE_SIZE, IMAGE_SIZE)\n            return frames, torch.zeros(self.num_frames), label","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-05T03:45:18.287209Z","iopub.execute_input":"2026-10-05T03:45:18.287845Z","iopub.status.idle":"2026-10-05T03:45:18.307006Z","shell.execute_reply.started":"2026-10-05T03:45:18.287817Z","shell.execute_reply":"2026-10-05T03:45:18.306131Z"}},"outputs":[],"execution_count":null},{"id":"8785fdab","cell_type":"markdown","source":"## Section 3: Model Architecture (Memory-Efficient for Kaggle)\nThe model combines a frequency encoder, ConvNeXt-V2 and Swin-V2 branches, hierarchical cross-attention fusion, a temporal transformer and a gated rPPG branch. Pretrained weights are used when found locally under `/kaggle/input/`; otherwise random initialization is used.","metadata":{}},{"id":"18e27600","cell_type":"code","source":"TOKEN_GRID: int = 7  # Spatial tokens per stage = TOKEN_GRID ** 2.\n\n\ndef to_nchw(feat: torch.Tensor, channels: int) -> torch.Tensor:\n    \"\"\"Converts channels-last feature maps (used by Swin) to NCHW.\n\n    Args:\n        feat: Feature map, either [B, C, H, W] or [B, H, W, C].\n        channels: Expected channel count.\n\n    Returns:\n        Feature map in [B, C, H, W] layout.\n    \"\"\"\n    if feat.dim() == 4 and feat.shape[1] != channels and feat.shape[-1] == channels:\n        return feat.permute(0, 3, 1, 2).contiguous()\n    return feat\n\n\ndef find_local_weights(name: str, root: Path = INPUT_ROOT, max_depth: int = 3) -> Optional[Path]:\n    if not root.exists():\n        return None\n    base_depth = len(root.parts)\n    \n    # More flexible matching\n    search_terms = {\n        \"convnextv2_tiny\": [\"convnextv2_tiny\", \"convnextv2-tiny\", \"convnext_v2_tiny\"],\n        \"swinv2_tiny_window8_256\": [\"swinv2_tiny\", \"swin_v2_tiny\", \"swinv2_tiny_patch4_window8\", \"swinv2_tiny_window8\"],\n    }\n    \n    terms = search_terms.get(name, [name])\n    \n    for current, dirs, files in os.walk(root):\n        for file_name in files:\n            lowered = file_name.lower()\n            if any(term in lowered for term in terms) and lowered.endswith((\".pth\", \".pt\", \".safetensors\", \".bin\")):\n                return Path(current) / file_name\n        if len(Path(current).parts) - base_depth >= max_depth:\n            dirs[:] = []\n    return None\n\n\ndef create_backbone(name: str, pretrained: bool, **kwargs: Any) -> nn.Module:\n    \"\"\"Creates a timm feature extractor with graceful offline fallbacks.\n\n    Args:\n        name: timm model name.\n        pretrained: Whether to try loading pretrained weights.\n        **kwargs: Extra arguments forwarded to `timm.create_model`.\n\n    Returns:\n        A features_only backbone.\n\n    Raises:\n        RuntimeError: If the model cannot be created at all.\n    \"\"\"\n    attempts: List[Dict[str, Any]] = []\n    if pretrained:\n        local = find_local_weights(name)\n        if local is not None:\n            attempts.append({\"pretrained\": True, \"pretrained_cfg_overlay\": {\"file\": str(local)}})\n        attempts.append({\"pretrained\": True})\n    attempts.append({\"pretrained\": False})\n    last_error: Optional[Exception] = None\n    for extra in attempts:\n        try:\n            model = timm.create_model(name, features_only=True, **extra, **kwargs)\n            print(f\"[INFO] {name}: pretrained={extra['pretrained']}\")\n            return model\n        except Exception as exc:\n            last_error = exc\n            print(f\"[WARN] {name} (pretrained={extra['pretrained']}) failed: {type(exc).__name__}\")\n    raise RuntimeError(f\"Could not create backbone {name}\") from last_error\n\n\ndef freeze_early_stages(module: nn.Module, num_stages: int) -> None:\n    \"\"\"Freezes the stem and the first `num_stages` stages of a backbone.\n\n    Args:\n        module: timm features_only backbone.\n        num_stages: Number of early stages to freeze.\n    \"\"\"\n    pattern = re.compile(r\"(?:stages|layers)[._](\\d+)\")\n    for name, param in module.named_parameters():\n        match = pattern.search(name)\n        if match is not None:\n            param.requires_grad = int(match.group(1)) >= num_stages\n        elif \"stem\" in name or \"patch_embed\" in name:\n            param.requires_grad = num_stages == 0\n\n\nclass FrequencyEncoder(nn.Module):\n    \"\"\"Encodes FFT magnitude/phase and DCT coefficients with a shallow CNN.\"\"\"\n\n    def __init__(self, image_size: int = IMAGE_SIZE, out_channels: int = 64) -> None:\n        \"\"\"Initializes the encoder.\n\n        Args:\n            image_size: Expected square input resolution.\n            out_channels: Number of output channels.\n        \"\"\"\n        super().__init__()\n        self.register_buffer(\"dct_matrix\", self._dct_matrix(image_size), persistent=False)\n        self.input_norm = nn.BatchNorm2d(9)\n        self.net = nn.Sequential(\n            nn.Conv2d(9, 32, 3, stride=2, padding=1, bias=False), nn.BatchNorm2d(32), nn.ReLU(inplace=True),\n            nn.Conv2d(32, 48, 3, stride=2, padding=1, bias=False), nn.BatchNorm2d(48), nn.ReLU(inplace=True),\n            nn.Conv2d(48, out_channels, 3, padding=1, bias=False), nn.BatchNorm2d(out_channels), nn.ReLU(inplace=True),\n        )\n\n    @staticmethod\n    def _dct_matrix(size: int) -> torch.Tensor:\n        \"\"\"Builds the orthonormal DCT-II basis (torch.fft has no dct).\n\n        Args:\n            size: Transform length.\n\n        Returns:\n            Matrix of shape [size, size].\n        \"\"\"\n        k = torch.arange(size, dtype=torch.float32).unsqueeze(1)\n        n = torch.arange(size, dtype=torch.float32).unsqueeze(0)\n        basis = torch.cos(math.pi * (2 * n + 1) * k / (2 * size)) * math.sqrt(2.0 / size)\n        basis[0] = basis[0] / math.sqrt(2.0)\n        return basis\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        \"\"\"Extracts frequency-domain features.\n\n        Args:\n            x: Images of shape [B, 3, H, W].\n\n        Returns:\n            Features of shape [B, out_channels, H/4, W/4].\n        \"\"\"\n        with torch.autocast(device_type=x.device.type, enabled=False):\n            x = x.float()\n            spectrum = torch.fft.fftshift(torch.fft.fft2(x, norm=\"ortho\"), dim=(-2, -1))\n            magnitude = torch.log1p(spectrum.abs())\n            phase = torch.angle(spectrum) / math.pi\n            dct_h = self.dct_matrix if x.shape[-2] == self.dct_matrix.shape[0] else self._dct_matrix(x.shape[-2]).to(x.device)\n            dct_w = self.dct_matrix if x.shape[-1] == self.dct_matrix.shape[0] else self._dct_matrix(x.shape[-1]).to(x.device)\n            dct = torch.matmul(torch.matmul(dct_h, x), dct_w.t())\n            dct = torch.sign(dct) * torch.log1p(dct.abs())\n            return self.net(self.input_norm(torch.cat([magnitude, phase, dct], dim=1)))\n\n\nclass HierarchicalCrossAttention(nn.Module):\n    \"\"\"Multi-head cross-attention with a learnable gate and residual norm.\"\"\"\n\n    def __init__(self, dim: int, num_heads: int = 8, dropout: float = 0.1) -> None:\n        \"\"\"Initializes the module.\n\n        Args:\n            dim: Token dimension.\n            num_heads: Number of attention heads.\n            dropout: Attention dropout probability.\n        \"\"\"\n        super().__init__()\n        if dim % num_heads != 0:\n            raise ValueError(\"dim must be divisible by num_heads\")\n        self.num_heads = num_heads\n        self.head_dim = dim // num_heads\n        self.dropout = dropout\n        self.q_proj = nn.Linear(dim, dim)\n        self.kv_proj = nn.Linear(dim, 2 * dim)\n        self.out_proj = nn.Linear(dim, dim)\n        self.gate = nn.Sequential(nn.Linear(2 * dim, dim), nn.Sigmoid())\n        self.norm = nn.LayerNorm(dim)\n        self._use_sdpa = hasattr(F, \"scaled_dot_product_attention\")\n\n    def forward(self, query: torch.Tensor, context: torch.Tensor) -> torch.Tensor:\n        \"\"\"Attends from `query` tokens to `context` tokens.\n\n        Args:\n            query: Tensor [B, N, D].\n            context: Tensor [B, M, D].\n\n        Returns:\n            Updated query tokens [B, N, D].\n        \"\"\"\n        b, n, d = query.shape\n        m = context.shape[1]\n        q = self.q_proj(query).view(b, n, self.num_heads, self.head_dim).transpose(1, 2)\n        k, v = self.kv_proj(context).view(b, m, 2, self.num_heads, self.head_dim).permute(2, 0, 3, 1, 4)\n        p = self.dropout if self.training else 0.0\n        if self._use_sdpa:\n            attended = F.scaled_dot_product_attention(q, k, v, dropout_p=p)\n        else:\n            scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)\n            attended = torch.matmul(F.dropout(scores.softmax(-1), p, self.training), v)\n        attended = self.out_proj(attended.transpose(1, 2).reshape(b, n, d))\n        gate = self.gate(torch.cat([query, attended], dim=-1))\n        return self.norm(query + gate * attended)\n\n\nclass StageFusion(nn.Module):\n    \"\"\"Bidirectional cross-attention fusion of ConvNeXt and Swin tokens.\"\"\"\n\n    def __init__(self, dim: int, num_heads: int, dropout: float) -> None:\n        \"\"\"Initializes the fusion block.\n\n        Args:\n            dim: Token dimension.\n            num_heads: Number of attention heads.\n            dropout: Attention dropout probability.\n        \"\"\"\n        super().__init__()\n        self.conv_to_swin = HierarchicalCrossAttention(dim, num_heads, dropout)\n        self.swin_to_conv = HierarchicalCrossAttention(dim, num_heads, dropout)\n        self.merge = nn.Sequential(nn.Linear(2 * dim, dim), nn.GELU(), nn.LayerNorm(dim))\n\n    def forward(self, conv_tokens: torch.Tensor, swin_tokens: torch.Tensor) -> torch.Tensor:\n        \"\"\"Fuses two token sets.\n\n        Args:\n            conv_tokens: ConvNeXt tokens [B, N, D].\n            swin_tokens: Swin tokens [B, N, D].\n\n        Returns:\n            Fused tokens [B, N, D].\n        \"\"\"\n        c = self.conv_to_swin(conv_tokens, swin_tokens)\n        s = self.swin_to_conv(swin_tokens, conv_tokens)\n        return self.merge(torch.cat([c, s], dim=-1))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-05T03:45:18.308059Z","iopub.execute_input":"2026-10-05T03:45:18.308359Z","iopub.status.idle":"2026-10-05T03:45:18.334585Z","shell.execute_reply.started":"2026-10-05T03:45:18.308335Z","shell.execute_reply":"2026-10-05T03:45:18.333907Z"}},"outputs":[],"execution_count":null},{"id":"9cddca9f","cell_type":"code","source":"class ConvNeXtV2Branch(nn.Module):\n    \"\"\"ConvNeXt-V2 tiny feature extractor returning four stage maps.\"\"\"\n\n    def __init__(self, pretrained: bool = PRETRAINED, freeze_stages: int = FREEZE_STAGES) -> None:\n        \"\"\"Initializes the branch.\n\n        Args:\n            pretrained: Whether to load pretrained weights.\n            freeze_stages: Number of early stages to freeze.\n        \"\"\"\n        super().__init__()\n        self.backbone = create_backbone(\"convnextv2_tiny\", pretrained)\n        self.channels: List[int] = list(self.backbone.feature_info.channels())\n        if freeze_stages > 0:\n            freeze_early_stages(self.backbone, freeze_stages)\n\n    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]:\n        \"\"\"Extracts multi-scale features.\n\n        Args:\n            x: Images [N, 3, H, W].\n\n        Returns:\n            Tuple of four NCHW feature maps.\n        \"\"\"\n        return tuple(to_nchw(f, c) for f, c in zip(self.backbone(x), self.channels))\n\n\nclass SwinTransformerV2Branch(nn.Module):\n    \"\"\"Swin-V2 tiny (window 8, 256px) feature extractor returning four stage maps.\"\"\"\n\n    INPUT_SIZE: int = 256\n\n    def __init__(self, pretrained: bool = PRETRAINED, freeze_stages: int = FREEZE_STAGES) -> None:\n        \"\"\"Initializes the branch.\n\n        Args:\n            pretrained: Whether to load pretrained weights.\n            freeze_stages: Number of early stages to freeze.\n        \"\"\"\n        super().__init__()\n        self.backbone = create_backbone(\"swinv2_tiny_window8_256\", pretrained)\n        self.channels: List[int] = list(self.backbone.feature_info.channels())\n        if freeze_stages > 0:\n            freeze_early_stages(self.backbone, freeze_stages)\n\n    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, ...]:\n        \"\"\"Extracts multi-scale features.\n\n        Args:\n            x: Images [N, 3, H, W]; resized internally to the native 256px.\n\n        Returns:\n            Tuple of four NCHW feature maps.\n        \"\"\"\n        if x.shape[-1] != self.INPUT_SIZE or x.shape[-2] != self.INPUT_SIZE:\n            x = F.interpolate(x, size=(self.INPUT_SIZE, self.INPUT_SIZE), mode=\"bilinear\", align_corners=False)\n        return tuple(to_nchw(f, c) for f, c in zip(self.backbone(x), self.channels))\n\n\nclass TemporalEncoderLayer(nn.Module):\n    \"\"\"Pre-norm transformer layer that keeps its last attention map.\"\"\"\n\n    def __init__(self, dim: int, num_heads: int, dropout: float = 0.1) -> None:\n        \"\"\"Initializes the layer.\n\n        Args:\n            dim: Token dimension.\n            num_heads: Number of attention heads.\n            dropout: Dropout probability.\n        \"\"\"\n        super().__init__()\n        self.norm1 = nn.LayerNorm(dim)\n        self.attn = nn.MultiheadAttention(dim, num_heads, dropout=dropout, batch_first=True)\n        self.norm2 = nn.LayerNorm(dim)\n        self.ff = nn.Sequential(nn.Linear(dim, dim * 2), nn.GELU(), nn.Dropout(dropout), nn.Linear(dim * 2, dim))\n        self.drop = nn.Dropout(dropout)\n        self.last_attn: Optional[torch.Tensor] = None\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        \"\"\"Applies self-attention and feed-forward blocks.\n\n        Args:\n            x: Tokens [B, L, D].\n\n        Returns:\n            Updated tokens [B, L, D].\n        \"\"\"\n        h = self.norm1(x)\n        out, weights = self.attn(h, h, h, need_weights=True, average_attn_weights=True)\n        self.last_attn = weights.detach()\n        x = x + self.drop(out)\n        return x + self.drop(self.ff(self.norm2(x)))\n\n\nclass BiologicalEncoder(nn.Module):\n    \"\"\"Encodes an rPPG trace using temporal convolutions and its magnitude spectrum.\"\"\"\n\n    def __init__(self, num_frames: int, out_dim: int) -> None:\n        \"\"\"Initializes the encoder.\n\n        Args:\n            num_frames: Length of the rPPG signal.\n            out_dim: Output feature size.\n        \"\"\"\n        super().__init__()\n        self.temporal = nn.Sequential(\n            nn.Conv1d(1, 16, 5, padding=2), nn.ReLU(inplace=True),\n            nn.Conv1d(16, 32, 5, padding=2), nn.ReLU(inplace=True),\n            nn.Conv1d(32, 64, 3, padding=1), nn.ReLU(inplace=True),\n        )\n        self.spectral = nn.Sequential(nn.Linear(num_frames // 2 + 1, 32), nn.ReLU(inplace=True))\n        self.proj = nn.Sequential(nn.Linear(64 * 2 + 32, out_dim), nn.GELU(), nn.LayerNorm(out_dim))\n\n    def forward(self, rppg: torch.Tensor) -> torch.Tensor:\n        \"\"\"Encodes rPPG signals.\n\n        Args:\n            rppg: Signals [B, T].\n\n        Returns:\n            Biological features [B, out_dim].\n        \"\"\"\n        rppg = torch.nan_to_num(rppg.float())\n        spectrum = torch.log1p(torch.fft.rfft(rppg, dim=-1).abs())\n        hidden = self.temporal(rppg.unsqueeze(1))\n        return self.proj(torch.cat([hidden.mean(-1), hidden.amax(-1), self.spectral(spectrum)], dim=-1))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-05T03:45:18.335534Z","iopub.execute_input":"2026-10-05T03:45:18.335911Z","iopub.status.idle":"2026-10-05T03:45:18.362446Z","shell.execute_reply.started":"2026-10-05T03:45:18.335891Z","shell.execute_reply":"2026-10-05T03:45:18.361738Z"}},"outputs":[],"execution_count":null},{"id":"4c37c904","cell_type":"code","source":"class ForensicXModel(nn.Module):\n    \"\"\"ForensicX: dual-backbone, frequency-aware, temporal and biological detector.\"\"\"\n\n    def __init__(\n        self,\n        num_frames: int = NUM_FRAMES,\n        dim: int = 128,\n        feat_dim: int = 256,\n        num_heads: int = 8,\n        temporal_layers: int = 2,\n        dropout: float = 0.1,\n        pretrained: bool = PRETRAINED,\n        freeze_stages: int = FREEZE_STAGES,\n        use_checkpoint: bool = False,\n    ) -> None:\n        \"\"\"Initializes all branches and heads.\n\n        Args:\n            num_frames: Frames per video.\n            dim: Token dimension used for cross-attention fusion.\n            feat_dim: Frame/video feature dimension.\n            num_heads: Attention heads.\n            temporal_layers: Number of temporal transformer layers.\n            dropout: Dropout probability.\n            pretrained: Whether to load pretrained backbones.\n            freeze_stages: Early backbone stages to freeze.\n            use_checkpoint: Enable gradient checkpointing.\n        \"\"\"\n        super().__init__()\n        self.num_frames = num_frames\n        self.use_checkpoint = use_checkpoint\n        self.convnext = ConvNeXtV2Branch(pretrained, freeze_stages)\n        self.swin = SwinTransformerV2Branch(pretrained, freeze_stages)\n        self.freq_encoder = FrequencyEncoder(IMAGE_SIZE, 64)\n\n        def projection(channels: int) -> nn.Sequential:\n            return nn.Sequential(nn.Conv2d(channels, dim, 1), nn.GroupNorm(1, dim))\n\n        self.conv_proj = nn.ModuleList([projection(c) for c in self.convnext.channels])\n        self.swin_proj = nn.ModuleList([projection(c) for c in self.swin.channels])\n        self.freq_proj = projection(64)\n        self.fusion = nn.ModuleList([StageFusion(dim, num_heads, dropout) for _ in self.convnext.channels])\n        self.frame_proj = nn.Sequential(\n            nn.Linear(dim * len(self.convnext.channels), feat_dim), nn.GELU(), nn.LayerNorm(feat_dim)\n        )\n\n        self.cls_token = nn.Parameter(torch.zeros(1, 1, feat_dim))\n        self.pos_embed = nn.Parameter(torch.zeros(1, num_frames + 1, feat_dim))\n        nn.init.trunc_normal_(self.pos_embed, std=0.02)\n        nn.init.trunc_normal_(self.cls_token, std=0.02)\n        self.temporal_encoder = nn.ModuleDict(\n            {\"layers\": nn.ModuleList([TemporalEncoderLayer(feat_dim, num_heads, dropout) for _ in range(temporal_layers)])}\n        )\n        self.temporal_norm = nn.LayerNorm(feat_dim)\n\n        self.bio_encoder = BiologicalEncoder(num_frames, feat_dim)\n        self.bio_gate = nn.Sequential(nn.Linear(2 * feat_dim, feat_dim), nn.Sigmoid())\n        self.vis_to_bio = nn.Linear(feat_dim, feat_dim)\n\n        self.classifier = nn.Sequential(\n            nn.LayerNorm(feat_dim), nn.Dropout(dropout), nn.Linear(feat_dim, feat_dim // 2),\n            nn.GELU(), nn.Dropout(dropout), nn.Linear(feat_dim // 2, 2),\n        )\n        self.evidence_head = nn.Sequential(nn.LayerNorm(feat_dim), nn.Linear(feat_dim, 2), nn.Softplus())\n        self.last_aux: Dict[str, torch.Tensor] = {}\n\n    def enable_gradient_checkpointing(self) -> None:\n        \"\"\"Recomputes backbone and fusion activations during backward to save memory.\"\"\"\n        self.use_checkpoint = True\n\n    @staticmethod\n    def _tokens(proj: nn.Module, feat: torch.Tensor) -> torch.Tensor:\n        \"\"\"Projects a feature map and pools it into a fixed token grid.\n\n        Args:\n            proj: Projection module.\n            feat: Feature map [N, C, H, W].\n\n        Returns:\n            Tokens [N, TOKEN_GRID**2, D].\n        \"\"\"\n        pooled = F.adaptive_avg_pool2d(proj(feat), TOKEN_GRID)\n        return rearrange(pooled, \"n c h w -> n (h w) c\")\n\n    def _run(self, module: nn.Module, *inputs: torch.Tensor) -> Any:\n        \"\"\"Runs a module, with non-reentrant checkpointing when enabled.\n\n        Args:\n            module: Module to execute.\n            *inputs: Positional tensor inputs.\n\n        Returns:\n            Module output.\n        \"\"\"\n        if self.use_checkpoint and self.training and torch.is_grad_enabled():\n            return checkpoint(module, *inputs, use_reentrant=False)\n        return module(*inputs)\n\n    def forward(\n        self, x: torch.Tensor, rppg_signal: Optional[torch.Tensor] = None\n    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n        \"\"\"Runs the full detector.\n\n        Args:\n            x: Video frames [B, T, 3, H, W].\n            rppg_signal: rPPG traces [B, T]; zeros are used when None.\n\n        Returns:\n            Tuple of logits [B, 2], evidence [B, 2] and temporal features [B, T, feat_dim].\n        \"\"\"\n        b, t = x.shape[:2]\n        frames = rearrange(x, \"b t c h w -> (b t) c h w\")\n        conv_feats = self._run(self.convnext, frames)\n        swin_feats = self._run(self.swin, frames)\n        freq_tokens = self._tokens(self.freq_proj, self.freq_encoder(frames))\n\n        stage_vectors: List[torch.Tensor] = []\n        previous: Optional[torch.Tensor] = None\n        for i, (cf, sf) in enumerate(zip(conv_feats, swin_feats)):\n            conv_tok = self._tokens(self.conv_proj[i], cf)\n            swin_tok = self._tokens(self.swin_proj[i], sf)\n            if i == 0:\n                conv_tok = conv_tok + freq_tokens\n            if previous is not None:\n                conv_tok, swin_tok = conv_tok + previous, swin_tok + previous\n            previous = self._run(self.fusion[i], conv_tok, swin_tok)\n            stage_vectors.append(previous.mean(dim=1))\n        del conv_feats, swin_feats, frames\n\n        frame_feat = self.frame_proj(torch.cat(stage_vectors, dim=-1)).view(b, t, -1)\n        tokens = torch.cat([self.cls_token.expand(b, -1, -1), frame_feat], dim=1)\n        tokens = tokens + self.pos_embed[:, : t + 1]\n        for layer in self.temporal_encoder[\"layers\"]:\n            tokens = layer(tokens)\n        tokens = self.temporal_norm(tokens)\n        video_feat, temporal_feat = tokens[:, 0], tokens[:, 1:]\n\n        if rppg_signal is None:\n            rppg_signal = torch.zeros(b, t, device=x.device)\n        bio_feat = self.bio_encoder(rppg_signal)\n        gate = self.bio_gate(torch.cat([video_feat, bio_feat.to(video_feat.dtype)], dim=-1))\n        fused = video_feat + gate * bio_feat.to(video_feat.dtype)\n\n        self.last_aux = {\n            \"embedding\": fused,\n            \"vis_aligned\": self.vis_to_bio(video_feat),\n            \"bio_feat\": bio_feat,\n        }\n        return self.classifier(fused), self.evidence_head(fused), temporal_feat\n\n\ndef count_parameters(model: nn.Module) -> Tuple[int, int]:\n    \"\"\"Counts total and trainable parameters.\n\n    Args:\n        model: Any torch module.\n\n    Returns:\n        Tuple of (total, trainable) parameter counts.\n    \"\"\"\n    total = sum(p.numel() for p in model.parameters())\n    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)\n    return total, trainable","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-05T03:45:18.363351Z","iopub.execute_input":"2026-10-05T03:45:18.36355Z","iopub.status.idle":"2026-10-05T03:45:18.385791Z","shell.execute_reply.started":"2026-10-05T03:45:18.363535Z","shell.execute_reply":"2026-10-05T03:45:18.385251Z"}},"outputs":[],"execution_count":null},{"id":"18118dec","cell_type":"markdown","source":"## Section 4: Loss Functions\nFocal loss handles class imbalance, supervised contrastive loss improves generalization, and temporal and biological terms enforce consistency on real videos. An evidential loss trains the uncertainty head, and `TotalLoss` combines all terms with learnable homoscedastic weights.","metadata":{}},{"id":"7d4e4891","cell_type":"code","source":"def dirichlet_stats(evidence: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"Converts evidence to Dirichlet expected probabilities and uncertainty.\n\n    Args:\n        evidence: Non-negative evidence [B, K].\n\n    Returns:\n        Tuple of expected probabilities [B, K] and uncertainty mass [B].\n    \"\"\"\n    alpha = evidence.float() + 1.0\n    strength = alpha.sum(dim=-1, keepdim=True)\n    return alpha / strength, (alpha.shape[-1] / strength).squeeze(-1)\n\n\nclass FocalLoss(nn.Module):\n    \"\"\"Multi-class focal loss for imbalanced real/fake data.\"\"\"\n\n    def __init__(self, gamma: float = 2.0, alpha: Optional[Sequence[float]] = None, label_smoothing: float = 0.0) -> None:\n        \"\"\"Initializes the loss.\n\n        Args:\n            gamma: Focusing parameter.\n            alpha: Optional per-class weights.\n            label_smoothing: Label smoothing factor.\n        \"\"\"\n        super().__init__()\n        self.gamma = gamma\n        self.label_smoothing = label_smoothing\n        self.register_buffer(\"alpha\", torch.tensor(alpha, dtype=torch.float32) if alpha is not None else None)\n\n    def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:\n        \"\"\"Computes the focal loss.\n\n        Args:\n            logits: Class logits [B, K].\n            labels: Class indices [B].\n\n        Returns:\n            Scalar loss.\n        \"\"\"\n        ce = F.cross_entropy(logits.float(), labels, reduction=\"none\", label_smoothing=self.label_smoothing)\n        loss = (1.0 - torch.exp(-ce)).pow(self.gamma) * ce\n        if self.alpha is not None:\n            loss = self.alpha[labels] * loss\n        return loss.mean()\n\n\nclass SupervisedContrastiveLoss(nn.Module):\n    def __init__(self, temperature: float = 0.07) -> None:\n        super().__init__()\n        self.temperature = temperature\n\n    def forward(self, embeddings: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:\n        z = F.normalize(embeddings.float(), dim=-1)\n        n = z.shape[0]\n        if n < 2:\n            return z.new_zeros(())\n\n        with torch.autocast(device_type=z.device.type, enabled=False):\n            sim = (z @ z.t()) / self.temperature\n            self_mask = torch.eye(n, dtype=torch.bool, device=z.device)\n            sim = sim.masked_fill(self_mask, -1e4)\n            log_prob = sim - torch.logsumexp(sim, dim=1, keepdim=True)\n            positives = (labels[:, None] == labels[None, :]) & ~self_mask\n            counts = positives.sum(dim=1)\n            valid = counts > 0\n            if not valid.any():\n                return z.new_zeros(())\n            return -((log_prob * positives).sum(dim=1)[valid] / counts[valid]).mean()\n\n\nclass TemporalConsistencyLoss(nn.Module):\n    \"\"\"Smooth frame features for real videos, encourage inconsistency for fakes.\"\"\"\n\n    def __init__(self, margin: float = 0.5) -> None:\n        \"\"\"Initializes the loss.\n\n        Args:\n            margin: Minimum desired frame-to-frame change for fake videos.\n        \"\"\"\n        super().__init__()\n        self.margin = margin\n\n    def forward(self, features: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:\n        \"\"\"Computes the loss.\n\n        Args:\n            features: Temporal features [B, T, D].\n            labels: Class indices [B] (0 = real, 1 = fake).\n\n        Returns:\n            Scalar loss.\n        \"\"\"\n        if features.shape[1] < 2:\n            return features.new_zeros(()).float()\n        f = F.normalize(features.float(), dim=-1)\n        diff = (f[:, 1:] - f[:, :-1]).pow(2).sum(-1).mean(dim=1)\n        real = (labels == 0).float()\n        return (real * diff + (1.0 - real) * F.relu(self.margin - diff)).mean()\n\n\nclass BiologicalConsistencyLoss(nn.Module):\n    \"\"\"Aligns visual and rPPG features for real videos and separates them for fakes.\"\"\"\n\n    def __init__(self, margin: float = 0.2) -> None:\n        \"\"\"Initializes the loss.\n\n        Args:\n            margin: Maximum tolerated similarity for fake videos.\n        \"\"\"\n        super().__init__()\n        self.margin = margin\n\n    def forward(self, visual: torch.Tensor, bio: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:\n        \"\"\"Computes the loss.\n\n        Args:\n            visual: Visual features aligned to the biological space [B, D].\n            bio: Biological features [B, D].\n            labels: Class indices [B].\n\n        Returns:\n            Scalar loss.\n        \"\"\"\n        sim = F.cosine_similarity(visual.float(), bio.float(), dim=-1)\n        real = (labels == 0).float()\n        return (real * (1.0 - sim) + (1.0 - real) * F.relu(sim - self.margin)).mean()\n\n\nclass EvidentialLoss(nn.Module):\n    \"\"\"Dirichlet-based evidential loss with annealed KL regularization.\"\"\"\n\n    def __init__(self, num_classes: int = 2, annealing_epochs: int = 10) -> None:\n        \"\"\"Initializes the loss.\n\n        Args:\n            num_classes: Number of classes.\n            annealing_epochs: Epochs over which the KL weight ramps to 1.\n        \"\"\"\n        super().__init__()\n        self.num_classes = num_classes\n        self.annealing_epochs = annealing_epochs\n        self.epoch = 0\n\n    def set_epoch(self, epoch: int) -> None:\n        \"\"\"Sets the current epoch for KL annealing.\n\n        Args:\n            epoch: Zero-based epoch index.\n        \"\"\"\n        self.epoch = epoch\n\n    @staticmethod\n    def _kl_to_uniform(alpha: torch.Tensor) -> torch.Tensor:\n        \"\"\"KL divergence between Dir(alpha) and the uniform Dirichlet.\n\n        Args:\n            alpha: Concentration parameters [B, K].\n\n        Returns:\n            KL values [B].\n        \"\"\"\n        ones = torch.ones_like(alpha)\n        total = alpha.sum(-1, keepdim=True)\n        log_norm = (\n            torch.lgamma(total) - torch.lgamma(alpha).sum(-1, keepdim=True)\n            + torch.lgamma(ones).sum(-1, keepdim=True) - torch.lgamma(ones.sum(-1, keepdim=True))\n        )\n        digamma_term = ((alpha - ones) * (torch.digamma(alpha) - torch.digamma(total))).sum(-1, keepdim=True)\n        return (log_norm + digamma_term).squeeze(-1)\n\n    def forward(self, evidence: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:\n        \"\"\"Computes the evidential loss.\n\n        Args:\n            evidence: Non-negative evidence [B, K].\n            labels: Class indices [B].\n\n        Returns:\n            Scalar loss.\n        \"\"\"\n        target = F.one_hot(labels, self.num_classes).float()\n        alpha = evidence.float() + 1.0\n        strength = alpha.sum(-1, keepdim=True)\n        prob = alpha / strength\n        error = (target - prob).pow(2).sum(-1)\n        variance = (prob * (1.0 - prob) / (strength + 1.0)).sum(-1)\n        misleading = target + (1.0 - target) * alpha\n        coef = min(1.0, self.epoch / max(1, self.annealing_epochs))\n        return (error + variance + coef * self._kl_to_uniform(misleading)).mean()\n\n\nclass TotalLoss(nn.Module):\n    \"\"\"Combines all objectives with learnable uncertainty-based weights.\"\"\"\n\n    NAMES: Tuple[str, ...] = (\"focal\", \"supcon\", \"temporal\", \"bio\", \"evidential\")\n\n    def __init__(self) -> None:\n        \"\"\"Initializes component losses and learnable log-variances.\"\"\"\n        super().__init__()\n        self.focal = FocalLoss()\n        self.supcon = SupervisedContrastiveLoss()\n        self.temporal = TemporalConsistencyLoss()\n        self.bio = BiologicalConsistencyLoss()\n        self.evidential = EvidentialLoss()\n        self.log_vars = nn.Parameter(torch.zeros(len(self.NAMES)))\n\n    def forward(\n        self,\n        logits: torch.Tensor,\n        evidence: torch.Tensor,\n        temporal_features: torch.Tensor,\n        aux: Dict[str, torch.Tensor],\n        labels: torch.Tensor,\n    ) -> Tuple[torch.Tensor, Dict[str, torch.Tensor]]:\n        \"\"\"Computes the weighted total loss.\n\n        Args:\n            logits: Class logits [B, 2].\n            evidence: Evidence [B, 2].\n            temporal_features: Temporal features [B, T, D].\n            aux: Auxiliary model outputs (`embedding`, `vis_aligned`, `bio_feat`).\n            labels: Class indices [B].\n\n        Returns:\n            Tuple of total loss and a dict of detached component losses.\n        \"\"\"\n        losses = torch.stack([\n            self.focal(logits, labels),\n            self.supcon(aux[\"embedding\"], labels),\n            self.temporal(temporal_features, labels),\n            self.bio(aux[\"vis_aligned\"], aux[\"bio_feat\"], labels),\n            self.evidential(evidence, labels),\n        ])\n        log_vars = self.log_vars.clamp(-4.0, 4.0)\n        total = (torch.exp(-log_vars) * losses + log_vars).sum()\n        return total, {n: l.detach() for n, l in zip(self.NAMES, losses)}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-05T03:45:18.388177Z","iopub.execute_input":"2026-10-05T03:45:18.388524Z","iopub.status.idle":"2026-10-05T03:45:18.416108Z","shell.execute_reply.started":"2026-10-05T03:45:18.388506Z","shell.execute_reply":"2026-10-05T03:45:18.415409Z"}},"outputs":[],"execution_count":null},{"id":"c1ef9781","cell_type":"markdown","source":"## Section 5: Training Pipeline (Kaggle-Optimized)\nThe `Trainer` uses mixed precision, gradient clipping and early stopping, and saves a checkpoint every epoch to `/kaggle/working/checkpoints/` so a session timeout loses nothing. Metrics are appended to `training_log.csv` and training resumes automatically from the latest checkpoint.","metadata":{}},{"id":"7ea357fb","cell_type":"code","source":"class OutOfMemoryAbort(RuntimeError):\n    \"\"\"Raised when repeated CUDA OOM errors require a smaller batch size.\"\"\"\n\n\ndef make_grad_scaler(enabled: bool) -> Any:\n    \"\"\"Creates a gradient scaler compatible with old and new torch versions.\n\n    Args:\n        enabled: Whether loss scaling is active.\n\n    Returns:\n        A GradScaler instance.\n    \"\"\"\n    try:\n        return torch.amp.GradScaler(\"cuda\", enabled=enabled)\n    except (AttributeError, TypeError):\n        return torch.cuda.amp.GradScaler(enabled=enabled)\n\n\ndef torch_load(path: Path, device: torch.device) -> Any:\n    \"\"\"Loads a checkpoint across torch versions.\n\n    Args:\n        path: Checkpoint path.\n        device: Target map location.\n\n    Returns:\n        The deserialized object.\n    \"\"\"\n    try:\n        return torch.load(path, map_location=device, weights_only=False)\n    except TypeError:\n        return torch.load(path, map_location=device)\n\n\ndef free_memory() -> None:\n    \"\"\"Runs garbage collection and clears the CUDA cache.\"\"\"\n    gc.collect()\n    if torch.cuda.is_available():\n        torch.cuda.empty_cache()\n\n\nclass Trainer:\n    \"\"\"Training loop with AMP, checkpointing, CSV logging and early stopping.\"\"\"\n\n    def __init__(\n        self,\n        model: nn.Module,\n        criterion: TotalLoss,\n        optimizer: torch.optim.Optimizer,\n        scheduler: Any,\n        device: torch.device = DEVICE,\n        use_amp: bool = USE_AMP,\n        checkpoint_dir: Path = KAGGLE_PATHS[\"checkpoints\"],\n        log_path: Path = KAGGLE_PATHS[\"log_csv\"],\n        patience: int = EARLY_STOP_PATIENCE,\n        grad_clip: float = 1.0,\n        keep_last: int = 2,\n        max_oom_skips: int = 3,\n        time_budget_hours: float = TIME_BUDGET_HOURS,\n    ) -> None:\n        \"\"\"Initializes the trainer.\n\n        Args:\n            model: ForensicX model.\n            criterion: Combined loss module.\n            optimizer: Optimizer covering model and loss parameters.\n            scheduler: Epoch-level learning-rate scheduler.\n            device: Compute device.\n            use_amp: Whether to use mixed precision.\n            checkpoint_dir: Directory for checkpoints.\n            log_path: CSV log path.\n            patience: Early-stopping patience in epochs.\n            grad_clip: Max gradient norm.\n            keep_last: Number of per-epoch checkpoints to keep (disk is limited).\n            max_oom_skips: OOM batches tolerated per epoch before aborting.\n            time_budget_hours: Wall-clock budget; training stops cleanly after it.\n        \"\"\"\n        self.model = model.to(device)\n        self.criterion = criterion.to(device)\n        self.optimizer = optimizer\n        self.scheduler = scheduler\n        self.device = device\n        self.use_amp = use_amp\n        self.checkpoint_dir = Path(checkpoint_dir)\n        self.log_path = Path(log_path)\n        self.patience = patience\n        self.grad_clip = grad_clip\n        self.keep_last = keep_last\n        self.max_oom_skips = max_oom_skips\n        self.time_budget = time_budget_hours * 3600\n        self.scaler = make_grad_scaler(use_amp)\n        self.history: List[Dict[str, float]] = []\n        self.best_score = -float(\"inf\")\n        self.epochs_without_improvement = 0\n        self.start_epoch = 0\n        self.session_start = time.time()\n        self.checkpoint_dir.mkdir(parents=True, exist_ok=True)\n\n    @property\n    def best_path(self) -> Path:\n        \"\"\"Path of the best-model weights file.\"\"\"\n        return self.checkpoint_dir / \"best_model.pt\"\n\n    def _autocast(self) -> Any:\n        \"\"\"Returns the autocast context for the current device.\"\"\"\n        return torch.autocast(device_type=self.device.type, dtype=torch.float16, enabled=self.use_amp)\n\n    def _params(self) -> List[torch.Tensor]:\n        \"\"\"Returns all optimized parameters.\"\"\"\n        return [p for group in self.optimizer.param_groups for p in group[\"params\"]]\n\n    def save_checkpoint(self, epoch: int, tag: Optional[str] = None) -> Path:\n        \"\"\"Atomically saves the full training state.\n\n        Args:\n            epoch: Epoch that just finished.\n            tag: Optional file stem (e.g. \"emergency\"); defaults to `epoch_XXX`.\n\n        Returns:\n            Path of the written checkpoint.\n        \"\"\"\n        path = self.checkpoint_dir / f\"{tag or f'epoch_{epoch:03d}'}.pt\"\n        state = {\n            \"epoch\": epoch,\n            \"model\": self.model.state_dict(),\n            \"criterion\": self.criterion.state_dict(),\n            \"optimizer\": self.optimizer.state_dict(),\n            \"scheduler\": self.scheduler.state_dict(),\n            \"scaler\": self.scaler.state_dict(),\n            \"history\": self.history,\n            \"best_score\": self.best_score,\n            \"epochs_without_improvement\": self.epochs_without_improvement,\n        }\n        tmp = path.with_suffix(\".tmp\")\n        torch.save(state, tmp)\n        os.replace(tmp, path)\n        if tag is None:\n            for old in sorted(self.checkpoint_dir.glob(\"epoch_*.pt\"))[: -self.keep_last]:\n                old.unlink(missing_ok=True)\n        return path\n\n    def load_latest_checkpoint(self) -> bool:\n        \"\"\"Restores the newest per-epoch checkpoint if one exists.\n\n        Returns:\n            True when a checkpoint was restored.\n        \"\"\"\n        candidates = sorted(self.checkpoint_dir.glob(\"epoch_*.pt\"))\n        if not candidates:\n            return False\n        try:\n            state = torch_load(candidates[-1], self.device)\n            self.model.load_state_dict(state[\"model\"])\n            self.criterion.load_state_dict(state[\"criterion\"])\n            self.optimizer.load_state_dict(state[\"optimizer\"])\n            self.scheduler.load_state_dict(state[\"scheduler\"])\n            self.scaler.load_state_dict(state[\"scaler\"])\n            self.history = list(state[\"history\"])\n            self.best_score = float(state[\"best_score\"])\n            self.epochs_without_improvement = int(state[\"epochs_without_improvement\"])\n            self.start_epoch = int(state[\"epoch\"]) + 1\n            print(f\"[INFO] Resumed from {candidates[-1].name} (next epoch {self.start_epoch + 1})\")\n            del state\n            free_memory()\n            return True\n        except Exception as exc:\n            print(f\"[WARN] Could not resume from {candidates[-1].name}: {exc}\")\n            return False\n\n    def _write_log(self) -> None:\n        \"\"\"Writes the training history to the CSV log.\"\"\"\n        pd.DataFrame(self.history).to_csv(self.log_path, index=False)\n\n    def train_epoch(self, loader: DataLoader, epoch: int) -> Dict[str, float]:\n        \"\"\"Runs one training epoch.\n\n        Args:\n            loader: Training DataLoader.\n            epoch: Zero-based epoch index.\n\n        Returns:\n            Dict with `train_loss` and `train_acc`.\n\n        Raises:\n            OutOfMemoryAbort: If too many batches hit CUDA OOM.\n        \"\"\"\n        self.model.train()\n        self.criterion.train()\n        self.criterion.evidential.set_epoch(epoch)\n        total_loss, correct, seen, oom_skips = 0.0, 0, 0, 0\n        bar = tqdm(loader, desc=f\"Train {epoch + 1}\", leave=False)\n        for frames, rppg, labels in bar:\n            frames = frames.to(self.device, non_blocking=True)\n            rppg = rppg.to(self.device, non_blocking=True)\n            labels = labels.to(self.device, non_blocking=True)\n            self.optimizer.zero_grad(set_to_none=True)\n            try:\n                with self._autocast():\n                    logits, evidence, temporal = self.model(frames, rppg)\n                    loss, _ = self.criterion(logits, evidence, temporal, self.model.last_aux, labels)\n                if not torch.isfinite(loss):\n                    continue\n                self.scaler.scale(loss).backward()\n                self.scaler.unscale_(self.optimizer)\n                nn.utils.clip_grad_norm_(self._params(), self.grad_clip)\n                self.scaler.step(self.optimizer)\n                self.scaler.update()\n            except torch.cuda.OutOfMemoryError:\n                oom_skips += 1\n                self.optimizer.zero_grad(set_to_none=True)\n                logits = evidence = temporal = loss = None\n                free_memory()\n                print(f\"[WARN] CUDA OOM at batch size {labels.shape[0]} ({oom_skips}/{self.max_oom_skips})\")\n                if oom_skips >= self.max_oom_skips:\n                    raise OutOfMemoryAbort(\"Repeated CUDA OOM; reduce the batch size.\")\n                continue\n            batch = labels.shape[0]\n            total_loss += loss.item() * batch\n            correct += (logits.argmax(dim=1) == labels).sum().item()\n            seen += batch\n            bar.set_postfix(loss=f\"{total_loss / seen:.4f}\", acc=f\"{correct / seen:.3f}\")\n        denominator = max(seen, 1)\n        return {\"train_loss\": total_loss / denominator, \"train_acc\": correct / denominator}\n\n    @torch.no_grad()\n    def validate_epoch(self, loader: DataLoader, epoch: int) -> Dict[str, float]:\n        \"\"\"Evaluates the model without gradients.\n\n        Args:\n            loader: Validation DataLoader.\n            epoch: Zero-based epoch index.\n\n        Returns:\n            Dict with `val_loss`, `val_acc`, `val_f1`, `val_auc` and `val_ece`.\n        \"\"\"\n        self.model.eval()\n        self.criterion.eval()\n        total_loss, seen = 0.0, 0\n        probs: List[np.ndarray] = []\n        labels_all: List[np.ndarray] = []\n        for frames, rppg, labels in tqdm(loader, desc=f\"Val {epoch + 1}\", leave=False):\n            frames = frames.to(self.device, non_blocking=True)\n            rppg = rppg.to(self.device, non_blocking=True)\n            labels = labels.to(self.device, non_blocking=True)\n            with self._autocast():\n                logits, evidence, temporal = self.model(frames, rppg)\n                loss, _ = self.criterion(logits, evidence, temporal, self.model.last_aux, labels)\n            total_loss += loss.item() * labels.shape[0]\n            seen += labels.shape[0]\n            probs.append(logits.float().softmax(dim=1)[:, 1].cpu().numpy())\n            labels_all.append(labels.cpu().numpy())\n        metrics = compute_metrics(np.concatenate(labels_all), np.concatenate(probs))\n        return {\n            \"val_loss\": total_loss / max(seen, 1),\n            \"val_acc\": metrics[\"accuracy\"],\n            \"val_f1\": metrics[\"f1\"],\n            \"val_auc\": metrics[\"auc\"],\n            \"val_ece\": metrics[\"ece\"],\n        }\n\n    def train(self, train_loader: DataLoader, val_loader: DataLoader, num_epochs: int = NUM_EPOCHS, resume: bool = True) -> pd.DataFrame:\n        \"\"\"Runs the full training loop.\n\n        Args:\n            train_loader: Training DataLoader.\n            val_loader: Validation DataLoader.\n            num_epochs: Total number of epochs.\n            resume: Whether to resume from the latest checkpoint.\n\n        Returns:\n            Training history as a DataFrame.\n        \"\"\"\n        if resume:\n            self.load_latest_checkpoint()\n        for epoch in range(self.start_epoch, num_epochs):\n            started = time.time()\n            train_metrics = self.train_epoch(train_loader, epoch)\n            val_metrics = self.validate_epoch(val_loader, epoch)\n            self.scheduler.step()\n\n            row = {\"epoch\": epoch + 1, **train_metrics, **val_metrics,\n                   \"lr\": self.optimizer.param_groups[0][\"lr\"], \"seconds\": time.time() - started}\n            self.history.append(row)\n            self._write_log()\n\n            score = val_metrics[\"val_auc\"] if np.isfinite(val_metrics[\"val_auc\"]) else -val_metrics[\"val_loss\"]\n            if score > self.best_score:\n                self.best_score, self.epochs_without_improvement = score, 0\n                torch.save(self.model.state_dict(), self.best_path)\n            else:\n                self.epochs_without_improvement += 1\n            self.save_checkpoint(epoch)\n\n            print(\" | \".join(f\"{k}={v:.4f}\" if isinstance(v, float) else f\"{k}={v}\" for k, v in row.items()))\n            print_gpu_memory(f\"epoch {epoch + 1}\")\n            free_memory()\n\n            if self.epochs_without_improvement >= self.patience:\n                print(f\"[INFO] Early stopping at epoch {epoch + 1}.\")\n                break\n            if time.time() - self.session_start > self.time_budget:\n                print(\"[INFO] Time budget reached; stopping to stay within the Kaggle session limit.\")\n                break\n        return pd.DataFrame(self.history)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-05T03:45:18.417291Z","iopub.execute_input":"2026-10-05T03:45:18.417559Z","iopub.status.idle":"2026-10-05T03:45:18.448933Z","shell.execute_reply.started":"2026-10-05T03:45:18.41754Z","shell.execute_reply":"2026-10-05T03:45:18.448415Z"}},"outputs":[],"execution_count":null},{"id":"1d15b30a","cell_type":"markdown","source":"## Section 6: Evaluation Metrics\nAccuracy, F1 and AUC come from scikit-learn, with a custom expected calibration error. Helpers collect predictions, aggregate clip scores per video and run cross-dataset evaluation.","metadata":{}},{"id":"32df7be5","cell_type":"code","source":"def expected_calibration_error(probs: np.ndarray, labels: np.ndarray, n_bins: int = 15) -> float:\n    \"\"\"Computes the expected calibration error for binary predictions.\n\n    Args:\n        probs: Probability of the fake class [N].\n        labels: Ground-truth labels [N].\n        n_bins: Number of confidence bins.\n\n    Returns:\n        ECE in [0, 1].\n    \"\"\"\n    probs, labels = np.asarray(probs, dtype=np.float64), np.asarray(labels)\n    confidence = np.maximum(probs, 1.0 - probs)\n    correct = ((probs >= 0.5).astype(int) == labels).astype(np.float64)\n    edges = np.linspace(0.5, 1.0, n_bins + 1)\n    ece = 0.0\n    for low, high in zip(edges[:-1], edges[1:]):\n        mask = (confidence > low) & (confidence <= high)\n        if mask.any():\n            ece += mask.mean() * abs(correct[mask].mean() - confidence[mask].mean())\n    return float(ece)\n\n\ndef compute_metrics(labels: np.ndarray, probs: np.ndarray, threshold: float = 0.5) -> Dict[str, float]:\n    \"\"\"Computes accuracy, F1, AUC and ECE.\n\n    Args:\n        labels: Ground-truth labels [N].\n        probs: Probability of the fake class [N].\n        threshold: Decision threshold.\n\n    Returns:\n        Dict with `accuracy`, `f1`, `auc` and `ece` (AUC is NaN for a single class).\n    \"\"\"\n    labels, probs = np.asarray(labels), np.asarray(probs)\n    preds = (probs >= threshold).astype(int)\n    try:\n        auc = float(roc_auc_score(labels, probs))\n    except ValueError:\n        auc = float(\"nan\")\n    return {\n        \"accuracy\": float(accuracy_score(labels, preds)),\n        \"f1\": float(f1_score(labels, preds, zero_division=0)),\n        \"auc\": auc,\n        \"ece\": expected_calibration_error(probs, labels),\n    }\n\n\n@torch.no_grad()\ndef collect_predictions(\n    model: nn.Module, loader: DataLoader, device: torch.device = DEVICE, use_amp: bool = USE_AMP, desc: str = \"Predict\"\n) -> Dict[str, np.ndarray]:\n    \"\"\"Runs the model over a loader.\n\n    Args:\n        model: Trained model.\n        loader: DataLoader (must not shuffle if order matters).\n        device: Compute device.\n        use_amp: Whether to use mixed precision.\n        desc: Progress-bar label.\n\n    Returns:\n        Dict with `probs` (fake probability), `labels` and evidential `uncertainty`.\n    \"\"\"\n    model.eval()\n    probs, labels_all, uncertainty = [], [], []\n    for frames, rppg, labels in tqdm(loader, desc=desc, leave=False):\n        with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=use_amp):\n            logits, evidence, _ = model(frames.to(device), rppg.to(device))\n        probs.append(logits.float().softmax(dim=1)[:, 1].cpu().numpy())\n        uncertainty.append(dirichlet_stats(evidence)[1].cpu().numpy())\n        labels_all.append(labels.numpy())\n    return {\"probs\": np.concatenate(probs), \"labels\": np.concatenate(labels_all), \"uncertainty\": np.concatenate(uncertainty)}\n\n\ndef video_level_prediction(\n    probs: np.ndarray, labels: np.ndarray, video_ids: Sequence[str], method: str = \"mean\"\n) -> Tuple[np.ndarray, np.ndarray]:\n    \"\"\"Aggregates clip-level scores into one score per video.\n\n    Args:\n        probs: Clip-level fake probabilities [N].\n        labels: Clip-level labels [N].\n        video_ids: Video identifier for each clip [N].\n        method: Aggregation: \"mean\", \"median\" or \"max\".\n\n    Returns:\n        Tuple of video-level probabilities and labels.\n    \"\"\"\n    if method not in (\"mean\", \"median\", \"max\"):\n        raise ValueError(f\"Unknown aggregation method: {method}\")\n    frame = pd.DataFrame({\"id\": list(video_ids), \"prob\": probs, \"label\": labels})\n    grouped = frame.groupby(\"id\", sort=False).agg(prob=(\"prob\", method), label=(\"label\", \"first\"))\n    return grouped[\"prob\"].to_numpy(), grouped[\"label\"].to_numpy()\n\n\ndef cross_dataset_evaluation(\n    model: nn.Module, loaders: Dict[str, DataLoader], device: torch.device = DEVICE, use_amp: bool = USE_AMP\n) -> pd.DataFrame:\n    \"\"\"Evaluates a model on several datasets at the video level.\n\n    Args:\n        model: Trained model.\n        loaders: Mapping from dataset name to a non-shuffled DataLoader.\n        device: Compute device.\n        use_amp: Whether to use mixed precision.\n\n    Returns:\n        DataFrame with one row of metrics per dataset.\n    \"\"\"\n    rows: List[Dict[str, Any]] = []\n    for name, loader in loaders.items():\n        try:\n            preds = collect_predictions(model, loader, device, use_amp, desc=f\"Eval {name}\")\n            ids = loader.dataset.video_paths\n            probs, labels = video_level_prediction(preds[\"probs\"], preds[\"labels\"], ids)\n            rows.append({\"dataset\": name, \"num_videos\": len(labels), **compute_metrics(labels, probs),\n                         \"mean_uncertainty\": float(preds[\"uncertainty\"].mean())})\n        except Exception as exc:\n            print(f\"[WARN] Evaluation failed for {name}: {exc}\")\n        free_memory()\n    return pd.DataFrame(rows)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-05T03:45:18.44961Z","iopub.execute_input":"2026-10-05T03:45:18.449854Z","iopub.status.idle":"2026-10-05T03:45:18.474061Z","shell.execute_reply.started":"2026-10-05T03:45:18.449837Z","shell.execute_reply":"2026-10-05T03:45:18.473375Z"}},"outputs":[],"execution_count":null},{"id":"712c8c94","cell_type":"markdown","source":"## Section 7: Explainability (XAI)\nGrad-CAM++ localizes the spatial evidence, Integrated Gradients gives pixel attributions via a Riemann sum, and attention rollout shows which frames drive the decision. Figures are saved to `/kaggle/working/xai_outputs/`.","metadata":{}},{"id":"84199532","cell_type":"code","source":"def get_gradcam_target_layer(model: ForensicXModel) -> nn.Module:\n    \"\"\"Returns the last ConvNeXt stage, used as the Grad-CAM++ target.\n\n    Args:\n        model: ForensicX model.\n\n    Returns:\n        Target module.\n    \"\"\"\n    return list(model.convnext.backbone.children())[-1]\n\n\nclass GradCAMPlusPlus:\n    \"\"\"Grad-CAM++ with forward-hook activation and gradient capture.\"\"\"\n\n    def __init__(self, model: ForensicXModel, target_layer: nn.Module) -> None:\n        \"\"\"Registers the hook on the target layer.\n\n        Args:\n            model: ForensicX model.\n            target_layer: Layer whose activations are explained.\n        \"\"\"\n        self.model = model\n        self.activations: Optional[torch.Tensor] = None\n        self.gradients: Optional[torch.Tensor] = None\n        self._handle = target_layer.register_forward_hook(self._forward_hook)\n\n    def _forward_hook(self, module: nn.Module, inputs: Any, output: torch.Tensor) -> None:\n        \"\"\"Stores activations and registers a gradient hook on the output.\"\"\"\n        self.activations = output.detach()\n        if output.requires_grad:\n            output.register_hook(lambda grad: setattr(self, \"gradients\", grad.detach()))\n\n    def remove(self) -> None:\n        \"\"\"Removes the registered hook.\"\"\"\n        self._handle.remove()\n\n    def __call__(self, x: torch.Tensor, rppg: torch.Tensor, class_idx: Optional[int] = None) -> Tuple[np.ndarray, int]:\n        \"\"\"Computes per-frame Grad-CAM++ maps for one video.\n\n        Args:\n            x: Video [1, T, 3, H, W].\n            rppg: rPPG trace [1, T].\n            class_idx: Target class; defaults to the predicted class.\n\n        Returns:\n            Tuple of maps [T, H, W] in [0, 1] and the explained class index.\n        \"\"\"\n        self.model.eval()\n        self.model.zero_grad(set_to_none=True)\n        with torch.enable_grad():\n            inputs = x.clone().requires_grad_(True)  # Guarantees a graph even with frozen stages.\n            logits, _, _ = self.model(inputs, rppg)\n            class_idx = int(logits.argmax(dim=1).item()) if class_idx is None else class_idx\n            logits[0, class_idx].backward()\n        acts, grads = self.activations.float(), self.gradients.float()\n        grads_sq, grads_cu = grads.pow(2), grads.pow(3)\n        denom = 2 * grads_sq + acts.sum(dim=(2, 3), keepdim=True) * grads_cu\n        alpha = grads_sq / torch.where(denom != 0, denom, torch.ones_like(denom))\n        weights = (alpha * F.relu(grads)).sum(dim=(2, 3), keepdim=True)\n        cam = F.relu((weights * acts).sum(dim=1, keepdim=True))\n        cam = F.interpolate(cam, size=x.shape[-2:], mode=\"bilinear\", align_corners=False).squeeze(1)\n        flat = cam.flatten(1)\n        low, high = flat.min(1)[0][:, None, None], flat.max(1)[0][:, None, None]\n        return ((cam - low) / (high - low + 1e-8)).cpu().numpy(), class_idx\n\n\nclass IntegratedGradients:\n    \"\"\"Integrated Gradients using a Riemann-sum approximation.\"\"\"\n\n    def __init__(self, model: ForensicXModel, steps: int = 16, method: str = \"midpoint\") -> None:\n        \"\"\"Initializes the explainer.\n\n        Args:\n            model: ForensicX model.\n            steps: Number of Riemann steps.\n            method: \"midpoint\", \"left\" or \"right\" Riemann sum.\n        \"\"\"\n        if method not in (\"midpoint\", \"left\", \"right\"):\n            raise ValueError(f\"Unknown method: {method}\")\n        self.model, self.steps, self.method = model, steps, method\n\n    def __call__(self, x: torch.Tensor, rppg: torch.Tensor, class_idx: int) -> np.ndarray:\n        \"\"\"Computes attributions against a zero baseline.\n\n        Args:\n            x: Video [1, T, 3, H, W].\n            rppg: rPPG trace [1, T].\n            class_idx: Class whose logit is explained.\n\n        Returns:\n            Attributions [T, 3, H, W].\n        \"\"\"\n        self.model.eval()\n        offsets = {\"midpoint\": 0.5, \"left\": 0.0, \"right\": 1.0}[self.method]\n        total = torch.zeros_like(x)\n        for step in range(self.steps):\n            scale = (step + offsets) / self.steps\n            point = (x * scale).detach().requires_grad_(True)\n            with torch.enable_grad():\n                logits, _, _ = self.model(point, rppg)\n                grad = torch.autograd.grad(logits[0, class_idx], point)[0]\n            total += grad\n            del point, grad, logits\n        return (x * total / self.steps).squeeze(0).detach().cpu().numpy()\n\n\nclass TemporalAttribution:\n    \"\"\"Frame importance from attention rollout over the temporal transformer.\"\"\"\n\n    def __init__(self, model: ForensicXModel) -> None:\n        \"\"\"Initializes the explainer.\n\n        Args:\n            model: ForensicX model.\n        \"\"\"\n        self.model = model\n\n    @torch.no_grad()\n    def __call__(self, x: torch.Tensor, rppg: torch.Tensor) -> np.ndarray:\n        \"\"\"Computes CLS-to-frame attention rollout.\n\n        Args:\n            x: Video [1, T, 3, H, W].\n            rppg: rPPG trace [1, T].\n\n        Returns:\n            Frame importance [T] summing to 1.\n        \"\"\"\n        self.model.eval()\n        self.model(x, rppg)\n        rollout: Optional[torch.Tensor] = None\n        for layer in self.model.temporal_encoder[\"layers\"]:\n            attn = layer.last_attn.float()\n            attn = 0.5 * attn + 0.5 * torch.eye(attn.shape[-1], device=attn.device)\n            attn = attn / attn.sum(dim=-1, keepdim=True)\n            rollout = attn if rollout is None else attn @ rollout\n        importance = rollout[0, 0, 1:]\n        return (importance / importance.sum()).cpu().numpy()\n\n\ndef denormalize(frames: torch.Tensor) -> np.ndarray:\n    \"\"\"Converts normalized frames back to uint8 RGB.\n\n    Args:\n        frames: Tensor [T, 3, H, W].\n\n    Returns:\n        Array [T, H, W, 3] of dtype uint8.\n    \"\"\"\n    mean = torch.tensor(IMAGENET_MEAN, device=frames.device).view(1, 3, 1, 1)\n    std = torch.tensor(IMAGENET_STD, device=frames.device).view(1, 3, 1, 1)\n    image = (frames * std + mean).clamp(0, 1).permute(0, 2, 3, 1).cpu().numpy()\n    return (image * 255).astype(np.uint8)\n\n\ndef visualize_explanations(\n    frames: np.ndarray,\n    cams: np.ndarray,\n    ig_maps: np.ndarray,\n    temporal_attn: np.ndarray,\n    title: str,\n    save_path: Path,\n    max_frames: int = 6,\n) -> Path:\n    \"\"\"Creates and saves a combined explanation figure.\n\n    Args:\n        frames: RGB frames [T, H, W, 3].\n        cams: Grad-CAM++ maps [T, H, W].\n        ig_maps: Integrated Gradients attributions [T, 3, H, W].\n        temporal_attn: Frame importance [T].\n        title: Figure title.\n        save_path: Output PNG path.\n        max_frames: Maximum number of frames displayed.\n\n    Returns:\n        The saved path.\n    \"\"\"\n    idx = np.unique(np.linspace(0, len(frames) - 1, min(max_frames, len(frames))).round().astype(int))\n    ig = np.abs(ig_maps).sum(axis=1)\n    ig = np.clip(ig / (np.percentile(ig, 99) + 1e-8), 0, 1)\n\n    fig = plt.figure(figsize=(3 * len(idx), 11))\n    grid = fig.add_gridspec(4, len(idx), height_ratios=[1, 1, 1, 0.7])\n    for col, i in enumerate(idx):\n        heat = cv2.applyColorMap((cams[i] * 255).astype(np.uint8), cv2.COLORMAP_JET)[..., ::-1]\n        overlay = (0.55 * frames[i] + 0.45 * heat).astype(np.uint8)\n        for row, (image, kwargs, label) in enumerate([\n            (frames[i], {}, \"Input\"), (overlay, {}, \"Grad-CAM++\"), (ig[i], {\"cmap\": \"inferno\"}, \"Integrated Gradients\"),\n        ]):\n            ax = fig.add_subplot(grid[row, col])\n            ax.imshow(image, **kwargs)\n            ax.axis(\"off\")\n            ax.set_title(f\"{label} (t={i})\" if col == 0 else f\"t={i}\", fontsize=9)\n    ax = fig.add_subplot(grid[3, :])\n    ax.bar(np.arange(len(temporal_attn)), temporal_attn, color=\"tab:blue\")\n    ax.set(xlabel=\"Frame\", ylabel=\"Attention rollout\", title=\"Temporal attribution\")\n    fig.suptitle(title, fontsize=13)\n    fig.tight_layout()\n    save_path.parent.mkdir(parents=True, exist_ok=True)\n    fig.savefig(save_path, dpi=120, bbox_inches=\"tight\")\n    plt.close(fig)\n    return save_path","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-05T03:45:18.474957Z","iopub.execute_input":"2026-10-05T03:45:18.475237Z","iopub.status.idle":"2026-10-05T03:45:18.499988Z","shell.execute_reply.started":"2026-10-05T03:45:18.475214Z","shell.execute_reply":"2026-10-05T03:45:18.499341Z"}},"outputs":[],"execution_count":null},{"id":"8bc2a463","cell_type":"markdown","source":"## Section 8: Uncertainty Quantification\nMonte Carlo dropout estimates epistemic uncertainty at inference, the evidential head yields a Dirichlet-based uncertainty in one pass, and split conformal prediction provides prediction sets with a coverage guarantee.","metadata":{}},{"id":"2f18a5d5","cell_type":"code","source":"def enable_dropout(model: nn.Module) -> None:\n    \"\"\"Switches dropout and attention-dropout layers to train mode, leaving BatchNorm in eval.\n\n    Args:\n        model: Model in eval mode.\n    \"\"\"\n    for module in model.modules():\n        if isinstance(module, (nn.Dropout, nn.MultiheadAttention, HierarchicalCrossAttention)):\n            module.train()\n\n\n@torch.no_grad()\ndef mc_dropout_predict(\n    model: nn.Module, x: torch.Tensor, rppg: torch.Tensor, n_samples: int = 20\n) -> Dict[str, torch.Tensor]:\n    \"\"\"Runs stochastic forward passes with dropout enabled.\n\n    Args:\n        model: ForensicX model.\n        x: Videos [B, T, 3, H, W].\n        rppg: rPPG traces [B, T].\n        n_samples: Number of stochastic passes.\n\n    Returns:\n        Dict with `mean_probs` [B, 2], `std_probs` [B, 2] and predictive `entropy` [B].\n    \"\"\"\n    model.eval()\n    enable_dropout(model)\n    try:\n        samples = torch.stack([model(x, rppg)[0].float().softmax(dim=-1) for _ in range(n_samples)])\n    finally:\n        model.eval()\n    mean = samples.mean(dim=0)\n    entropy = -(mean * torch.log(mean + 1e-8)).sum(dim=-1)\n    return {\"mean_probs\": mean, \"std_probs\": samples.std(dim=0), \"entropy\": entropy}\n\n\nclass EvidentialHead(nn.Module):\n    \"\"\"Standalone evidential head mapping features to Dirichlet evidence.\"\"\"\n\n    def __init__(self, in_features: int, num_classes: int = 2) -> None:\n        \"\"\"Initializes the head.\n\n        Args:\n            in_features: Input feature size.\n            num_classes: Number of classes.\n        \"\"\"\n        super().__init__()\n        self.net = nn.Sequential(nn.LayerNorm(in_features), nn.Linear(in_features, num_classes), nn.Softplus())\n\n    def forward(self, features: torch.Tensor) -> torch.Tensor:\n        \"\"\"Computes non-negative evidence.\n\n        Args:\n            features: Features [B, D].\n\n        Returns:\n            Evidence [B, K].\n        \"\"\"\n        return self.net(features)\n\n    @staticmethod\n    def predict(evidence: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:\n        \"\"\"Returns expected probabilities and uncertainty mass.\n\n        Args:\n            evidence: Evidence [B, K].\n\n        Returns:\n            Tuple of probabilities [B, K] and uncertainty [B].\n        \"\"\"\n        return dirichlet_stats(evidence)\n\n\nclass ConformalPrediction:\n    \"\"\"Split conformal prediction for classification.\"\"\"\n\n    def __init__(self, alpha: float = 0.1) -> None:\n        \"\"\"Initializes the predictor.\n\n        Args:\n            alpha: Target miscoverage rate (coverage = 1 - alpha).\n        \"\"\"\n        self.alpha = alpha\n        self.threshold: Optional[float] = None\n\n    @staticmethod\n    def _as_matrix(probs: np.ndarray) -> np.ndarray:\n        \"\"\"Converts fake-class probabilities [N] to class probabilities [N, 2].\"\"\"\n        probs = np.asarray(probs, dtype=np.float64)\n        return np.stack([1.0 - probs, probs], axis=1) if probs.ndim == 1 else probs\n\n    def calibrate(self, probs: np.ndarray, labels: np.ndarray) -> float:\n        \"\"\"Computes the conformal quantile on a calibration set.\n\n        Args:\n            probs: Fake probabilities [N] or class probabilities [N, K].\n            labels: Ground-truth labels [N].\n\n        Returns:\n            The nonconformity threshold.\n        \"\"\"\n        matrix = self._as_matrix(probs)\n        scores = 1.0 - matrix[np.arange(len(labels)), np.asarray(labels, dtype=int)]\n        n = len(scores)\n        level = min(1.0, math.ceil((n + 1) * (1.0 - self.alpha)) / n)\n        self.threshold = float(np.quantile(scores, level, method=\"higher\"))\n        return self.threshold\n\n    def predict_sets(self, probs: np.ndarray) -> np.ndarray:\n        \"\"\"Builds prediction sets.\n\n        Args:\n            probs: Fake probabilities [N] or class probabilities [N, K].\n\n        Returns:\n            Boolean membership matrix [N, K].\n        \"\"\"\n        if self.threshold is None:\n            raise RuntimeError(\"Call calibrate() first.\")\n        return (1.0 - self._as_matrix(probs)) <= self.threshold\n\n    def evaluate(self, probs: np.ndarray, labels: np.ndarray) -> Dict[str, float]:\n        \"\"\"Measures empirical coverage and mean set size.\n\n        Args:\n            probs: Fake probabilities [N] or class probabilities [N, K].\n            labels: Ground-truth labels [N].\n\n        Returns:\n            Dict with `coverage` and `avg_set_size`.\n        \"\"\"\n        sets = self.predict_sets(probs)\n        covered = sets[np.arange(len(labels)), np.asarray(labels, dtype=int)]\n        return {\"coverage\": float(covered.mean()), \"avg_set_size\": float(sets.sum(axis=1).mean())}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-05T03:45:18.500965Z","iopub.execute_input":"2026-10-05T03:45:18.501232Z","iopub.status.idle":"2026-10-05T03:45:18.515484Z","shell.execute_reply.started":"2026-10-05T03:45:18.501175Z","shell.execute_reply":"2026-10-05T03:45:18.514796Z"}},"outputs":[],"execution_count":null},{"id":"8d930b8f","cell_type":"markdown","source":"## Section 9: Data Loading and Preprocessing\nLoaders handle the FF++, Celeb-DF and DFDC folder layouts (with a generic fallback) and return a DataFrame of `path, label, dataset, group`. Splits are group-aware to avoid identity leakage, and DataLoaders use `num_workers=2` with an automatic fallback to 0 if multiprocessing fails.","metadata":{}},{"id":"76003386","cell_type":"code","source":"FF_FAKE_TOKENS: Tuple[str, ...] = (\n    \"manipulated\", \"deepfakes\", \"face2face\", \"faceswap\", \"neuraltextures\", \"faceshifter\", \"deepfakedetection\", \"fake\",\n)\n\n\ndef list_videos(root: Path) -> List[Path]:\n    \"\"\"Recursively lists video files.\n\n    Args:\n        root: Directory to scan.\n\n    Returns:\n        Sorted list of video paths.\n    \"\"\"\n    videos = [Path(c) / f for c, _, files in os.walk(root) for f in files if f.lower().endswith(VIDEO_EXTS)]\n    return sorted(videos)\n\n\ndef infer_label(path: Path, root: Path) -> Optional[int]:\n    \"\"\"Infers the label (0 real, 1 fake) from folder names.\n\n    Args:\n        path: Video path.\n        root: Dataset root, stripped so its own name does not bias the result.\n\n    Returns:\n        0, 1 or None when the path is ambiguous.\n    \"\"\"\n    try:\n        text = \"/\".join(path.relative_to(root).parts[:-1]).lower()\n    except ValueError:\n        text = str(path.parent).lower()\n    if any(token in text for token in FF_FAKE_TOKENS) or \"synthesis\" in text:\n        return 1\n    if any(token in text for token in (\"original\", \"real\", \"youtube\", \"actors\")):\n        return 0\n    return None\n\n\ndef _frame(paths: List[Path], labels: List[int], groups: List[str], name: str) -> pd.DataFrame:\n    \"\"\"Builds the standard dataset DataFrame.\"\"\"\n    return pd.DataFrame({\"path\": [str(p) for p in paths], \"label\": labels, \"dataset\": name, \"group\": groups})\n\n\ndef load_generic(root: Path, name: str) -> pd.DataFrame:\n    \"\"\"Loads any dataset organised in real/fake-named folders.\n\n    Args:\n        root: Dataset root.\n        name: Dataset key stored in the DataFrame.\n\n    Returns:\n        DataFrame with columns path, label, dataset, group.\n    \"\"\"\n    paths, labels = [], []\n    for path in list_videos(root):\n        label = infer_label(path, root)\n        if label is not None:\n            paths.append(path)\n            labels.append(label)\n    return _frame(paths, labels, [p.stem.split(\"_\")[0] for p in paths], name)\n\n\ndef load_ffpp(root: Path) -> pd.DataFrame:\n    \"\"\"Loads FaceForensics++ (original_sequences / manipulated_sequences).\n\n    Args:\n        root: Dataset root.\n\n    Returns:\n        DataFrame grouped by the source/target identity id in the file name.\n    \"\"\"\n    return load_generic(root, \"ffpp\")\n\n\ndef load_celeb_df(root: Path) -> pd.DataFrame:\n    \"\"\"Loads Celeb-DF (Celeb-real, YouTube-real, Celeb-synthesis).\n\n    Args:\n        root: Dataset root.\n\n    Returns:\n        DataFrame restricted to the official test list when it is present.\n    \"\"\"\n    df = load_generic(root, \"celeb_df\")\n    test_lists = [Path(c) / f for c, _, fs in os.walk(root) for f in fs if f == \"List_of_testing_videos.txt\"]\n    if test_lists and len(df):\n        official = {line.split()[-1].strip() for line in test_lists[0].read_text().splitlines() if line.strip()}\n        keys = df[\"path\"].map(lambda p: \"/\".join(Path(p).parts[-2:]))\n        restricted = df[keys.isin(official)]\n        if restricted[\"label\"].nunique() == 2:\n            df = restricted\n    return df.reset_index(drop=True)\n\n\ndef load_dfdc(root: Path) -> pd.DataFrame:\n    \"\"\"Loads DFDC using metadata.json files (REAL/FAKE labels).\n\n    Args:\n        root: Dataset root.\n\n    Returns:\n        DataFrame where fakes share a group with their original video.\n    \"\"\"\n    paths, labels, groups = [], [], []\n    for current, _, files in os.walk(root):\n        if \"metadata.json\" not in files:\n            continue\n        try:\n            metadata = json.loads((Path(current) / \"metadata.json\").read_text())\n        except (OSError, json.JSONDecodeError) as exc:\n            print(f\"[WARN] Bad DFDC metadata in {current}: {exc}\")\n            continue\n        for file_name, info in metadata.items():\n            video = Path(current) / file_name\n            if video.exists():\n                paths.append(video)\n                labels.append(int(str(info.get(\"label\", \"\")).upper() == \"FAKE\"))\n                groups.append(str(info.get(\"original\") or file_name))\n    if not paths:\n        return load_generic(root, \"dfdc\")\n    return _frame(paths, labels, groups, \"dfdc\")\n\n\ndef balanced_subsample(df: pd.DataFrame, max_per_class: Optional[int], seed: int = SEED) -> pd.DataFrame:\n    \"\"\"Limits the number of videos per class.\n\n    Args:\n        df: Dataset DataFrame.\n        max_per_class: Cap per label; None keeps everything.\n        seed: Sampling seed.\n\n    Returns:\n        Subsampled DataFrame.\n    \"\"\"\n    if max_per_class is None or df.empty:\n        return df\n    parts = [g.sample(n=min(len(g), max_per_class), random_state=seed) for _, g in df.groupby(\"label\")]\n    return pd.concat(parts).reset_index(drop=True)\n\n\ndef create_synthetic_dataset(root: Path, n_videos: int = 48, num_frames: int = 20, size: int = 160) -> pd.DataFrame:\n    \"\"\"Writes tiny synthetic videos so the pipeline can be smoke-tested without data.\n\n    Args:\n        root: Output directory.\n        n_videos: Number of videos (half real, half fake).\n        num_frames: Frames per video.\n        size: Frame side length in pixels.\n\n    Returns:\n        DataFrame describing the synthetic videos.\n    \"\"\"\n    root.mkdir(parents=True, exist_ok=True)\n    rng = np.random.default_rng(SEED)\n    paths, labels = [], []\n    yy, xx = np.mgrid[0:size, 0:size]\n    for i in range(n_videos):\n        label = i % 2\n        path = root / f\"{'fake' if label else 'real'}_{i:03d}.avi\"\n        if not path.exists():\n            writer = cv2.VideoWriter(str(path), cv2.VideoWriter_fourcc(*\"MJPG\"), 15, (size, size))\n            base = rng.integers(40, 120, size=3)\n            for t in range(num_frames):\n                frame = np.empty((size, size, 3), dtype=np.uint8)\n                frame[:] = base\n                cx, cy = size // 2 + int(10 * math.sin(t / 3)), size // 2\n                mask = ((xx - cx) ** 2 / 40 ** 2 + (yy - cy) ** 2 / 55 ** 2) < 1\n                frame[mask] = (90, 140, 200 + int(8 * math.sin(t / 2)))\n                if label:\n                    frame = np.clip(frame.astype(np.int16) + rng.normal(0, 12, frame.shape) + ((xx + yy) % 4 < 2)[..., None] * 10, 0, 255).astype(np.uint8)\n                writer.write(frame)\n            writer.release()\n        paths.append(path)\n        labels.append(label)\n    return _frame(paths, labels, [p.stem for p in paths], \"synthetic\")\n\n\ndef load_all_datasets() -> Dict[str, pd.DataFrame]:\n    \"\"\"Loads every detected dataset, falling back to synthetic data.\n\n    Returns:\n        Mapping from dataset key to its (subsampled) DataFrame.\n    \"\"\"\n    loaders = {\"ffpp\": load_ffpp, \"celeb_df\": load_celeb_df, \"dfdc\": load_dfdc}\n    datasets: Dict[str, pd.DataFrame] = {}\n    for key, root in DATASET_PATHS.items():\n        try:\n            df = loaders[key](root)\n        except Exception as exc:\n            print(f\"[WARN] Could not load {key} from {root}: {exc}\")\n            continue\n        if df[\"label\"].nunique() == 2:\n            datasets[key] = balanced_subsample(df, MAX_VIDEOS_PER_CLASS)\n            print(f\"[INFO] {key}: {len(datasets[key])} videos ({dict(datasets[key]['label'].value_counts())})\")\n        else:\n            print(f\"[WARN] {key}: needs both classes, found {dict(df['label'].value_counts())}; skipped.\")\n    if not datasets:\n        print(\"[WARN] No usable dataset found; generating a synthetic smoke-test dataset.\")\n        datasets[\"synthetic\"] = create_synthetic_dataset(KAGGLE_PATHS[\"synthetic\"])\n    return datasets\n\n\ndef create_splits(\n    df: pd.DataFrame, val_size: float = 0.15, test_size: float = 0.15, seed: int = SEED\n) -> Tuple[pd.DataFrame, pd.DataFrame, pd.DataFrame]:\n    \"\"\"Creates group-aware train/val/test splits containing both classes.\n\n    Args:\n        df: Dataset DataFrame with a `group` column.\n        val_size: Validation fraction of the full data.\n        test_size: Test fraction of the full data.\n        seed: Base random seed.\n\n    Returns:\n        Tuple of train, validation and test DataFrames.\n    \"\"\"\n    for attempt in range(25):\n        outer = GroupShuffleSplit(n_splits=1, test_size=test_size, random_state=seed + attempt)\n        trainval_idx, test_idx = next(outer.split(df, groups=df[\"group\"]))\n        trainval = df.iloc[trainval_idx]\n        inner = GroupShuffleSplit(n_splits=1, test_size=val_size / (1 - test_size), random_state=seed + attempt)\n        train_idx, val_idx = next(inner.split(trainval, groups=trainval[\"group\"]))\n        train, val, test = trainval.iloc[train_idx], trainval.iloc[val_idx], df.iloc[test_idx]\n        if all(part[\"label\"].nunique() == 2 for part in (train, val, test)):\n            break\n    return train.reset_index(drop=True), val.reset_index(drop=True), test.reset_index(drop=True)\n\n\ndef prepare_data() -> Tuple[pd.DataFrame, pd.DataFrame, Dict[str, pd.DataFrame]]:\n    \"\"\"Builds train/val splits from the primary dataset and cross-dataset test sets.\n\n    Returns:\n        Tuple of train DataFrame, validation DataFrame and a dict of test DataFrames.\n    \"\"\"\n    datasets = load_all_datasets()\n    primary = \"ffpp\" if \"ffpp\" in datasets else next(iter(datasets))\n    train, val, test = create_splits(datasets[primary])\n    test_sets = {f\"{primary}_test\": test}\n    test_sets.update({name: df for name, df in datasets.items() if name != primary})\n    print(f\"[INFO] Primary={primary} | train={len(train)} val={len(val)} | tests={ {k: len(v) for k, v in test_sets.items()} }\")\n    return train, val, test_sets\n\n\ndef multiprocessing_available() -> bool:\n    \"\"\"Checks whether worker processes can be spawned in this environment.\n\n    Returns:\n        True if a one-process pool works.\n    \"\"\"\n    try:\n        with mp.get_context().Pool(1) as pool:\n            return pool.map(abs, [-1]) == [1]\n    except Exception:\n        return False\n\n\ndef make_loader(dataset: Dataset, batch_size: int, shuffle: bool, num_workers: int) -> DataLoader:\n    \"\"\"Creates a DataLoader with Kaggle-friendly settings.\n\n    Args:\n        dataset: Dataset instance.\n        batch_size: Batch size.\n        shuffle: Whether to shuffle (training only).\n        num_workers: Worker processes; 0 disables multiprocessing.\n\n    Returns:\n        Configured DataLoader.\n    \"\"\"\n    kwargs: Dict[str, Any] = {}\n    if num_workers > 0:\n        kwargs.update(persistent_workers=True, prefetch_factor=2, worker_init_fn=worker_init_fn)\n    return DataLoader(\n        dataset, batch_size=batch_size, shuffle=shuffle, num_workers=num_workers,\n        pin_memory=torch.cuda.is_available(), drop_last=shuffle and len(dataset) > batch_size, **kwargs,\n    )\n\n\ndef build_loaders(\n    train_df: pd.DataFrame, val_df: pd.DataFrame, test_sets: Dict[str, pd.DataFrame], batch_size: int = BATCH_SIZE\n) -> Tuple[DataLoader, DataLoader, Dict[str, DataLoader]]:\n    \"\"\"Builds train, validation and test DataLoaders.\n\n    Args:\n        train_df: Training DataFrame.\n        val_df: Validation DataFrame.\n        test_sets: Mapping from name to test DataFrame.\n        batch_size: Training batch size (evaluation uses the same size).\n\n    Returns:\n        Tuple of train loader, validation loader and test loaders.\n    \"\"\"\n    workers = NUM_WORKERS if multiprocessing_available() else 0\n    detector = FaceDetector(\"haar\")\n    print(f\"[INFO] workers={workers} | face detector={detector.backend} | batch_size={batch_size}\")\n\n    def dataset(df: pd.DataFrame, train: bool, quality: str) -> ForensicXDataset:\n        return ForensicXDataset(df[\"path\"].tolist(), df[\"label\"].tolist(), build_transforms(train), NUM_FRAMES, quality, detector)\n\n    train_loader = make_loader(dataset(train_df, True, \"MIXED\"), batch_size, True, workers)\n    val_loader = make_loader(dataset(val_df, False, \"HQ\"), batch_size, False, workers)\n    test_loaders = {name: make_loader(dataset(df, False, \"HQ\"), batch_size, False, workers) for name, df in test_sets.items()}\n    return train_loader, val_loader, test_loaders","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-05T03:45:18.516437Z","iopub.execute_input":"2026-10-05T03:45:18.516697Z","iopub.status.idle":"2026-10-05T03:45:18.55468Z","shell.execute_reply.started":"2026-10-05T03:45:18.516668Z","shell.execute_reply":"2026-10-05T03:45:18.554041Z"}},"outputs":[],"execution_count":null},{"id":"0ea420b3","cell_type":"markdown","source":"## Section 10: Main Execution\nRuns the full workflow: data preparation, model and optimizer setup, training with per-epoch checkpoints and an OOM batch-size fallback, cross-dataset evaluation, conformal calibration and XAI figures. On any error, an emergency checkpoint is written to `/kaggle/working/checkpoints/` before re-raising.","metadata":{}},{"id":"70d76aa9","cell_type":"code","source":"def run_xai_examples(model: ForensicXModel, df: pd.DataFrame, n_per_class: int = 1) -> List[Path]:\n    \"\"\"Generates XAI figures for a few test videos.\n\n    Args:\n        model: Trained model.\n        df: Test DataFrame.\n        n_per_class: Number of videos per class to explain.\n\n    Returns:\n        Paths of the saved figures.\n    \"\"\"\n    dataset = ForensicXDataset(df[\"path\"].tolist(), df[\"label\"].tolist(), build_transforms(False), NUM_FRAMES, \"HQ\", FaceDetector(\"haar\"))\n    cam_explainer = GradCAMPlusPlus(model, get_gradcam_target_layer(model))\n    saved: List[Path] = []\n    try:\n        for label in (0, 1):\n            for index in df.index[df[\"label\"] == label][:n_per_class]:\n                frames, rppg, _ = dataset[int(index)]\n                x, r = frames.unsqueeze(0).to(DEVICE), rppg.unsqueeze(0).to(DEVICE)\n                cams, cls = cam_explainer(x, r)\n                ig = IntegratedGradients(model)(x, r, cls)\n                temporal = TemporalAttribution(model)(x, r)\n                title = f\"{'Real' if label == 0 else 'Fake'} video | predicted: {'Fake' if cls == 1 else 'Real'}\"\n                path = KAGGLE_PATHS[\"xai\"] / f\"xai_{'real' if label == 0 else 'fake'}_{int(index):04d}.png\"\n                saved.append(visualize_explanations(denormalize(x[0]), cams, ig, temporal, title, path))\n                free_memory()\n    finally:\n        cam_explainer.remove()\n    return saved\n\n\ndef main() -> None:\n    \"\"\"Runs the complete ForensicX training and evaluation pipeline.\"\"\"\n    trainer: Optional[Trainer] = None\n    try:\n        set_seed()\n        create_output_dirs()\n        train_df, val_df, test_sets = prepare_data()\n\n        model = ForensicXModel(num_frames=NUM_FRAMES, use_checkpoint=USE_GRAD_CHECKPOINT).to(DEVICE)\n        total, trainable = count_parameters(model)\n        print(f\"[INFO] Parameters: total={total / 1e6:.1f}M trainable={trainable / 1e6:.1f}M\")\n        print_gpu_memory(\"model loaded\")\n\n        criterion = TotalLoss().to(DEVICE)\n        params = [p for p in model.parameters() if p.requires_grad] + list(criterion.parameters())\n        optimizer = torch.optim.AdamW(params, lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY)\n        scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=NUM_EPOCHS)\n        trainer = Trainer(model, criterion, optimizer, scheduler)\n\n        batch_size = BATCH_SIZE\n        for attempt in range(4):\n            train_loader, val_loader, test_loaders = build_loaders(train_df, val_df, test_sets, batch_size)\n            try:\n                history = trainer.train(train_loader, val_loader, NUM_EPOCHS)\n                break\n            except OutOfMemoryAbort:\n                batch_size = max(1, batch_size // 2)\n                print(f\"[WARN] OOM fallback: retrying with batch size {batch_size}\")\n                del train_loader, val_loader, test_loaders\n                free_memory()\n        else:\n            raise RuntimeError(\"Training failed even with the smallest batch size.\")\n\n        if trainer.best_path.exists():\n            model.load_state_dict(torch_load(trainer.best_path, DEVICE))\n        results = cross_dataset_evaluation(model, test_loaders)\n        print(results.to_string(index=False))\n        results.to_csv(KAGGLE_PATHS[\"results\"] / \"cross_dataset_results.csv\", index=False)\n\n        calib = collect_predictions(model, val_loader, desc=\"Conformal calibration\")\n        conformal = ConformalPrediction(alpha=0.1)\n        conformal.calibrate(calib[\"probs\"], calib[\"labels\"])\n        primary_test = next(iter(test_loaders))\n        test_preds = collect_predictions(model, test_loaders[primary_test], desc=\"Conformal test\")\n        conformal_stats = conformal.evaluate(test_preds[\"probs\"], test_preds[\"labels\"])\n        print(f\"[INFO] Conformal (alpha=0.1) on {primary_test}: {conformal_stats}\")\n\n        figures = run_xai_examples(model, test_sets[primary_test])\n        print(f\"[INFO] Saved {len(figures)} XAI figures to {KAGGLE_PATHS['xai']}\")\n\n        torch.save(model.state_dict(), KAGGLE_PATHS[\"working\"] / \"forensicx_final.pt\")\n        summary = {\n            \"config\": {\"image_size\": IMAGE_SIZE, \"num_frames\": NUM_FRAMES, \"batch_size\": batch_size,\n                       \"epochs_run\": len(history), \"learning_rate\": LEARNING_RATE},\n            \"best_score\": trainer.best_score,\n            \"conformal\": conformal_stats,\n            \"cross_dataset\": results.to_dict(orient=\"records\"),\n        }\n        (KAGGLE_PATHS[\"results\"] / \"summary.json\").write_text(json.dumps(summary, indent=2, default=float))\n        print_gpu_memory(\"done\")\n        print(f\"[INFO] All outputs saved under {KAGGLE_PATHS['working']}\")\n    except (Exception, KeyboardInterrupt):\n        traceback.print_exc()\n        if trainer is not None:\n            path = trainer.save_checkpoint(max(len(trainer.history) - 1, 0), tag=\"emergency\")\n            print(f\"[INFO] Emergency checkpoint saved to {path}\")\n        raise\n\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-10-05T03:45:18.555724Z","iopub.execute_input":"2026-10-05T03:45:18.556112Z","iopub.status.idle":"2026-10-05T07:34:58.69987Z","shell.execute_reply.started":"2026-10-05T03:45:18.556083Z","shell.execute_reply":"2026-10-05T07:34:58.699149Z"}},"outputs":[],"execution_count":null}]}