backend/train_image_improved.py

"""
train_image_improved.py — Improved EfficientNetV2-S image deepfake detector training
Optimised for small datasets (~2000 images). Key improvements over original:

 1. Stronger augmentation (RandomErasing, GaussianBlur, Random90°rotations)
 2. Label smoothing to prevent overconfident predictions
 3. Warmup + CosineAnnealing LR schedule (better convergence)
 4. Gradient accumulation (effective larger batch without RAM cost)
 5. Early stopping with patience
 6. Test-time augmentation (TTA) during validation for better AUC estimates
 7. Longer Phase-2 fine-tune with layer-wise LR decay

Dataset structure (same as original):
  datasets/images/
    train/
      real/  ← real face crops (JPG/PNG)
      fake/  ← deepfake face crops (JPG/PNG)
    val/
      real/
      fake/

Usage:
  python train_image_improved.py \
    --data_dir   ../datasets/images \
    --epochs     40 \
    --finetune_epochs 20 \
    --batch_size 16 \
    --lr         2e-4 \
    --save_dir   models
"""

import argparse, os, time, random
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
from torch.optim import AdamW
from torch.optim.lr_scheduler import OneCycleLR
from sklearn.metrics import roc_auc_score, accuracy_score, f1_score
from PIL import Image
from tqdm import tqdm
import torchvision.transforms as T

from image_detector import EfficientNetV2Detector, inference_transform

# ─────────────────────────────────────────────────────────────────
SEED     = 42
IMG_EXTS = {".jpg", ".jpeg", ".png", ".webp", ".bmp"}
MEAN     = [0.485, 0.456, 0.406]
STD      = [0.229, 0.224, 0.225]


def seed_everything(seed: int = SEED):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


# ─────────────────────────────────────────────────────────────────
# IMPROVED Augmentation Pipeline
# ─────────────────────────────────────────────────────────────────
strong_train_transform = T.Compose([
    T.Resize((256, 256)),
    T.RandomCrop(224),
    T.RandomHorizontalFlip(p=0.5),
    T.RandomApply([T.RandomRotation(degrees=10)], p=0.4),
    T.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.05),
    T.RandomApply([T.GaussianBlur(kernel_size=5, sigma=(0.1, 2.0))], p=0.3),
    T.RandomGrayscale(p=0.05),
    T.ToTensor(),
    T.Normalize(mean=MEAN, std=STD),
    T.RandomErasing(p=0.25, scale=(0.02, 0.15)),   # simulate occlusions
])

# TTA transforms (multiple crops + flips averaged at val time)
tta_transforms = [
    T.Compose([T.Resize((224, 224)), T.ToTensor(), T.Normalize(MEAN, STD)]),
    T.Compose([T.Resize((256, 256)), T.CenterCrop(224), T.ToTensor(), T.Normalize(MEAN, STD)]),
    T.Compose([T.Resize((224, 224)), T.RandomHorizontalFlip(p=1.0), T.ToTensor(), T.Normalize(MEAN, STD)]),
]


# ─────────────────────────────────────────────────────────────────
# Dataset
# ─────────────────────────────────────────────────────────────────
class ImageFakeDataset(Dataset):
    def __init__(self, root: str, split: str, augment: bool = False):
        self.augment   = augment
        self.transform = strong_train_transform if augment else inference_transform
        self.samples   = []

        for label, name in [(0, "real"), (1, "fake")]:
            folder = Path(root) / split / name
            if not folder.exists():
                print(f"[WARN] Folder not found: {folder}")
                continue
            for f in folder.rglob("*"):
                if f.suffix.lower() in IMG_EXTS:
                    self.samples.append((str(f), label))

        random.shuffle(self.samples)
        real_n = sum(1 for _, l in self.samples if l == 0)
        fake_n = sum(1 for _, l in self.samples if l == 1)
        print(f"[Dataset/{split}] real={real_n}  fake={fake_n}  total={len(self.samples)}")

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        path, label = self.samples[idx]
        try:
            img = Image.open(path).convert("RGB")
        except Exception:
            img = Image.new("RGB", (224, 224), (128, 128, 128))
        tensor = self.transform(img)
        return tensor, torch.tensor(label, dtype=torch.float32)

    def get_weights(self):
        """Per-sample weights for WeightedRandomSampler (balances classes)."""
        labels  = [l for _, l in self.samples]
        n_real  = labels.count(0)
        n_fake  = labels.count(1)
        w_real  = 1.0 / n_real if n_real else 1.0
        w_fake  = 1.0 / n_fake if n_fake else 1.0
        return [w_real if l == 0 else w_fake for l in labels]


