""" 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)