#!/usr/bin/env python3 """ train_liveness_detector.py - Train a liveness detection model This script trains a neural network model for liveness detection using a dataset of real faces and presentation attacks. It supports various backbone architectures and training configurations. """ import os import sys import argparse import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import Dataset, DataLoader import numpy as np from PIL import Image import cv2 import json import time import datetime import random from typing import Dict, List, Tuple, Optional, Any from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score, roc_auc_score # Add project root to path sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) # Import MorphGuard modules from src.liveness.liveness_detector import LivenessDetectionModel from models_config import LIVENESS_MODELS # Try to import optional dependencies try: import wandb WANDB_AVAILABLE = True except ImportError: WANDB_AVAILABLE = False try: import albumentations as A ALBUMENTATIONS_AVAILABLE = True except ImportError: ALBUMENTATIONS_AVAILABLE = False class LivenessDataset(Dataset): """Dataset for liveness detection training""" def __init__(self, data_dir: str, split: str = 'train', img_size: int = 224, transform = None, augment: bool = True): """Initialize dataset Args: data_dir: Base data directory split: Dataset split ('train', 'val', or 'test') img_size: Image size for preprocessing transform: Optional custom transform augment: Whether to apply data augmentation (train split only) """ self.data_dir = data_dir self.split = split self.img_size = img_size self.custom_transform = transform self.augment = augment and split == 'train' # Define paths for real and spoof subsets self.real_dir = os.path.join(data_dir, split, 'real') self.spoof_dir = os.path.join(data_dir, split, 'spoof') # Ensure directories exist if not os.path.exists(self.real_dir): os.makedirs(self.real_dir, exist_ok=True) print(f"Created directory: {self.real_dir}") if not os.path.exists(self.spoof_dir): os.makedirs(self.spoof_dir, exist_ok=True) print(f"Created directory: {self.spoof_dir}") # Collect image paths and labels self.real_images = [os.path.join(self.real_dir, f) for f in os.listdir(self.real_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))] self.spoof_images = [os.path.join(self.spoof_dir, f) for f in os.listdir(self.spoof_dir) if f.lower().endswith(('.jpg', '.jpeg', '.png'))] # Combine into a single dataset self.image_paths = self.real_images + self.spoof_images self.labels = [0] * len(self.real_images) + [1] * len(self.spoof_images) # 0: real, 1: spoof # Create pairs of (path, label) and shuffle them self.data = list(zip(self.image_paths, self.labels)) if split == 'train': random.shuffle(self.data) # Print dataset stats print(f"{split} dataset: {len(self.real_images)} real, {len(self.spoof_images)} spoof images") # Set up transforms self._setup_transforms() def _setup_transforms(self): """Set up image transformations""" from torchvision import transforms # Standard normalization values for ImageNet self.mean = [0.485, 0.456, 0.406] self.std = [0.229, 0.224, 0.225] # Basic transformation for all splits self.transform = transforms.Compose([ transforms.Resize((self.img_size, self.img_size)), transforms.ToTensor(), transforms.Normalize(mean=self.mean, std=self.std) ]) # Augmentations for training if self.augment and ALBUMENTATIONS_AVAILABLE: self.aug_transform = A.Compose([ A.HorizontalFlip(p=0.5), A.RandomBrightnessContrast(p=0.5), A.ShiftScaleRotate(shift_limit=0.05, scale_limit=0.05, rotate_limit=10, p=0.5), A.OneOf([ A.MotionBlur(p=0.5), A.GaussianBlur(p=0.5), A.GaussNoise(p=0.5), ], p=0.3), A.OneOf([ A.OpticalDistortion(p=0.3), A.GridDistortion(p=0.3), ], p=0.2), A.CoarseDropout(max_holes=8, max_height=8, max_width=8, p=0.3), ]) def __len__(self): return len(self.data) def __getitem__(self, idx): img_path, label = self.data[idx] # Load and preprocess image try: # Read image img = Image.open(img_path).convert('RGB') # Apply augmentations if needed if self.augment and ALBUMENTATIONS_AVAILABLE: img_np = np.array(img) augmented = self.aug_transform(image=img_np) img = Image.fromarray(augmented['image']) # Apply transform if self.custom_transform: img_tensor = self.custom_transform(img) else: img_tensor = self.transform(img) return img_tensor, label except Exception as e: print(f"Error loading image {img_path}: {e}") # Return a random noise image as fallback img_tensor = torch.randn(3, self.img_size, self.img_size) img_tensor = (img_tensor - img_tensor.min()) / (img_tensor.max() - img_tensor.min()) img_tensor = transforms.Normalize(mean=self.mean, std=self.std)(img_tensor) return img_tensor, label def train_epoch(model, dataloader, criterion, optimizer, device, epoch, log_interval=10): """Train model for one epoch""" model.train() running_loss = 0.0 all_preds = [] all_labels = [] start_time = time.time() for batch_idx, (data, target) in enumerate(dataloader): data, target = data.to(device), target.to(device) # Forward pass optimizer.zero_grad() output = model(data) loss = criterion(output, target) # Backward pass loss.backward() optimizer.step() # Track metrics running_loss += loss.item() pred = output.argmax(dim=1, keepdim=True) all_preds.extend(pred.view(-1).cpu().numpy()) all_labels.extend(target.cpu().numpy()) # Log progress if batch_idx % log_interval == 0: print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(dataloader.dataset)}' f' ({100. * batch_idx / len(dataloader):.0f}%)]\tLoss: {loss.item():.6f}') # Compute final metrics train_loss = running_loss / len(dataloader) train_acc = accuracy_score(all_labels, all_preds) train_metrics = { 'loss': train_loss, 'accuracy': train_acc, 'precision': precision_score(all_labels, all_preds, average='binary'), 'recall': recall_score(all_labels, all_preds, average='binary'), 'f1': f1_score(all_labels, all_preds, average='binary') } # Log to console print(f'\nTrain Epoch: {epoch} completed in {time.time() - start_time:.2f}s') print(f'Train metrics: Loss: {train_metrics["loss"]:.4f}, Accuracy: {train_metrics["accuracy"]:.4f}, ' f'F1: {train_metrics["f1"]:.4f}') return train_metrics def validate(model, dataloader, criterion, device): """Validate model on validation set""" model.eval() val_loss = 0 all_preds = [] all_labels = [] all_probs = [] with torch.no_grad(): for data, target in dataloader: data, target = data.to(device), target.to(device) # Forward pass output = model(data) loss = criterion(output, target) # Track metrics val_loss += loss.item() pred = output.argmax(dim=1, keepdim=True) all_preds.extend(pred.view(-1).cpu().numpy()) all_labels.extend(target.cpu().numpy()) all_probs.extend(output[:, 1].cpu().numpy()) # Prob of spoof class # Compute final metrics val_loss /= len(dataloader) val_metrics = { 'loss': val_loss, 'accuracy': accuracy_score(all_labels, all_preds), 'precision': precision_score(all_labels, all_preds, average='binary'), 'recall': recall_score(all_labels, all_preds, average='binary'), 'f1': f1_score(all_labels, all_preds, average='binary'), 'auc': roc_auc_score(all_labels, all_probs) } # Log to console print(f'\nValidation metrics: Loss: {val_metrics["loss"]:.4f}, Accuracy: {val_metrics["accuracy"]:.4f}, ' f'F1: {val_metrics["f1"]:.4f}, AUC: {val_metrics["auc"]:.4f}') return val_metrics def log_metrics_to_db(model_name, epoch, train_metrics, val_metrics, session_id, batch_size, lr): """Log training and validation metrics to TimescaleDB""" try: import psycopg2 import datetime # Try to get DB settings from config try: import config db_settings = { "dbname": getattr(config, 'DB_NAME', 'morphguard'), "user": getattr(config, 'DB_USER', 'morphguard'), "password": getattr(config, 'DB_PASSWORD', 'morphguard'), "host": getattr(config, 'DB_HOST', 'localhost'), "port": getattr(config, 'DB_PORT', 5432) } except ImportError: # Default settings if config not found db_settings = { "dbname": "morphguard", "user": "morphguard", "password": "morphguard", "host": "localhost", "port": 5432 } # Connect to database conn = psycopg2.connect(**db_settings) cursor = conn.cursor() # Insert training metrics cursor.execute( """ INSERT INTO training_metrics (time, model_name, epoch, loss, accuracy, val_loss, val_accuracy, learning_rate, batch_size, training_session_id) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s) """, ( datetime.datetime.now(), model_name, epoch, train_metrics["loss"], train_metrics["accuracy"], val_metrics["loss"], val_metrics["accuracy"], lr, batch_size, session_id ) ) # Commit and close conn.commit() cursor.close() conn.close() print(f"Logged metrics to TimescaleDB for epoch {epoch}") return True except Exception as e: print(f"Error logging training metrics to DB: {e}") return False def save_checkpoint(model, optimizer, scheduler, epoch, metrics, backbone, save_path): """Save model checkpoint with metadata""" # Ensure directory exists os.makedirs(os.path.dirname(save_path), exist_ok=True) # Create checkpoint checkpoint = { 'epoch': epoch, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'scheduler_state_dict': scheduler.state_dict() if scheduler else None, 'metrics': metrics, 'backbone': backbone, 'timestamp': datetime.datetime.now().isoformat() } # Save checkpoint torch.save(checkpoint, save_path) print(f"Checkpoint saved to {save_path}") def main(args): """Main training function""" print(f"Starting liveness detection model training with backbone: {args.backbone}") # Set random seed for reproducibility random.seed(args.seed) np.random.seed(args.seed) torch.manual_seed(args.seed) if torch.cuda.is_available(): torch.cuda.manual_seed_all(args.seed) # Create unique session ID for this training run session_id = f"liveness_{args.backbone}_{datetime.datetime.now().strftime('%Y%m%d_%H%M%S')}" # Initialize WandB if available if WANDB_AVAILABLE and args.use_wandb: wandb.init( project="morphguard-liveness", name=session_id, config=args.__dict__ ) # Set device device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu") print(f"Using device: {device}") # Create output directory os.makedirs(os.path.dirname(args.save_path), exist_ok=True) # Create datasets and dataloaders train_dataset = LivenessDataset( data_dir=args.data_dir, split='train', img_size=args.img_size, augment=not args.no_augment ) val_dataset = LivenessDataset( data_dir=args.data_dir, split='val', img_size=args.img_size, augment=False ) train_loader = DataLoader( train_dataset, batch_size=args.batch_size, shuffle=True, num_workers=args.workers, pin_memory=True ) val_loader = DataLoader( val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=args.workers, pin_memory=True ) # Initialize model model = LivenessDetectionModel( backbone=args.backbone, pretrained=not args.no_pretrained, num_classes=2 ).to(device) # Define loss function and optimizer criterion = nn.CrossEntropyLoss() # Use different learning rates for backbone and classifier if transfer learning if not args.no_pretrained and args.finetune: # Parameters of newly constructed modules have requires_grad=True by default backbone_params = list(model.features.parameters()) classifier_params = list(model.classifier.parameters()) optimizer = optim.Adam([ {'params': backbone_params, 'lr': args.lr * 0.1}, {'params': classifier_params, 'lr': args.lr} ], weight_decay=args.weight_decay) print("Using transfer learning with different learning rates") else: optimizer = optim.Adam(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) # Learning rate scheduler scheduler = optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode='min', factor=0.5, patience=5, verbose=True ) # Training loop best_val_acc = 0.0 best_val_loss = float('inf') best_epoch = 0 # Create stats dictionary to track metrics stats = { 'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': [], 'val_f1': [], 'epochs': [] } for epoch in range(1, args.epochs + 1): # Update learning rate scheduler current_lr = optimizer.param_groups[0]['lr'] print(f"\nEpoch {epoch}/{args.epochs}, Learning rate: {current_lr}") # Train and validate train_metrics = train_epoch( model, train_loader, criterion, optimizer, device, epoch, args.log_interval ) val_metrics = validate(model, val_loader, criterion, device) # Update learning rate scheduler based on validation loss scheduler.step(val_metrics['loss']) # Log metrics stats['train_loss'].append(train_metrics['loss']) stats['train_acc'].append(train_metrics['accuracy']) stats['val_loss'].append(val_metrics['loss']) stats['val_acc'].append(val_metrics['accuracy']) stats['val_f1'].append(val_metrics['f1']) stats['epochs'].append(epoch) # Log to database log_metrics_to_db( model_name=f"liveness_{args.backbone}", epoch=epoch, train_metrics=train_metrics, val_metrics=val_metrics, session_id=session_id, batch_size=args.batch_size, lr=current_lr ) # Log to WandB if available if WANDB_AVAILABLE and args.use_wandb: wandb.log({ "epoch": epoch, "train_loss": train_metrics['loss'], "train_acc": train_metrics['accuracy'], "train_f1": train_metrics['f1'], "val_loss": val_metrics['loss'], "val_acc": val_metrics['accuracy'], "val_f1": val_metrics['f1'], "val_auc": val_metrics['auc'], "learning_rate": current_lr }) # Save checkpoint for best validation accuracy and loss if val_metrics['accuracy'] > best_val_acc: best_val_acc = val_metrics['accuracy'] best_epoch = epoch # Save best model save_path = args.save_path.replace('.pth', f'_best.pth') save_checkpoint( model, optimizer, scheduler, epoch, val_metrics, args.backbone, save_path ) print(f"New best model saved with validation accuracy: {best_val_acc:.4f}") # Save last checkpoint if epoch == args.epochs or epoch % args.save_interval == 0: save_path = args.save_path.replace('.pth', f'_epoch{epoch}.pth') save_checkpoint( model, optimizer, scheduler, epoch, val_metrics, args.backbone, save_path ) # Save final model save_checkpoint( model, optimizer, scheduler, args.epochs, val_metrics, args.backbone, args.save_path ) # Save training stats to JSON stats_path = os.path.join(os.path.dirname(args.save_path), 'liveness_training_stats.json') with open(stats_path, 'w') as f: json.dump(stats, f, indent=2) print(f"\nTraining completed. Best validation accuracy: {best_val_acc:.4f} at epoch {best_epoch}") print(f"Final model saved to {args.save_path}") print(f"Training stats saved to {stats_path}") if WANDB_AVAILABLE and args.use_wandb: wandb.finish() if __name__ == '__main__': parser = argparse.ArgumentParser(description='Train liveness detection model') # Dataset parameters parser.add_argument('--data-dir', type=str, default='./data', help='Data directory (default: ./data)') parser.add_argument('--img-size', type=int, default=224, help='Input image size (default: 224)') # Model parameters parser.add_argument('--backbone', type=str, default='efficientnet_b0', choices=LIVENESS_MODELS, help=f'Backbone architecture (default: efficientnet_b0)') parser.add_argument('--no-pretrained', action='store_true', help='Do not use pre-trained weights') parser.add_argument('--finetune', action='store_true', help='Use fine-tuning (different learning rates for backbone and classifier)') # Training parameters parser.add_argument('--batch-size', type=int, default=32, help='Batch size for training (default: 32)') parser.add_argument('--epochs', type=int, default=50, help='Number of epochs to train (default: 50)') parser.add_argument('--lr', type=float, default=0.001, help='Learning rate (default: 0.001)') parser.add_argument('--weight-decay', type=float, default=1e-5, help='Weight decay (default: 1e-5)') parser.add_argument('--no-augment', action='store_true', help='Disable data augmentation') parser.add_argument('--workers', type=int, default=4, help='Number of worker threads for data loading (default: 4)') parser.add_argument('--seed', type=int, default=42, help='Random seed (default: 42)') parser.add_argument('--save-path', type=str, default='./models/liveness/liveness_model.pth', help='Path to save the trained model (default: ./models/liveness/liveness_model.pth)') parser.add_argument('--save-interval', type=int, default=10, help='Interval to save model checkpoints (default: 10)') parser.add_argument('--log-interval', type=int, default=10, help='Interval to log training progress (default: 10)') # CUDA parameters parser.add_argument('--no-cuda', action='store_true', help='Disable CUDA training') # Logging parameters parser.add_argument('--use-wandb', action='store_true', help='Log metrics with Weights & Biases') args = parser.parse_args() main(args)