# ─────────────────────────────────────────────────────────────────
# Label-smoothed BCE loss
# ─────────────────────────────────────────────────────────────────
class LabelSmoothingBCE(nn.Module):
    """BCEWithLogitsLoss with label smoothing for better generalisation."""
    def __init__(self, smoothing: float = 0.1, pos_weight=None):
        super().__init__()
        self.smoothing  = smoothing
        self.pos_weight = pos_weight

    def forward(self, logits, targets):
        targets = targets * (1 - self.smoothing) + 0.5 * self.smoothing
        return nn.functional.binary_cross_entropy_with_logits(
            logits, targets,
            pos_weight=self.pos_weight,
        )


# ─────────────────────────────────────────────────────────────────
# TTA inference helper
# ─────────────────────────────────────────────────────────────────
def tta_predict(model, pil_img, device):
    probs = []
    for tf in tta_transforms:
        t = tf(pil_img).unsqueeze(0).to(device)
        with torch.no_grad():
            p = torch.sigmoid(model(t)).item()
        probs.append(p)
    return np.mean(probs)


# ─────────────────────────────────────────────────────────────────
# Epoch runner
# ─────────────────────────────────────────────────────────────────
def run_epoch(model, loader, criterion, optimizer, device, training: bool,
              scheduler=None, accumulation_steps: int = 1):
    model.train(training)
    total_loss, all_probs, all_labels = 0.0, [], []
    desc = "Train" if training else "Val  "
    pbar = tqdm(loader, desc=desc, leave=False, dynamic_ncols=True)

    optimizer_step_count = 0
    if training:
        optimizer.zero_grad()

    for step, (imgs, labels) in enumerate(pbar):
        imgs   = imgs.to(device)
        labels = labels.to(device)

        with torch.set_grad_enabled(training):
            logits = model(imgs).squeeze(1)
            loss   = criterion(logits, labels)

            if training:
                (loss / accumulation_steps).backward()
                if (step + 1) % accumulation_steps == 0:
                    nn.utils.clip_grad_norm_(model.parameters(), 1.0)
                    optimizer.step()
                    if scheduler is not None:
                        scheduler.step()
                    optimizer.zero_grad()
                    optimizer_step_count += 1

        total_loss += loss.item() * len(labels)
        probs = torch.sigmoid(logits).detach().cpu().numpy()
        all_probs.extend(probs.tolist())
        all_labels.extend(labels.cpu().numpy().tolist())
        pbar.set_postfix(loss=f"{(total_loss / max(len(all_labels), 1)):.4f}")

    # flush remaining gradient
    if training and (len(loader) % accumulation_steps != 0):
        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        optimizer.step()
        optimizer.zero_grad()

    avg_loss = total_loss / max(len(loader.dataset), 1)
    int_labels = [round(l) for l in all_labels]
    preds      = [1 if p >= 0.5 else 0 for p in all_probs]
    acc = accuracy_score(int_labels, preds)
    f1  = f1_score(int_labels, preds, zero_division=0)
    try:
        auc = roc_auc_score(int_labels, all_probs)
    except Exception:
        auc = 0.5
    return avg_loss, acc, auc, f1


