import argparse from pathlib import Path import torch import torch.nn as nn from torch.utils.data import DataLoader from torchvision import transforms from sklearn.metrics import ( accuracy_score, precision_score, recall_score, f1_score, roc_auc_score, confusion_matrix, ) from video_data import VideoDataset from video_model import ResNetLSTM def build_transforms(): return transforms.Compose( [ transforms.Resize((224, 224)), transforms.ToTensor(), transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ] ) def accuracy_from_logits(logits, labels): preds = torch.argmax(logits, dim=1) return (preds == labels).float().mean().item() def compute_classification_metrics(y_true, y_pred, y_prob_fake): cm = confusion_matrix(y_true, y_pred, labels=[0, 1]) tn, fp, fn, tp = cm.ravel() metrics = { "accuracy": accuracy_score(y_true, y_pred), "precision_real": precision_score(y_true, y_pred, pos_label=0, zero_division=0), "precision_fake": precision_score(y_true, y_pred, pos_label=1, zero_division=0), "recall_real": recall_score(y_true, y_pred, pos_label=0, zero_division=0), "recall_fake": recall_score(y_true, y_pred, pos_label=1, zero_division=0), "f1_real": f1_score(y_true, y_pred, pos_label=0, zero_division=0), "f1_fake": f1_score(y_true, y_pred, pos_label=1, zero_division=0), "f1_macro": f1_score(y_true, y_pred, average="macro", zero_division=0), "f1_weighted": f1_score(y_true, y_pred, average="weighted", zero_division=0), "tn": int(tn), "fp": int(fp), "fn": int(fn), "tp": int(tp), "specificity": tn / (tn + fp) if (tn + fp) > 0 else 0.0, "sensitivity": tp / (tp + fn) if (tp + fn) > 0 else 0.0, } try: metrics["auc_roc"] = roc_auc_score(y_true, y_prob_fake) except Exception: metrics["auc_roc"] = float("nan") return metrics def train_one_epoch(model, loader, optimizer, criterion, device, frame_loss_weight): model.train() total_loss = 0.0 total_acc = 0.0 for videos, labels in loader: videos = videos.to(device) labels = labels.to(device) frame_logits, video_logits = model(videos) loss_video = criterion(video_logits, labels) b, t, c = frame_logits.shape frame_targets = labels.repeat_interleave(t) loss_frame = criterion(frame_logits.view(b * t, c), frame_targets) loss = loss_video + frame_loss_weight * loss_frame optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() total_acc += accuracy_from_logits(video_logits, labels) return total_loss / max(1, len(loader)), total_acc / max(1, len(loader)) def evaluate(model, loader, criterion, device): model.eval() total_loss = 0.0 total_acc = 0.0 all_labels = [] all_preds = [] all_prob_fake = [] with torch.no_grad(): for videos, labels in loader: videos = videos.to(device) labels = labels.to(device) _, video_logits = model(videos) loss = criterion(video_logits, labels) probs = torch.softmax(video_logits, dim=1) preds = torch.argmax(video_logits, dim=1) total_loss += loss.item() total_acc += accuracy_from_logits(video_logits, labels) all_labels.extend(labels.cpu().tolist()) all_preds.extend(preds.cpu().tolist()) all_prob_fake.extend(probs[:, 1].cpu().tolist()) metrics = compute_classification_metrics(all_labels, all_preds, all_prob_fake) return total_loss / max(1, len(loader)), total_acc / max(1, len(loader)), metrics def main(): parser = argparse.ArgumentParser() parser.add_argument("--dataset", type=str, default="SDFVD/SDFVD") parser.add_argument("--out", type=str, default="video_xception_lstm.pt") parser.add_argument("--epochs", type=int, default=50) parser.add_argument("--resume", type=str, default=None) parser.add_argument("--override_config", action="store_true") parser.add_argument("--batch_size", type=int, default=2) parser.add_argument("--lr", type=float, default=1e-4) parser.add_argument("--frames_per_video", type=int, default=16) parser.add_argument("--frame_stride", type=int, default=4) parser.add_argument("--max_videos_per_class", type=int, default=None) parser.add_argument("--val_split", type=float, default=0.2) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--num_workers", type=int, default=0) parser.add_argument("--no_face_detection", action="store_true") parser.add_argument("--frame_loss_weight", type=float, default=0.3) parser.add_argument("--temporal_pool", type=str, choices=["mean", "last"], default="mean") parser.add_argument("--backbone", type=str, choices=["resnet50", "xception"], default="xception") parser.add_argument("--hidden_size", type=int, default=256) parser.add_argument("--num_layers", type=int, default=1) parser.add_argument("--no_bidirectional", action="store_false", dest="bidirectional") parser.add_argument("--no_pretrained", action="store_true") parser.add_argument("--early_stop_patience", type=int, default=8) parser.add_argument("--early_stop_min_delta", type=float, default=1e-4) parser.add_argument("--cpu", action="store_true") args = parser.parse_args() device = torch.device("cpu" if args.cpu or not torch.cuda.is_available() else "cuda") resume_path = Path(args.resume) if args.resume else None checkpoint = None if resume_path and resume_path.exists(): checkpoint = torch.load(resume_path, map_location="cpu") if not args.override_config and isinstance(checkpoint, dict) and "config" in checkpoint: cfg = checkpoint["config"] args.hidden_size = cfg.get("hidden_size", args.hidden_size) args.num_layers = cfg.get("num_layers", args.num_layers) args.bidirectional = cfg.get("bidirectional", args.bidirectional) args.temporal_pool = cfg.get("temporal_pool", args.temporal_pool) args.backbone = cfg.get("backbone", args.backbone) args.no_pretrained = not cfg.get("pretrained", not args.no_pretrained) args.frames_per_video = cfg.get("frames_per_video", args.frames_per_video) args.frame_stride = cfg.get("frame_stride", args.frame_stride) args.no_face_detection = not cfg.get("face_detection", not args.no_face_detection) transform = build_transforms() train_ds = VideoDataset( root_dir=args.dataset, split="train", val_split=args.val_split, seed=args.seed, max_videos_per_class=args.max_videos_per_class, frames_per_video=args.frames_per_video, frame_stride=args.frame_stride, face_detection=not args.no_face_detection, transform=transform, ) val_ds = VideoDataset( root_dir=args.dataset, split="val", val_split=args.val_split, seed=args.seed, max_videos_per_class=args.max_videos_per_class, frames_per_video=args.frames_per_video, frame_stride=args.frame_stride, face_detection=not args.no_face_detection, transform=transform, ) train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers) val_loader = DataLoader(val_ds, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers) model = ResNetLSTM( hidden_size=args.hidden_size, num_layers=args.num_layers, bidirectional=args.bidirectional, temporal_pool=args.temporal_pool, pretrained=not args.no_pretrained, backbone_name=args.backbone, ) model.to(device) if checkpoint and "model_state" in checkpoint: model.load_state_dict(checkpoint["model_state"], strict=True) print(f"Resumed model weights from {resume_path}") criterion = nn.CrossEntropyLoss() optimizer = torch.optim.AdamW(model.parameters(), lr=args.lr) best_val_acc = 0.0 start_epoch = 0 if checkpoint: best_val_acc = float(checkpoint.get("best_val_acc", best_val_acc)) start_epoch = int(checkpoint.get("epoch", start_epoch)) out_path = Path(args.out) out_path.parent.mkdir(parents=True, exist_ok=True) total_epochs = start_epoch + args.epochs no_improve_epochs = 0 for epoch in range(args.epochs): train_loss, train_acc = train_one_epoch( model, train_loader, optimizer, criterion, device, args.frame_loss_weight ) val_loss, val_acc, val_metrics = evaluate(model, val_loader, criterion, device) print( f"Epoch {start_epoch + epoch + 1}/{total_epochs} | " f"Train Loss: {train_loss:.4f} Acc: {train_acc:.3f} | " f"Val Loss: {val_loss:.4f} Acc: {val_acc:.3f}" ) print( " Accuracy | " f"Train Accuracy: {train_acc:.4f} | " f"Validation Accuracy: {val_acc:.4f}" ) print( " Val Metrics | " f"AUC: {val_metrics['auc_roc']:.4f} | " f"F1(macro): {val_metrics['f1_macro']:.4f} | " f"F1(weighted): {val_metrics['f1_weighted']:.4f}" ) print( " Real Class | " f"Precision: {val_metrics['precision_real']:.4f} | " f"Recall: {val_metrics['recall_real']:.4f} | " f"F1: {val_metrics['f1_real']:.4f}" ) print( " Fake Class | " f"Precision: {val_metrics['precision_fake']:.4f} | " f"Recall: {val_metrics['recall_fake']:.4f} | " f"F1: {val_metrics['f1_fake']:.4f}" ) print( " Derived | " f"Sensitivity(TPR): {val_metrics['sensitivity']:.4f} | " f"Specificity(TNR): {val_metrics['specificity']:.4f}" ) print( " Confusion | " f"TN: {val_metrics['tn']} FP: {val_metrics['fp']} " f"FN: {val_metrics['fn']} TP: {val_metrics['tp']}" ) if val_acc >= best_val_acc + args.early_stop_min_delta: best_val_acc = val_acc no_improve_epochs = 0 torch.save( { "model_state": model.state_dict(), "epoch": start_epoch + epoch + 1, "best_val_acc": best_val_acc, "config": { "hidden_size": args.hidden_size, "num_layers": args.num_layers, "bidirectional": args.bidirectional, "temporal_pool": args.temporal_pool, "backbone": args.backbone, "pretrained": not args.no_pretrained, "frames_per_video": args.frames_per_video, "frame_stride": args.frame_stride, "face_detection": not args.no_face_detection, }, }, out_path, ) print(f"Saved best model to {out_path}") else: no_improve_epochs += 1 if no_improve_epochs >= args.early_stop_patience: print( "Early stopping triggered: " f"no validation improvement for {args.early_stop_patience} epoch(s)." ) break if __name__ == "__main__": main()