""" Training Script — IEEE Research Experiments Trains HybridDeepfakeDetector on FaceForensics++ or Celeb-DF v2 Usage: python train.py --data_dir /path/to/dataset --epochs 30 --batch_size 32 """ import argparse, time, os try: from tqdm import tqdm except ImportError: tqdm = None from pathlib import Path import numpy as np import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler import torchvision.transforms as T from PIL import Image from sklearn.metrics import roc_auc_score from model import HybridDeepfakeDetector # ── Dataset ───────────────────────────────────────────────────────── class DeepfakeDataset(Dataset): """ Expects directory layout: data_dir/ real/ ← real face crops (PNG/JPG) fake/ ← fake face crops """ TRAIN_TF = T.Compose([ T.Resize((224, 224)), T.RandomHorizontalFlip(), T.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.1), T.RandomRotation(10), T.GaussianBlur(kernel_size=3, sigma=(0.1, 1.5)), T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) VAL_TF = T.Compose([ T.Resize((224, 224)), T.ToTensor(), T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), ]) def __init__(self, data_dir: str, split: str = "train"): self.tf = self.TRAIN_TF if split == "train" else self.VAL_TF self.samples = [] for label, folder in [(0, "real"), (1, "fake")]: p = Path(data_dir) / folder if p.exists(): for img in p.rglob("*.jpg"): self.samples.append((str(img), label)) for img in p.rglob("*.png"): self.samples.append((str(img), label)) if not self.samples: raise ValueError(f"No images found in {data_dir}. " f"Ensure real/ and fake/ sub-directories exist under {data_dir}.") def __len__(self): return len(self.samples) def __getitem__(self, idx): path, label = self.samples[idx] img = Image.open(path).convert("RGB") return self.tf(img), torch.tensor(label, dtype=torch.float32) def get_sampler(dataset: DeepfakeDataset) -> WeightedRandomSampler: labels = [s[1] for s in dataset.samples] counts = [labels.count(0), labels.count(1)] weights = [1.0 / counts[l] for l in labels] return WeightedRandomSampler(weights, len(weights)) # ── Training loop ─────────────────────────────────────────────────── def train_epoch(model, loader, optimizer, criterion, device, scaler): model.train() total_loss, n = 0.0, 0 total_batches = len(loader) iterator = tqdm(loader, desc=" Training", unit="batch", bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}] loss={postfix}") \ if tqdm else loader for batch_idx, (imgs, labels) in enumerate(iterator): imgs, labels = imgs.to(device), labels.to(device) optimizer.zero_grad() with torch.amp.autocast(device_type=device, enabled=scaler is not None): logits = model(imgs).squeeze(1) loss = criterion(logits, labels) if scaler: scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() else: loss.backward() optimizer.step() total_loss += loss.item() * imgs.size(0) n += imgs.size(0) # Update tqdm with current avg loss if tqdm and hasattr(iterator, 'set_postfix_str'): iterator.set_postfix_str(f"{total_loss/n:.4f}") elif not tqdm and (batch_idx % 50 == 0 or batch_idx == total_batches - 1): pct = (batch_idx + 1) / total_batches * 100 print(f" Batch {batch_idx+1}/{total_batches} ({pct:.0f}%) avg_loss={total_loss/n:.4f}", flush=True) return total_loss / n @torch.no_grad() def evaluate(model, loader, device): model.eval() all_probs, all_labels = [], [] for imgs, labels in loader: imgs = imgs.to(device) probs = torch.sigmoid(model(imgs).squeeze(1)).cpu().numpy() all_probs.extend(probs) all_labels.extend(labels.numpy()) auc = roc_auc_score(all_labels, all_probs) preds = [1 if p >= 0.5 else 0 for p in all_probs] acc = np.mean(np.array(preds) == np.array(all_labels)) return {"auc": auc, "acc": acc} # ── Main ───────────────────────────────────────────────────────────── def main(): parser = argparse.ArgumentParser() parser.add_argument("--data_dir", default="../datasets/video_crops") parser.add_argument("--epochs", type=int, default=30) parser.add_argument("--batch_size", type=int, default=32) parser.add_argument("--lr", type=float, default=1e-4) parser.add_argument("--save_dir", default="models") parser.add_argument("--device", default=None) args = parser.parse_args() device = args.device or ("cuda" if torch.cuda.is_available() else "cpu") print(f"[Train] Device: {device}") # Datasets # data_dir should be the base data/ folder; train/val subdirs are appended internally train_split_dir = str(Path(args.data_dir) / "train") val_split_dir = str(Path(args.data_dir) / "val") train_ds = DeepfakeDataset(train_split_dir, "train") val_ds = DeepfakeDataset(val_split_dir, "val") sampler = get_sampler(train_ds) # num_workers=0 avoids Windows multiprocessing issues train_dl = DataLoader(train_ds, batch_size=args.batch_size, sampler=sampler, num_workers=0, pin_memory=False, drop_last=True) val_dl = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False, num_workers=0) print(f"[Train] Train samples: {len(train_ds)} | Val samples: {len(val_ds)}") # Model model = HybridDeepfakeDetector(pretrained=True).to(device) # Loss with label smoothing criterion = nn.BCEWithLogitsLoss(label_smoothing=0.1) if hasattr( nn.BCEWithLogitsLoss, "label_smoothing" ) else nn.BCEWithLogitsLoss() optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4) scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs) scaler = torch.cuda.amp.GradScaler() if device == "cuda" else None os.makedirs(args.save_dir, exist_ok=True) best_auc = 0.0 for epoch in range(1, args.epochs + 1): t0 = time.time() loss = train_epoch(model, train_dl, optimizer, criterion, device, scaler) mets = evaluate(model, val_dl, device) scheduler.step() print(f"Epoch {epoch:03d}/{args.epochs} " f"Loss={loss:.4f} AUC={mets['auc']*100:.2f}% " f"ACC={mets['acc']*100:.2f}% " f"[{time.time()-t0:.1f}s]") if mets["auc"] > best_auc: best_auc = mets["auc"] save_path = Path(args.save_dir) / "deepfake_model.pth" torch.save(model.state_dict(), save_path) print(f" ✓ Best model saved → {save_path} (AUC={best_auc*100:.2f}%)") print(f"\n[Done] Best AUC: {best_auc*100:.2f}%") if __name__ == "__main__": main()