# ─────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────
def main(args):
    seed_everything()
    device = "cuda" if torch.cuda.is_available() else "cpu"
    print(f"\n[Train] Device: {device}  |  Data: {args.data_dir}")
    if device == "cpu":
        print("[WARN] No GPU detected — training will be slow. Consider using Google Colab with GPU.")

    # ── Datasets ──────────────────────────────────────────────
    train_ds = ImageFakeDataset(args.data_dir, "train", augment=True)
    val_ds   = ImageFakeDataset(args.data_dir, "val",   augment=False)

    if len(train_ds) == 0:
        print("\n[WARN] No training images found! Expected:")
        print("  datasets/images/train/real/*.jpg")
        print("  datasets/images/train/fake/*.jpg")
        return

    # Balanced sampler — ensures each batch has ~50/50 real/fake
    weights  = train_ds.get_weights()
    sampler  = WeightedRandomSampler(weights, num_samples=len(weights), replacement=True)

    train_loader = DataLoader(
        train_ds, batch_size=args.batch_size, sampler=sampler,
        num_workers=0, pin_memory=(device == "cuda"), drop_last=True,
    )
    val_loader = DataLoader(
        val_ds, batch_size=args.batch_size, shuffle=False, num_workers=0,
    )

    # ── Model ─────────────────────────────────────────────────
    model = EfficientNetV2Detector(pretrained=True)

    # Phase 1: freeze backbone, only train classifier head
    for param in model.backbone.parameters():
        param.requires_grad = False
    model.to(device)

    # ── Loss ──────────────────────────────────────────────────
    real_n     = sum(1 for _, l in train_ds.samples if l == 0)
    fake_n     = sum(1 for _, l in train_ds.samples if l == 1)
    pos_weight = torch.tensor([real_n / max(fake_n, 1)], device=device)
    criterion  = LabelSmoothingBCE(smoothing=0.1, pos_weight=pos_weight)

    # ── Phase 1 optimizer + scheduler ─────────────────────────
    accum_steps = max(1, 32 // args.batch_size)   # effective batch ≈ 32
    total_steps_p1 = (len(train_loader) // accum_steps) * args.epochs

    optimizer = AdamW(
        filter(lambda p: p.requires_grad, model.parameters()),
        lr=args.lr, weight_decay=2e-4,
    )
    scheduler = OneCycleLR(
        optimizer, max_lr=args.lr,
        total_steps=total_steps_p1,
        pct_start=0.1,            # 10% warmup
        anneal_strategy="cos",
    )

    # ── Training loop – Phase 1 ───────────────────────────────
    save_dir = Path(args.save_dir)
    save_dir.mkdir(parents=True, exist_ok=True)
    best_auc, no_improve = 0.0, 0

    print("\n" + "=" * 70)
    print("  DeepShield — Improved Image Detector Training (EfficientNetV2-S)")
    print("=" * 70)
    print(f"  Phase 1: {args.epochs} epochs | frozen backbone | LR={args.lr}")
    print(f"  Grad accum steps: {accum_steps} -> effective batch ~{args.batch_size * accum_steps}")
    print(f"  Early stopping patience: {args.patience}")
    print("=" * 70)

    for epoch in range(1, args.epochs + 1):
        t0 = time.time()
        tr_loss, tr_acc, tr_auc, tr_f1 = run_epoch(
            model, train_loader, criterion, optimizer, device,
            training=True, scheduler=scheduler, accumulation_steps=accum_steps
        )
        va_loss, va_acc, va_auc, va_f1 = run_epoch(
            model, val_loader, criterion, None, device,
            training=False
        )
        elapsed = time.time() - t0

        print(
            f"  [P1] Ep {epoch:03d}/{args.epochs} "
            f"| train loss={tr_loss:.4f} acc={tr_acc:.3f} AUC={tr_auc:.3f} F1={tr_f1:.3f}"
            f" | val loss={va_loss:.4f} acc={va_acc:.3f} AUC={va_auc:.3f} F1={va_f1:.3f}"
            f" | {elapsed:.1f}s"
        )

        if va_auc > best_auc:
            best_auc  = va_auc
            no_improve = 0
            ckpt = save_dir / "image_model_best.pth"
            torch.save(model.state_dict(), ckpt)
            print(f"  [BEST] Model saved -> {ckpt}  (AUC={best_auc:.4f})")
        else:
            no_improve += 1
            if no_improve >= args.patience:
                print(f"\n  [STOP] Early stopping at epoch {epoch} (no improvement for {args.patience} epochs)")
                break

    # ── Phase 2: Unfreeze + fine-tune ─────────────────────────
    if args.finetune_epochs > 0:
        print(f"\n  Phase 2: Unfreezing all layers for {args.finetune_epochs} epochs...")

        # Reload best checkpoint before fine-tuning
        best_ckpt = save_dir / "image_model_best.pth"
        if best_ckpt.exists():
            model.load_state_dict(torch.load(str(best_ckpt), map_location=device))
            print(f"  Loaded best checkpoint (AUC={best_auc:.4f}) for fine-tuning")

        for param in model.backbone.parameters():
            param.requires_grad = True

        # Layer-wise LR: backbone gets 10x lower LR than head
        backbone_params = list(model.backbone.parameters())
        head_params     = list(model.classifier.parameters())
        param_groups    = [
            {"params": backbone_params, "lr": args.lr * 0.05},
            {"params": head_params,     "lr": args.lr * 0.5},
        ]

        total_steps_p2 = (len(train_loader) // accum_steps) * args.finetune_epochs
        optimizer2 = AdamW(param_groups, weight_decay=2e-4)
        scheduler2 = OneCycleLR(
            optimizer2,
            max_lr=[args.lr * 0.05, args.lr * 0.5],
            total_steps=total_steps_p2,
            pct_start=0.15,
            anneal_strategy="cos",
        )

        no_improve = 0
        for epoch in range(1, args.finetune_epochs + 1):
            t0 = time.time()
            tr_loss, tr_acc, tr_auc, tr_f1 = run_epoch(
                model, train_loader, criterion, optimizer2, device,
                training=True, scheduler=scheduler2, accumulation_steps=accum_steps
            )
            va_loss, va_acc, va_auc, va_f1 = run_epoch(
                model, val_loader, criterion, None, device, training=False
            )
            elapsed = time.time() - t0

            print(
                f"  [P2] Ep {epoch:03d}/{args.finetune_epochs} "
                f"| val loss={va_loss:.4f} acc={va_acc:.3f} AUC={va_auc:.3f} F1={va_f1:.3f}"
                f" | {elapsed:.1f}s"
            )

            if va_auc > best_auc:
                best_auc   = va_auc
                no_improve = 0
                ckpt = save_dir / "image_model_best.pth"
                torch.save(model.state_dict(), ckpt)
                print(f"  [BEST] Model updated -> {ckpt}  (AUC={best_auc:.4f})")
            else:
                no_improve += 1
                if no_improve >= args.patience:
                    print(f"\n  [STOP] Early stopping at fine-tune epoch {epoch}")
                    break

    torch.save(model.state_dict(), save_dir / "image_model_final.pth")
    print(f"\n{'='*70}")
    print(f"  Training Complete!")
    print(f"  Best Val AUC : {best_auc:.4f}")
    print(f"  Best model   : {save_dir / 'image_model_best.pth'}")
    print(f"  Final model  : {save_dir / 'image_model_final.pth'}")
    print(f"{'='*70}")
    print("  NOTE: the Flask app auto-loads 'image_model_best.pth'.")
    print("        Restart Flask after training to use the new weights.\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="DeepShield Improved Image Training")
    parser.add_argument("--data_dir",        type=str,   default="../datasets/images")
    parser.add_argument("--epochs",          type=int,   default=40,
                        help="Phase-1 epochs (frozen backbone)")
    parser.add_argument("--finetune_epochs", type=int,   default=20,
                        help="Phase-2 epochs (full fine-tune)")
    parser.add_argument("--batch_size",      type=int,   default=16)
    parser.add_argument("--lr",              type=float, default=2e-4)
    parser.add_argument("--save_dir",        type=str,   default="models")
    parser.add_argument("--patience",        type=int,   default=10,
                        help="Early stopping patience (epochs without val AUC improvement)")
    args = parser.parse_args()
    main(args)