{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🔍 AI-Generated Image Detector — Kaggle Training\n",
    "\n",
    "**Architecture:** SwinV2 + SRM High-Pass Filters + DCT + FFT Frequency Analysis\n",
    "\n",
    "**Dataset:** tristanzhang32/ai-generated-images-vs-real-images (60K images)\n",
    "\n",
    "**Hardware:** T4 x2 (2× 16GB VRAM)\n",
    "\n",
    "---\n",
    "\n",
    "### Setup Instructions\n",
    "1. **Add Dataset:** `tristanzhang32/ai-generated-images-vs-real-images`\n",
    "2. **Accelerator:** GPU T4 x2\n",
    "3. **Internet:** ON (needed for SwinV2 weights + pushing to HF Hub)\n",
    "4. **Kaggle Secret:** Add `HF_TOKEN` with your Hugging Face write token (for push_to_hub)\n",
    "5. **Run All**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "_cell_guid": "b1076dfc-b9ad-4769-8c92-a6c4dae69d19",
    "_uuid": "8f2839f25d086af736a60e9eeb907d3b93b6e0e5",
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# CELL 1: Install dependencies\n",
    "# ============================================================\n",
    "!pip install -q transformers accelerate evaluate trackio scikit-learn huggingface_hub"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# CELL 2: Imports + HF Login\n",
    "# ============================================================\n",
    "import os, io, math, random, json, glob\n",
    "import numpy as np\n",
    "from pathlib import Path\n",
    "from PIL import Image, ImageFilter\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.utils.data import Dataset, DataLoader, random_split\n",
    "\n",
    "from transformers import (\n",
    "    Swinv2Model,\n",
    "    TrainingArguments,\n",
    "    Trainer,\n",
    ")\n",
    "from torchvision.transforms import (\n",
    "    Compose, Normalize, Resize, CenterCrop, RandomResizedCrop,\n",
    "    RandomHorizontalFlip, ToTensor, ColorJitter,\n",
    ")\n",
    "import evaluate\n",
    "import trackio\n",
    "\n",
    "# Login to HF Hub (uses Kaggle Secret)\n",
    "from huggingface_hub import login\n",
    "from kaggle_secrets import UserSecretsClient\n",
    "try:\n",
    "    secrets = UserSecretsClient()\n",
    "    hf_token = secrets.get_secret(\"HF_TOKEN\")\n",
    "    login(token=hf_token)\n",
    "    print(\"✓ Logged in to Hugging Face Hub\")\n",
    "except Exception as e:\n",
    "    print(f\"⚠ Could not login to HF Hub: {e}\")\n",
    "    print(\"  Model will be saved locally only. Add HF_TOKEN to Kaggle Secrets to push to Hub.\")\n",
    "\n",
    "print(f\"PyTorch: {torch.__version__}\")\n",
    "print(f\"CUDA available: {torch.cuda.is_available()}\")\n",
    "print(f\"GPUs: {torch.cuda.device_count()}\")\n",
    "for i in range(torch.cuda.device_count()):\n",
    "    print(f\"  GPU {i}: {torch.cuda.get_device_name(i)} ({torch.cuda.get_device_properties(i).total_mem / 1e9:.1f} GB)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# CELL 3: Explore dataset structure\n",
    "# ============================================================\n",
    "BASE_PATH = Path(\"/kaggle/input/ai-generated-images-vs-real-images\")\n",
    "\n",
    "# Auto-detect folder structure\n",
    "if not BASE_PATH.exists():\n",
    "    # Try alternative paths\n",
    "    candidates = list(Path(\"/kaggle/input\").glob(\"**/real\")) + list(Path(\"/kaggle/input\").glob(\"**/Real\"))\n",
    "    if candidates:\n",
    "        BASE_PATH = candidates[0].parent\n",
    "        print(f\"Found dataset at: {BASE_PATH}\")\n",
    "    else:\n",
    "        # List everything in /kaggle/input\n",
    "        for p in sorted(Path(\"/kaggle/input\").rglob(\"*\")):\n",
    "            if p.is_dir():\n",
    "                n_files = len(list(p.iterdir()))\n",
    "                print(f\"  DIR  {p} ({n_files} items)\")\n",
    "        raise FileNotFoundError(\"Could not find dataset. Check the paths above and update BASE_PATH.\")\n",
    "\n",
    "print(f\"\\nDataset root: {BASE_PATH}\")\n",
    "print(f\"Contents:\")\n",
    "for item in sorted(BASE_PATH.iterdir()):\n",
    "    if item.is_dir():\n",
    "        n_files = len(list(item.rglob(\"*.*\")))\n",
    "        print(f\"  📁 {item.name}/ ({n_files} files)\")\n",
    "    else:\n",
    "        print(f\"  📄 {item.name} ({item.stat().st_size / 1e6:.1f} MB)\")\n",
    "\n",
    "# Detect class folders\n",
    "subdirs = sorted([d.name for d in BASE_PATH.iterdir() if d.is_dir()])\n",
    "print(f\"\\nSubdirectories: {subdirs}\")\n",
    "\n",
    "# Check for nested structure (e.g., train/real, train/fake)\n",
    "for sd in subdirs:\n",
    "    sub_subdirs = sorted([d.name for d in (BASE_PATH / sd).iterdir() if d.is_dir()])\n",
    "    if sub_subdirs:\n",
    "        print(f\"  {sd}/ contains: {sub_subdirs}\")\n",
    "        for ssd in sub_subdirs:\n",
    "            n = len(list((BASE_PATH / sd / ssd).rglob(\"*.*\")))\n",
    "            print(f\"    {ssd}/: {n} files\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# CELL 4: SRM Filter Bank (30 forensic high-pass filters)\n",
    "# ============================================================\n",
    "\n",
    "def get_srm_kernels():\n",
    "    \"\"\"30 SRM high-pass filter kernels (5x5) for forensic residual extraction.\n",
    "    Based on Fridrich & Kodovsky (2012) - Rich Models for Steganalysis.\n",
    "    These detect manipulation artifacts invisible in RGB space.\"\"\"\n",
    "    f1 = np.array([[0,0,0,0,0],[0,0,0,0,0],[0,0,-1,1,0],[0,0,0,0,0],[0,0,0,0,0]], dtype=np.float32)\n",
    "    f2 = np.array([[0,0,0,0,0],[0,0,0,0,0],[0,0,-1,0,0],[0,0,1,0,0],[0,0,0,0,0]], dtype=np.float32)\n",
    "    f3 = np.array([[0,0,0,0,0],[0,0,0,0,0],[0,1,-2,1,0],[0,0,0,0,0],[0,0,0,0,0]], dtype=np.float32)\n",
    "    f4 = np.array([[0,0,0,0,0],[0,0,1,0,0],[0,0,-2,0,0],[0,0,1,0,0],[0,0,0,0,0]], dtype=np.float32)\n",
    "    f5 = np.array([[0,0,0,0,0],[0,1,0,0,0],[0,0,-2,0,0],[0,0,0,1,0],[0,0,0,0,0]], dtype=np.float32)\n",
    "    f6 = np.array([[0,0,0,0,0],[0,0,0,1,0],[0,0,-2,0,0],[0,1,0,0,0],[0,0,0,0,0]], dtype=np.float32)\n",
    "    f7 = np.array([[0,0,0,0,0],[0,0,0,0,0],[0,-1,3,-3,1],[0,0,0,0,0],[0,0,0,0,0]], dtype=np.float32)\n",
    "    f8 = np.array([[0,0,0,0,0],[0,0,-1,0,0],[0,0,3,0,0],[0,0,-3,0,0],[0,0,1,0,0]], dtype=np.float32)\n",
    "    f9 = np.array([[0,0,0,0,0],[0,0,1,0,0],[0,1,-4,1,0],[0,0,1,0,0],[0,0,0,0,0]], dtype=np.float32)\n",
    "    f10 = np.array([[0,0,0,0,0],[0,1,1,1,0],[0,1,-8,1,0],[0,1,1,1,0],[0,0,0,0,0]], dtype=np.float32) / 3.0\n",
    "    f11 = np.array([[0,0,0,0,0],[0,-1,2,-1,0],[0,2,-4,2,0],[0,-1,2,-1,0],[0,0,0,0,0]], dtype=np.float32)\n",
    "    f12 = np.array([[0,0,0,0,0],[0,0,0,0,0],[-1,2,-2,2,-1],[0,0,0,0,0],[0,0,0,0,0]], dtype=np.float32) / 2.0\n",
    "    f13 = np.array([[0,0,-1,0,0],[0,0,2,0,0],[0,0,-2,0,0],[0,0,2,0,0],[0,0,-1,0,0]], dtype=np.float32) / 2.0\n",
    "    f14 = np.array([[0,0,-1,0,0],[0,0,2,0,0],[-1,2,-4,2,-1],[0,0,2,0,0],[0,0,-1,0,0]], dtype=np.float32) / 4.0\n",
    "    f15 = np.array([[-1,2,-2,2,-1],[2,-6,8,-6,2],[-2,8,-12,8,-2],[2,-6,8,-6,2],[-1,2,-2,2,-1]], dtype=np.float32) / 12.0\n",
    "    spam_h = np.array([[0,0,0,0,0],[0,0,0,0,0],[0,-1,2,-1,0],[0,0,0,0,0],[0,0,0,0,0]], dtype=np.float32)\n",
    "    spam_v = np.array([[0,0,0,0,0],[0,0,-1,0,0],[0,0,2,0,0],[0,0,-1,0,0],[0,0,0,0,0]], dtype=np.float32)\n",
    "    spam_d1 = np.array([[0,0,0,0,0],[0,-1,0,0,0],[0,0,2,0,0],[0,0,0,-1,0],[0,0,0,0,0]], dtype=np.float32)\n",
    "    spam_d2 = np.array([[0,0,0,0,0],[0,0,0,-1,0],[0,0,2,0,0],[0,-1,0,0,0],[0,0,0,0,0]], dtype=np.float32)\n",
    "    spam3_h = np.array([[0,0,0,0,0],[0,0,0,0,0],[1,-3,3,-1,0],[0,0,0,0,0],[0,0,0,0,0]], dtype=np.float32)\n",
    "    spam3_v = np.array([[0,0,1,0,0],[0,0,-3,0,0],[0,0,3,0,0],[0,0,-1,0,0],[0,0,0,0,0]], dtype=np.float32)\n",
    "    sq5_1 = np.array([[0,0,0,0,0],[0,0,1,0,0],[0,1,-4,1,0],[0,0,1,0,0],[0,0,0,0,0]], dtype=np.float32) / 2.0\n",
    "    sq5_2 = np.array([[0,0,0,0,0],[0,1,0,1,0],[0,0,-4,0,0],[0,1,0,1,0],[0,0,0,0,0]], dtype=np.float32) / 2.0\n",
    "    cross1 = np.array([[0,0,-1,0,0],[0,0,2,0,0],[0,0,-2,0,0],[0,0,2,0,0],[0,0,-1,0,0]], dtype=np.float32) / 2.0\n",
    "    cross2 = np.array([[0,0,0,0,0],[0,0,0,0,0],[-1,2,-2,2,-1],[0,0,0,0,0],[0,0,0,0,0]], dtype=np.float32) / 2.0\n",
    "    edge_d1 = np.array([[-1,0,0,0,0],[0,2,0,0,0],[0,0,-2,0,0],[0,0,0,2,0],[0,0,0,0,-1]], dtype=np.float32) / 2.0\n",
    "    edge_d2 = np.array([[0,0,0,0,-1],[0,0,0,2,0],[0,0,-2,0,0],[0,2,0,0,0],[-1,0,0,0,0]], dtype=np.float32) / 2.0\n",
    "    gabor_h = np.array([[0,0,0,0,0],[1,-1,0,-1,1],[0,0,0,0,0],[-1,1,0,1,-1],[0,0,0,0,0]], dtype=np.float32) / 4.0\n",
    "    gabor_v = np.array([[0,1,0,-1,0],[0,-1,0,1,0],[0,0,0,0,0],[0,-1,0,1,0],[0,1,0,-1,0]], dtype=np.float32) / 4.0\n",
    "    all_filters = [f1,f2,f3,f4,f5,f6,f7,f8,f9,f10,f11,f12,f13,f14,f15,\n",
    "                   spam_h,spam_v,spam_d1,spam_d2,spam3_h,spam3_v,\n",
    "                   sq5_1,sq5_2,cross1,cross2,edge_d1,edge_d2,gabor_h,gabor_v,f15]\n",
    "    return all_filters[:30]\n",
    "\n",
    "\n",
    "class SRMFilterBank(nn.Module):\n",
    "    \"\"\"30 fixed forensic high-pass filters. No learnable parameters.\"\"\"\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        kernels = get_srm_kernels()\n",
    "        weight = torch.stack([torch.tensor(k) for k in kernels]).unsqueeze(1)\n",
    "        weight = weight.repeat(1, 3, 1, 1) / 3.0\n",
    "        self.register_buffer('weight', weight)\n",
    "    def forward(self, x):\n",
    "        return F.conv2d(x, self.weight, padding=2)\n",
    "\n",
    "print(\"✓ SRM Filter Bank defined (30 kernels)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# CELL 5: DCT Frequency Analyzer\n",
    "# ============================================================\n",
    "\n",
    "class DCTFrequencyAnalyzer(nn.Module):\n",
    "    \"\"\"2D DCT on image patches → frequency band energy statistics.\n",
    "    AI images show different frequency distributions than natural images.\"\"\"\n",
    "    def __init__(self, patch_size=32, num_freq_bands=8):\n",
    "        super().__init__()\n",
    "        self.patch_size = patch_size\n",
    "        self.num_freq_bands = num_freq_bands\n",
    "        N = patch_size\n",
    "        dct_mat = torch.zeros(N, N)\n",
    "        for k in range(N):\n",
    "            for n in range(N):\n",
    "                if k == 0:\n",
    "                    dct_mat[k, n] = math.sqrt(1.0 / N)\n",
    "                else:\n",
    "                    dct_mat[k, n] = math.sqrt(2.0 / N) * math.cos(math.pi * (2*n + 1) * k / (2*N))\n",
    "        self.register_buffer('dct_mat', dct_mat)\n",
    "        self.register_buffer('dct_mat_t', dct_mat.t())\n",
    "\n",
    "    def dct2d(self, x):\n",
    "        return torch.matmul(torch.matmul(self.dct_mat, x), self.dct_mat_t)\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, C, H, W = x.shape\n",
    "        ps = self.patch_size\n",
    "        gray = 0.299 * x[:, 0] + 0.587 * x[:, 1] + 0.114 * x[:, 2]\n",
    "        h_patches = H // ps\n",
    "        w_patches = W // ps\n",
    "        gray = gray[:, :h_patches*ps, :w_patches*ps]\n",
    "        patches = gray.unfold(1, ps, ps).unfold(2, ps, ps)\n",
    "        B_p, hp, wp = patches.shape[:3]\n",
    "        patches = patches.reshape(B_p * hp * wp, ps, ps)\n",
    "        dct_patches = self.dct2d(patches)\n",
    "        dct_patches = dct_patches.reshape(B_p, hp * wp, ps, ps)\n",
    "\n",
    "        features = []\n",
    "        freq_y = torch.arange(ps, device=x.device).float()\n",
    "        freq_x = torch.arange(ps, device=x.device).float()\n",
    "        fy, fx = torch.meshgrid(freq_y, freq_x, indexing='ij')\n",
    "        freq_dist = torch.sqrt(fy**2 + fx**2)\n",
    "        max_freq = math.sqrt(2) * ps\n",
    "\n",
    "        for band in range(self.num_freq_bands):\n",
    "            lo = band * max_freq / self.num_freq_bands\n",
    "            hi = (band + 1) * max_freq / self.num_freq_bands\n",
    "            mask = ((freq_dist >= lo) & (freq_dist < hi)).float()\n",
    "            band_energy = (dct_patches ** 2 * mask.unsqueeze(0).unsqueeze(0)).sum(dim=(-2, -1))\n",
    "            features.append(band_energy.mean(dim=1, keepdim=True))\n",
    "            features.append(band_energy.std(dim=1, keepdim=True))\n",
    "\n",
    "        total_energy = (dct_patches ** 2).sum(dim=(-2, -1))\n",
    "        weighted_freq = (dct_patches ** 2 * freq_dist.unsqueeze(0).unsqueeze(0)).sum(dim=(-2, -1))\n",
    "        spectral_centroid = weighted_freq / (total_energy + 1e-8)\n",
    "        features.append(spectral_centroid.mean(dim=1, keepdim=True))\n",
    "        features.append(spectral_centroid.std(dim=1, keepdim=True))\n",
    "\n",
    "        mid = ps // 2\n",
    "        low_mask = (freq_dist < mid).float()\n",
    "        high_mask = (freq_dist >= mid).float()\n",
    "        low_energy = (dct_patches ** 2 * low_mask).sum(dim=(-2, -1))\n",
    "        high_energy = (dct_patches ** 2 * high_mask).sum(dim=(-2, -1))\n",
    "        hl_ratio = high_energy / (low_energy + 1e-8)\n",
    "        features.append(hl_ratio.mean(dim=1, keepdim=True))\n",
    "        features.append(hl_ratio.std(dim=1, keepdim=True))\n",
    "\n",
    "        dc_values = dct_patches[:, :, 0, 0]\n",
    "        features.append(dc_values.mean(dim=1, keepdim=True))\n",
    "        features.append(dc_values.std(dim=1, keepdim=True))\n",
    "\n",
    "        return torch.cat(features, dim=1)\n",
    "\n",
    "print(\"✓ DCT Frequency Analyzer defined\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# CELL 6: FFT Spectral Analyzer\n",
    "# ============================================================\n",
    "\n",
    "class FFTSpectralAnalyzer(nn.Module):\n",
    "    \"\"\"Azimuthally averaged FFT power spectrum.\n",
    "    Natural images follow 1/f^2 power law. AI images deviate.\"\"\"\n",
    "    def __init__(self, num_bins=32):\n",
    "        super().__init__()\n",
    "        self.num_bins = num_bins\n",
    "\n",
    "    def forward(self, x):\n",
    "        B, C, H, W = x.shape\n",
    "        gray = 0.299 * x[:, 0] + 0.587 * x[:, 1] + 0.114 * x[:, 2]\n",
    "        hann_y = torch.hann_window(H, device=x.device)\n",
    "        hann_x = torch.hann_window(W, device=x.device)\n",
    "        window = hann_y.unsqueeze(1) * hann_x.unsqueeze(0)\n",
    "        gray = gray * window.unsqueeze(0)\n",
    "\n",
    "        fft = torch.fft.fft2(gray)\n",
    "        fft_shift = torch.fft.fftshift(fft)\n",
    "        power = torch.abs(fft_shift) ** 2\n",
    "\n",
    "        cy, cx = H // 2, W // 2\n",
    "        y = torch.arange(H, device=x.device).float() - cy\n",
    "        xx = torch.arange(W, device=x.device).float() - cx\n",
    "        yy, xx = torch.meshgrid(y, xx, indexing='ij')\n",
    "        radius = torch.sqrt(yy**2 + xx**2)\n",
    "\n",
    "        max_radius = min(cy, cx)\n",
    "        bin_width = max_radius / self.num_bins\n",
    "\n",
    "        features = []\n",
    "        for i in range(self.num_bins):\n",
    "            r_lo = i * bin_width\n",
    "            r_hi = (i + 1) * bin_width\n",
    "            mask = ((radius >= r_lo) & (radius < r_hi)).float()\n",
    "            count = mask.sum() + 1e-8\n",
    "            bin_power = (power * mask.unsqueeze(0)).sum(dim=(-2, -1)) / count\n",
    "            features.append(bin_power.unsqueeze(1))\n",
    "\n",
    "        radial_spectrum = torch.cat(features, dim=1)\n",
    "        log_spectrum = torch.log1p(radial_spectrum)\n",
    "\n",
    "        log_freq = torch.log1p(torch.arange(self.num_bins, device=x.device).float() + 1)\n",
    "        log_freq = log_freq.unsqueeze(0).expand(B, -1)\n",
    "\n",
    "        xm = log_freq - log_freq.mean(dim=1, keepdim=True)\n",
    "        ym = log_spectrum - log_spectrum.mean(dim=1, keepdim=True)\n",
    "        slope = (xm * ym).sum(dim=1, keepdim=True) / ((xm**2).sum(dim=1, keepdim=True) + 1e-8)\n",
    "        intercept = log_spectrum.mean(dim=1, keepdim=True) - slope * log_freq.mean(dim=1, keepdim=True)\n",
    "\n",
    "        predicted = slope * log_freq + intercept\n",
    "        residuals = log_spectrum - predicted\n",
    "        residual_std = residuals.std(dim=1, keepdim=True)\n",
    "        residual_max = residuals.max(dim=1, keepdim=True)[0]\n",
    "\n",
    "        return torch.cat([log_spectrum, slope, intercept, residual_std, residual_max], dim=1)\n",
    "\n",
    "print(\"✓ FFT Spectral Analyzer defined\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# CELL 7: Full Model — FrequencyAwareDetector\n",
    "# ============================================================\n",
    "\n",
    "class FrequencyAwareDetector(nn.Module):\n",
    "    \"\"\"\n",
    "    4-Branch Fusion Detector:\n",
    "      Branch 1: SwinV2-Tiny (semantic features)          → 768d\n",
    "      Branch 2: SRM filter bank → CNN (forensic residuals) → 256d\n",
    "      Branch 3: DCT frequency band statistics              → 22d\n",
    "      Branch 4: FFT radial power spectrum                  → 36d\n",
    "      Fusion MLP: (768 + 128) → 512 → 128 → 2\n",
    "    \"\"\"\n",
    "    def __init__(self, backbone_name=\"microsoft/swinv2-tiny-patch4-window8-256\",\n",
    "                 num_labels=2, dct_patch_size=32, num_freq_bands=8, fft_bins=32):\n",
    "        super().__init__()\n",
    "        self.num_labels = num_labels\n",
    "        self.supports_gradient_checkpointing = True\n",
    "\n",
    "        # Branch 1: SwinV2 backbone\n",
    "        self.backbone = Swinv2Model.from_pretrained(backbone_name)\n",
    "        backbone_dim = self.backbone.config.hidden_size  # 768\n",
    "\n",
    "        # Branch 2: SRM forensic features\n",
    "        self.srm = SRMFilterBank()\n",
    "        self.srm_encoder = nn.Sequential(\n",
    "            nn.Conv2d(30, 64, 3, stride=2, padding=1), nn.BatchNorm2d(64), nn.GELU(),\n",
    "            nn.Conv2d(64, 128, 3, stride=2, padding=1), nn.BatchNorm2d(128), nn.GELU(),\n",
    "            nn.Conv2d(128, 256, 3, stride=2, padding=1), nn.BatchNorm2d(256), nn.GELU(),\n",
    "            nn.AdaptiveAvgPool2d((1, 1)), nn.Flatten(),\n",
    "        )\n",
    "        srm_dim = 256\n",
    "\n",
    "        # Branch 3: DCT\n",
    "        self.dct_analyzer = DCTFrequencyAnalyzer(patch_size=dct_patch_size, num_freq_bands=num_freq_bands)\n",
    "        dct_dim = num_freq_bands * 2 + 6  # 22\n",
    "\n",
    "        # Branch 4: FFT\n",
    "        self.fft_analyzer = FFTSpectralAnalyzer(num_bins=fft_bins)\n",
    "        fft_dim = fft_bins + 4  # 36\n",
    "\n",
    "        # Frequency projection\n",
    "        freq_total_dim = srm_dim + dct_dim + fft_dim  # 314\n",
    "        self.freq_proj = nn.Sequential(\n",
    "            nn.Linear(freq_total_dim, 256), nn.GELU(), nn.Dropout(0.3),\n",
    "            nn.Linear(256, 128),\n",
    "        )\n",
    "\n",
    "        # Fusion classifier\n",
    "        self.classifier = nn.Sequential(\n",
    "            nn.Linear(backbone_dim + 128, 512), nn.GELU(), nn.Dropout(0.3),\n",
    "            nn.Linear(512, 128), nn.GELU(), nn.Dropout(0.1),\n",
    "            nn.Linear(128, num_labels),\n",
    "        )\n",
    "\n",
    "        self.loss_fn = nn.CrossEntropyLoss(label_smoothing=0.1)\n",
    "\n",
    "    def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None):\n",
    "        self.backbone.gradient_checkpointing_enable(gradient_checkpointing_kwargs=gradient_checkpointing_kwargs)\n",
    "\n",
    "    def gradient_checkpointing_disable(self):\n",
    "        self.backbone.gradient_checkpointing_disable()\n",
    "\n",
    "    def forward(self, pixel_values, labels=None, **kwargs):\n",
    "        semantic_feats = self.backbone(pixel_values=pixel_values).pooler_output\n",
    "        srm_feats = self.srm_encoder(self.srm(pixel_values))\n",
    "        dct_feats = self.dct_analyzer(pixel_values)\n",
    "        fft_feats = self.fft_analyzer(pixel_values)\n",
    "\n",
    "        freq_feats = torch.cat([srm_feats, dct_feats, fft_feats], dim=1)\n",
    "        freq_proj = self.freq_proj(freq_feats)\n",
    "\n",
    "        fused = torch.cat([semantic_feats, freq_proj], dim=1)\n",
    "        logits = self.classifier(fused)\n",
    "\n",
    "        loss = None\n",
    "        if labels is not None:\n",
    "            loss = self.loss_fn(logits, labels)\n",
    "\n",
    "        return {\"loss\": loss, \"logits\": logits}\n",
    "\n",
    "\n",
    "# Quick test\n",
    "print(\"Building model...\")\n",
    "model = FrequencyAwareDetector()\n",
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"✓ Model built: {total_params:,} parameters ({total_params/1e6:.1f}M)\")\n",
    "del model\n",
    "torch.cuda.empty_cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# CELL 8: Dataset — load from Kaggle local path\n",
    "# ============================================================\n",
    "\n",
    "class AIImageDataset(Dataset):\n",
    "    \"\"\"\n",
    "    Loads images from folder structure. Auto-detects:\n",
    "      - real/ + fake/  (or REAL/ + FAKE/)\n",
    "      - train/real/ + train/fake/ (nested)\n",
    "      - any folder with 'real' in name → label 0, others → label 1\n",
    "    \"\"\"\n",
    "    def __init__(self, root_path, transform=None, is_train=True):\n",
    "        self.transform = transform\n",
    "        self.is_train = is_train\n",
    "        self.samples = []  # list of (path, label)\n",
    "\n",
    "        root = Path(root_path)\n",
    "        subdirs = sorted([d for d in root.iterdir() if d.is_dir()])\n",
    "        subdir_names = [d.name.lower() for d in subdirs]\n",
    "\n",
    "        # Detect structure\n",
    "        if 'test' in subdir_names and 'train' in subdir_names:\n",
    "            # Nested: train/real, train/fake, test/real, test/fake\n",
    "            print(\"Detected train/test split structure\")\n",
    "            for split_dir in subdirs:\n",
    "                self._load_class_folders(split_dir)\n",
    "        else:\n",
    "            self._load_class_folders(root)\n",
    "\n",
    "        random.shuffle(self.samples)\n",
    "        labels = [s[1] for s in self.samples]\n",
    "        print(f\"  Total: {len(self.samples)} images\")\n",
    "        print(f\"  Real (0): {labels.count(0)}, Fake (1): {labels.count(1)}\")\n",
    "\n",
    "    def _load_class_folders(self, parent):\n",
    "        extensions = {'.jpg', '.jpeg', '.png', '.webp', '.bmp', '.tiff'}\n",
    "        for class_dir in sorted(parent.iterdir()):\n",
    "            if not class_dir.is_dir():\n",
    "                continue\n",
    "            name_lower = class_dir.name.lower()\n",
    "            # Determine label\n",
    "            if any(k in name_lower for k in ['real', 'authentic', 'natural', 'photo']):\n",
    "                label = 0\n",
    "            elif any(k in name_lower for k in ['fake', 'ai', 'generated', 'synthetic', 'midjourney',\n",
    "                                                'stable', 'dall', 'diffusion']):\n",
    "                label = 1\n",
    "            else:\n",
    "                print(f\"  ⚠ Unknown folder '{class_dir.name}' — skipping. \"\n",
    "                      f\"Rename to 'real' or 'fake' to include.\")\n",
    "                continue\n",
    "\n",
    "            files = [f for f in class_dir.rglob('*') if f.suffix.lower() in extensions]\n",
    "            self.samples.extend([(str(f), label) for f in files])\n",
    "            label_name = 'real' if label == 0 else 'fake'\n",
    "            print(f\"  📁 {class_dir.name}/: {len(files)} files → label={label_name}\")\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.samples)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        path, label = self.samples[idx]\n",
    "        try:\n",
    "            img = Image.open(path).convert('RGB')\n",
    "        except Exception:\n",
    "            # Return a black image on error to avoid crashing\n",
    "            img = Image.new('RGB', (256, 256), (0, 0, 0))\n",
    "\n",
    "        if self.is_train:\n",
    "            img = self._augment(img)\n",
    "\n",
    "        if self.transform:\n",
    "            pixel_values = self.transform(img)\n",
    "        else:\n",
    "            pixel_values = ToTensor()(img)\n",
    "\n",
    "        return {'pixel_values': pixel_values, 'labels': label}\n",
    "\n",
    "    def _augment(self, img):\n",
    "        \"\"\"Social media simulation: JPEG compression, blur, downscale.\"\"\"\n",
    "        if random.random() < 0.10:\n",
    "            quality = random.randint(30, 95)\n",
    "            buf = io.BytesIO()\n",
    "            img.save(buf, format='JPEG', quality=quality)\n",
    "            buf.seek(0)\n",
    "            img = Image.open(buf).convert('RGB')\n",
    "        if random.random() < 0.10:\n",
    "            img = img.filter(ImageFilter.GaussianBlur(radius=random.uniform(0.1, 2.0)))\n",
    "        if random.random() < 0.05:\n",
    "            w, h = img.size\n",
    "            s = random.uniform(0.5, 0.9)\n",
    "            img = img.resize((int(w*s), int(h*s)), Image.BILINEAR).resize((w, h), Image.BILINEAR)\n",
    "        return img\n",
    "\n",
    "\n",
    "# ---------- Load dataset ----------\n",
    "print(\"Loading dataset...\")\n",
    "IMG_SIZE = 256\n",
    "\n",
    "train_transform = Compose([\n",
    "    RandomResizedCrop((IMG_SIZE, IMG_SIZE), scale=(0.8, 1.0)),\n",
    "    RandomHorizontalFlip(),\n",
    "    ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1),\n",
    "    ToTensor(),\n",
    "    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "])\n",
    "\n",
    "eval_transform = Compose([\n",
    "    Resize((IMG_SIZE + 32, IMG_SIZE + 32)),\n",
    "    CenterCrop((IMG_SIZE, IMG_SIZE)),\n",
    "    ToTensor(),\n",
    "    Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n",
    "])\n",
    "\n",
    "full_dataset = AIImageDataset(BASE_PATH, transform=None, is_train=False)\n",
    "\n",
    "# Stratified 90/10 split\n",
    "n_total = len(full_dataset)\n",
    "n_eval = max(int(0.10 * n_total), 100)\n",
    "n_train = n_total - n_eval\n",
    "\n",
    "generator = torch.Generator().manual_seed(42)\n",
    "train_indices, eval_indices = random_split(range(n_total), [n_train, n_eval], generator=generator)\n",
    "\n",
    "class SubsetWithTransform(Dataset):\n",
    "    def __init__(self, base_dataset, indices, transform, is_train):\n",
    "        self.base = base_dataset\n",
    "        self.indices = list(indices)\n",
    "        self.transform = transform\n",
    "        self.is_train = is_train\n",
    "    def __len__(self):\n",
    "        return len(self.indices)\n",
    "    def __getitem__(self, idx):\n",
    "        real_idx = self.indices[idx]\n",
    "        path, label = self.base.samples[real_idx]\n",
    "        try:\n",
    "            img = Image.open(path).convert('RGB')\n",
    "        except Exception:\n",
    "            img = Image.new('RGB', (256, 256), (0, 0, 0))\n",
    "        if self.is_train:\n",
    "            img = self.base._augment(img)\n",
    "        pixel_values = self.transform(img)\n",
    "        return {'pixel_values': pixel_values, 'labels': label}\n",
    "\n",
    "train_dataset = SubsetWithTransform(full_dataset, train_indices, train_transform, is_train=True)\n",
    "eval_dataset = SubsetWithTransform(full_dataset, eval_indices, eval_transform, is_train=False)\n",
    "\n",
    "print(f\"\\n✓ Train: {len(train_dataset)}, Eval: {len(eval_dataset)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# CELL 9: Training Configuration\n",
    "# ============================================================\n",
    "\n",
    "# -------- CONFIGURATION — edit these --------\n",
    "HUB_MODEL_ID = \"Reju983/ai-generated-image-detector\"   # ← your HF repo\n",
    "NUM_EPOCHS = 5\n",
    "BATCH_SIZE = 12          # per GPU → effective = 12 * 2 GPUs * 2 accum = 48\n",
    "GRAD_ACCUM = 2\n",
    "LEARNING_RATE = 2e-5\n",
    "OUTPUT_DIR = \"/kaggle/working/ai-image-detector\"\n",
    "# -------------------------------------------\n",
    "\n",
    "print(f\"Effective batch size: {BATCH_SIZE} × {torch.cuda.device_count()} GPUs × {GRAD_ACCUM} accum = \"\n",
    "      f\"{BATCH_SIZE * torch.cuda.device_count() * GRAD_ACCUM}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# CELL 10: Custom Trainer + Metrics\n",
    "# ============================================================\n",
    "\n",
    "class FreqDetectorTrainer(Trainer):\n",
    "    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):\n",
    "        labels = inputs.pop(\"labels\")\n",
    "        outputs = model(pixel_values=inputs[\"pixel_values\"], labels=labels)\n",
    "        loss = outputs[\"loss\"]\n",
    "        return (loss, outputs) if return_outputs else loss\n",
    "\n",
    "\n",
    "def data_collator(features):\n",
    "    pixel_values = torch.stack([f[\"pixel_values\"] for f in features])\n",
    "    labels = torch.tensor([f[\"labels\"] for f in features], dtype=torch.long)\n",
    "    return {\"pixel_values\": pixel_values, \"labels\": labels}\n",
    "\n",
    "\n",
    "accuracy_metric = evaluate.load(\"accuracy\")\n",
    "\n",
    "def compute_metrics(eval_pred):\n",
    "    preds, labels = eval_pred\n",
    "    if isinstance(preds, dict):\n",
    "        preds = preds[\"logits\"]\n",
    "    preds = np.argmax(preds, axis=1)\n",
    "    acc = accuracy_metric.compute(predictions=preds, references=labels)\n",
    "    real_mask = labels == 0\n",
    "    fake_mask = labels == 1\n",
    "    real_acc = float((preds[real_mask] == labels[real_mask]).mean()) if real_mask.sum() > 0 else 0.0\n",
    "    fake_acc = float((preds[fake_mask] == labels[fake_mask]).mean()) if fake_mask.sum() > 0 else 0.0\n",
    "    return {\n",
    "        \"accuracy\": acc[\"accuracy\"],\n",
    "        \"real_accuracy\": real_acc,\n",
    "        \"fake_accuracy\": fake_acc,\n",
    "    }\n",
    "\n",
    "print(\"✓ Trainer + Metrics ready\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# CELL 11: 🚀 TRAIN!\n",
    "# ============================================================\n",
    "\n",
    "# Init tracking\n",
    "trackio.init(project=\"ai-image-detector\", name=\"kaggle-swinv2-srm-dct-fft\")\n",
    "\n",
    "# Build model fresh\n",
    "print(\"Building model...\")\n",
    "model = FrequencyAwareDetector(\n",
    "    backbone_name=\"microsoft/swinv2-tiny-patch4-window8-256\",\n",
    "    num_labels=2, dct_patch_size=32, num_freq_bands=8, fft_bins=32,\n",
    ")\n",
    "total_params = sum(p.numel() for p in model.parameters())\n",
    "print(f\"Parameters: {total_params:,}\")\n",
    "\n",
    "# Training args — optimized for Kaggle T4 x2\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=OUTPUT_DIR,\n",
    "    remove_unused_columns=False,\n",
    "\n",
    "    # Schedule\n",
    "    num_train_epochs=NUM_EPOCHS,\n",
    "    learning_rate=LEARNING_RATE,\n",
    "    warmup_ratio=0.1,\n",
    "    weight_decay=0.01,\n",
    "    lr_scheduler_type=\"cosine\",\n",
    "\n",
    "    # Batch\n",
    "    per_device_train_batch_size=BATCH_SIZE,\n",
    "    per_device_eval_batch_size=BATCH_SIZE,\n",
    "    gradient_accumulation_steps=GRAD_ACCUM,\n",
    "\n",
    "    # Precision — T4 supports fp16 (NOT bf16)\n",
    "    fp16=True,\n",
    "    bf16=False,\n",
    "\n",
    "    # Eval & Save\n",
    "    eval_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model=\"accuracy\",\n",
    "    greater_is_better=True,\n",
    "    save_total_limit=2,\n",
    "\n",
    "    # Logging\n",
    "    logging_steps=50,\n",
    "    logging_strategy=\"steps\",\n",
    "    logging_first_step=True,\n",
    "    disable_tqdm=False,  # keep tqdm in notebooks\n",
    "    report_to=\"trackio\",\n",
    "    run_name=\"kaggle-swinv2-srm-dct-fft\",\n",
    "\n",
    "    # Memory\n",
    "    gradient_checkpointing=True,\n",
    "    dataloader_num_workers=4,\n",
    "    dataloader_pin_memory=True,\n",
    "\n",
    "    # Hub\n",
    "    push_to_hub=True,\n",
    "    hub_model_id=HUB_MODEL_ID,\n",
    "    hub_strategy=\"end\",\n",
    "\n",
    "    label_names=[\"labels\"],\n",
    ")\n",
    "\n",
    "trainer = FreqDetectorTrainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    data_collator=data_collator,\n",
    "    train_dataset=train_dataset,\n",
    "    eval_dataset=eval_dataset,\n",
    "    compute_metrics=compute_metrics,\n",
    ")\n",
    "\n",
    "print(\"\\n\" + \"=\"*60)\n",
    "print(\"Starting training...\")\n",
    "print(f\"  Train samples: {len(train_dataset)}\")\n",
    "print(f\"  Eval samples:  {len(eval_dataset)}\")\n",
    "print(f\"  Epochs: {NUM_EPOCHS}\")\n",
    "print(f\"  Batch: {BATCH_SIZE} x {torch.cuda.device_count()} GPUs x {GRAD_ACCUM} accum\")\n",
    "print(f\"  LR: {LEARNING_RATE}\")\n",
    "print(\"=\"*60 + \"\\n\")\n",
    "\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# CELL 12: Evaluate + Push to Hub\n",
    "# ============================================================\n",
    "\n",
    "print(\"\\n\" + \"=\"*60)\n",
    "print(\"Final Evaluation:\")\n",
    "metrics = trainer.evaluate()\n",
    "for k, v in metrics.items():\n",
    "    print(f\"  {k}: {v}\")\n",
    "print(\"=\"*60)\n",
    "\n",
    "# Save model state dict\n",
    "save_dir = os.path.join(OUTPUT_DIR, \"final_model\")\n",
    "os.makedirs(save_dir, exist_ok=True)\n",
    "torch.save(model.state_dict(), os.path.join(save_dir, \"model_state_dict.pt\"))\n",
    "print(f\"\\n✓ Model state dict saved to {save_dir}/model_state_dict.pt\")\n",
    "\n",
    "# Save config\n",
    "config = {\n",
    "    \"architecture\": \"FrequencyAwareDetector\",\n",
    "    \"backbone_name\": \"microsoft/swinv2-tiny-patch4-window8-256\",\n",
    "    \"num_labels\": 2, \"dct_patch_size\": 32, \"num_freq_bands\": 8, \"fft_bins\": 32,\n",
    "    \"id2label\": {\"0\": \"real\", \"1\": \"ai_generated\"},\n",
    "    \"label2id\": {\"real\": 0, \"ai_generated\": 1},\n",
    "    \"training_metrics\": metrics,\n",
    "}\n",
    "with open(os.path.join(save_dir, \"detector_config.json\"), \"w\") as f:\n",
    "    json.dump(config, f, indent=2)\n",
    "\n",
    "# Push to Hub\n",
    "try:\n",
    "    print(\"\\nPushing to Hub...\")\n",
    "    trainer.push_to_hub(\n",
    "        commit_message=f\"Kaggle training: acc={metrics.get('eval_accuracy', 'N/A')}\",\n",
    "        tags=[\"image-classification\", \"ai-image-detection\", \"deepfake-detection\",\n",
    "              \"frequency-analysis\", \"swinv2\", \"srm\", \"dct\", \"fft\"],\n",
    "    )\n",
    "\n",
    "    from huggingface_hub import HfApi\n",
    "    api = HfApi()\n",
    "    api.upload_file(\n",
    "        path_or_fileobj=os.path.join(save_dir, \"model_state_dict.pt\"),\n",
    "        path_in_repo=\"model_state_dict.pt\",\n",
    "        repo_id=HUB_MODEL_ID,\n",
    "    )\n",
    "    api.upload_file(\n",
    "        path_or_fileobj=os.path.join(save_dir, \"detector_config.json\"),\n",
    "        path_in_repo=\"detector_config.json\",\n",
    "        repo_id=HUB_MODEL_ID,\n",
    "    )\n",
    "    print(f\"\\n✓ Model pushed to https://huggingface.co/{HUB_MODEL_ID}\")\n",
    "except Exception as e:\n",
    "    print(f\"\\n⚠ Hub push failed: {e}\")\n",
    "    print(f\"  Model saved locally at: {save_dir}\")\n",
    "    print(f\"  Download from Kaggle Output tab.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "trusted": true
   },
   "outputs": [],
   "source": [
    "# ============================================================\n",
    "# CELL 13: Quick Inference Test\n",
    "# ============================================================\n",
    "\n",
    "model.eval()\n",
    "device = next(model.parameters()).device\n",
    "\n",
    "# Test on a few images from the eval set\n",
    "print(\"\\n🔍 Sample predictions:\\n\")\n",
    "for i in range(min(10, len(eval_dataset))):\n",
    "    item = eval_dataset[i]\n",
    "    pv = item['pixel_values'].unsqueeze(0).to(device)\n",
    "    true_label = item['labels']\n",
    "\n",
    "    with torch.no_grad():\n",
    "        out = model(pixel_values=pv)\n",
    "        probs = torch.softmax(out['logits'], dim=1)\n",
    "        pred = probs.argmax(dim=1).item()\n",
    "        conf = probs[0][pred].item()\n",
    "\n",
    "    labels = {0: 'Real', 1: 'AI-Gen'}\n",
    "    correct = '✓' if pred == true_label else '✗'\n",
    "    print(f\"  {correct} True={labels[true_label]:6s} | Pred={labels[pred]:6s} ({conf:.1%})\")\n",
    "\n",
    "print(\"\\n🎉 Training complete!\")"
   ]
  }
 ],
 "metadata": {
  "kaggle": {
   "accelerator": "nvidiaTeslaT4",
   "dataSources": [
    {
     "datasetId": 0,
     "sourceId": 0,
     "sourceType": "datasetVersion"
    }
   ],
   "dockerImageVersionId": 30000,
   "isGpuEnabled": true,
   "isInternetEnabled": true,
   "language": "python",
   "sourceType": "notebook"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.10.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}