scripts/train_liveness_detector.py

#!/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)