src/train.py

"""
Training script for AI Image Detection
"""

import os
import json
import yaml
import argparse
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 DataLoader
from torch.optim.lr_scheduler import CosineAnnealingLR, StepLR
from torch.utils.tensorboard import SummaryWriter
from tqdm import tqdm

from src.dataset import AIImageDataset, get_transforms
from src.model import AIImageClassifier


def load_config(config_path='config.yaml'):
    """Load configuration from YAML file"""
    with open(config_path, 'r') as f:
        config = yaml.safe_load(f)
    return config


def setup_directories(config):
    """Create necessary directories"""
    os.makedirs(config['output']['model_save_path'], exist_ok=True)
    os.makedirs(config['output']['checkpoint_path'], exist_ok=True)
    os.makedirs(config['output']['results_path'], exist_ok=True)
    os.makedirs(f"{config['output']['results_path']}/logs", exist_ok=True)


def train_epoch(model, train_loader, criterion, optimizer, device, epoch, total_epochs):
    """Train for one epoch"""
    model.train()
    running_loss = 0.0
    correct = 0
    total = 0
    
    pbar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{total_epochs}")
    
    for images, labels in pbar:
        images = images.to(device)
        labels = labels.to(device)
        
        # Forward pass
        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, labels)
        
        # Backward pass
        loss.backward()
        optimizer.step()
        
        # Statistics
        running_loss += loss.item()
        _, predicted = torch.max(outputs.data, 1)
        total += labels.size(0)
        correct += (predicted == labels).sum().item()
        
        pbar.set_postfix({
            'loss': running_loss / (pbar.n + 1),
            'acc': 100 * correct / total
        })
    
    epoch_loss = running_loss / len(train_loader)
    epoch_acc = 100 * correct / total
    
    return epoch_loss, epoch_acc


def validate(model, val_loader, criterion, device):
    """Validate model on validation set"""
    model.eval()
    running_loss = 0.0
    correct = 0
    total = 0
    
    with torch.no_grad():
        for images, labels in tqdm(val_loader, desc="Validating"):
            images = images.to(device)
            labels = labels.to(device)
            
            outputs = model(images)
            loss = criterion(outputs, labels)
            
            running_loss += loss.item()
            _, predicted = torch.max(outputs.data, 1)
            total += labels.size(0)
            correct += (predicted == labels).sum().item()
    
    epoch_loss = running_loss / len(val_loader)
    epoch_acc = 100 * correct / total
    
    return epoch_loss, epoch_acc


def train(config=None, config_path='config.yaml', resume=None):
    """Main training function"""
    
    if config is None:
        config = load_config(config_path)
    
    # Coerce types for common numeric config values to avoid YAML parsing issues
    try:
        config['training']['batch_size'] = int(config['training']['batch_size'])
    except Exception:
        config['training']['batch_size'] = int(float(config['training']['batch_size']))

    config['training']['num_epochs'] = int(config['training']['num_epochs'])
    config['training']['learning_rate'] = float(config['training']['learning_rate'])
    config['training']['weight_decay'] = float(config['training']['weight_decay'])

    # Ensure num_workers and image size are integers
    config['num_workers'] = int(config.get('num_workers', 0))
    config['preprocessing']['image_size'] = int(config['preprocessing']['image_size'])

    # Normalize boolean-like strings for pretrained flag
    if isinstance(config['model'].get('pretrained'), str):
        config['model']['pretrained'] = config['model']['pretrained'].lower() in ('true', '1', 'yes')
    
    print("=== AI Image Detection - Training ===")
    print(f"Config: {config_path}")
    
    # Setup
    setup_directories(config)
    device = torch.device('cuda' if torch.cuda.is_available() and config['device'] == 'cuda' else 'cpu')
    print(f"Device: {device}")
    
    # Create model
    model = AIImageClassifier(
        model_name=config['model']['name'],
        num_classes=config['model']['num_classes'],
        pretrained=config['model']['pretrained'],
        dropout=config['model']['dropout']
    )
    model = model.to(device)
    print(f"Model: {config['model']['name']}")
    
    # Load data
    train_transform = get_transforms(
        image_size=config['preprocessing']['image_size'],
        mode='train',
        normalize_mean=config['preprocessing']['normalize_mean'],
        normalize_std=config['preprocessing']['normalize_std']
    )
    val_transform = get_transforms(
        image_size=config['preprocessing']['image_size'],
        mode='val',
        normalize_mean=config['preprocessing']['normalize_mean'],
        normalize_std=config['preprocessing']['normalize_std']
    )
    
    train_dataset = AIImageDataset(config['data']['train_path'], transform=train_transform)
    val_dataset = AIImageDataset(config['data']['val_path'], transform=val_transform)
    
    train_loader = DataLoader(
        train_dataset,
        batch_size=config['training']['batch_size'],
        shuffle=True,
        num_workers=config['num_workers']
    )
    val_loader = DataLoader(
        val_dataset,
        batch_size=config['training']['batch_size'],
        shuffle=False,
        num_workers=config['num_workers']
    )
    
    # Loss and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters(), lr=config['training']['learning_rate'], 
                          weight_decay=config['training']['weight_decay'])
    
    if config['training']['scheduler'] == 'cosine':
        scheduler = CosineAnnealingLR(optimizer, T_max=config['training']['num_epochs'])
    else:
        scheduler = StepLR(optimizer, step_size=10, gamma=0.1)
    
    # TensorBoard
    writer = SummaryWriter(f"{config['output']['results_path']}/logs")

    # Load existing training history if present (useful when resuming)
    history_path = f"{config['output']['results_path']}/training_history.json"
    if os.path.exists(history_path):
        try:
            with open(history_path, 'r') as f:
                history = json.load(f)
        except Exception:
            history = {'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': []}
    else:
        history = {'train_loss': [], 'train_acc': [], 'val_loss': [], 'val_acc': []}

    # Determine starting epoch (supports resume via checkpoint or existing history)
    start_epoch = len(history.get('train_loss', []))

    # Training bookkeeping
    best_val_loss = float('inf')
    patience_counter = 0

    # If resume checkpoint provided, attempt to load model/optimizer/scheduler state
    if resume:
        if os.path.exists(resume):
            try:
                ckpt = torch.load(resume, map_location=device)
                # If full checkpoint dict with keys
                if isinstance(ckpt, dict):
                    if 'model_state_dict' in ckpt:
                        model.load_state_dict(ckpt['model_state_dict'])
                    else:
                        try:
                            model.load_state_dict(ckpt)
                        except Exception:
                            pass
                    if 'optimizer_state_dict' in ckpt:
                        try:
                            optimizer.load_state_dict(ckpt['optimizer_state_dict'])
                        except Exception:
                            pass
                    if 'scheduler_state_dict' in ckpt and ckpt['scheduler_state_dict'] is not None:
                        try:
                            scheduler.load_state_dict(ckpt['scheduler_state_dict'])
                        except Exception:
                            pass
                    if 'best_val_loss' in ckpt:
                        best_val_loss = ckpt.get('best_val_loss', best_val_loss)
                    # If checkpoint contains epoch, resume from next
                    if 'epoch' in ckpt:
                        start_epoch = ckpt.get('epoch', 0) + 1
                else:
                    # assume it's a state_dict
                    try:
                        model.load_state_dict(ckpt)
                    except Exception:
                        pass
                print(f"Resuming training from checkpoint: {resume}, start_epoch={start_epoch}")
            except Exception as e:
                print(f"Warning: failed to load resume checkpoint {resume}: {e}")
        else:
            print(f"Warning: resume checkpoint {resume} not found")
    

    for epoch in range(start_epoch, config['training']['num_epochs']):
        train_loss, train_acc = train_epoch(
            model, train_loader, criterion, optimizer, device, epoch, config['training']['num_epochs']
        )
        val_loss, val_acc = validate(model, val_loader, criterion, device)
        scheduler.step()
        
        # Logging
        print(f"Epoch {epoch+1}/{config['training']['num_epochs']}")
        print(f"  Train Loss: {train_loss:.4f}, Acc: {train_acc:.2f}%")
        print(f"  Val Loss: {val_loss:.4f}, Acc: {val_acc:.2f}%")
        
        writer.add_scalar('Loss/train', train_loss, epoch)
        writer.add_scalar('Loss/val', val_loss, epoch)
        writer.add_scalar('Accuracy/train', train_acc, epoch)
        writer.add_scalar('Accuracy/val', val_acc, epoch)
        
        history['train_loss'].append(train_loss)
        history['train_acc'].append(train_acc)
        history['val_loss'].append(val_loss)
        history['val_acc'].append(val_acc)
        
        # Save best model (save full checkpoint including optimizer/scheduler state)
        if val_loss < best_val_loss:
            best_val_loss = val_loss
            patience_counter = 0
            checkpoint_path = f"{config['output']['checkpoint_path']}/best_model.pth"
            try:
                torch.save({
                    'epoch': epoch,
                    'model_state_dict': model.state_dict(),
                    'optimizer_state_dict': optimizer.state_dict(),
                    'scheduler_state_dict': scheduler.state_dict() if 'scheduler' in locals() else None,
                    'best_val_loss': best_val_loss,
                }, checkpoint_path)
            except Exception:
                # fallback to saving model weights only
                torch.save(model.state_dict(), checkpoint_path)
            print(f"  Saved best model to {checkpoint_path}")
        else:
            patience_counter += 1
        
        # Early stopping
        if patience_counter >= config['training']['patience']:
            print(f"Early stopping at epoch {epoch+1}")
            break
    
    writer.close()
    
    # Save training history
    history_path = f"{config['output']['results_path']}/training_history.json"
    with open(history_path, 'w') as f:
        json.dump(history, f, indent=2)
    print(f"Training history saved to {history_path}")
    
    # Save final model
    final_model_path = f"{config['output']['model_save_path']}/final_model.pth"
    torch.save(model.state_dict(), final_model_path)
    print(f"Final model saved to {final_model_path}")
    
    print("=== Training Complete ===")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Train AI Image Detection Model')
    parser.add_argument('--config', type=str, default='config.yaml', help='Path to config file')
    parser.add_argument('--resume', type=str, default=None, help='Path to checkpoint to resume from')
    args = parser.parse_args()
    
    try:
        train(config_path=args.config, resume=args.resume)
    except Exception as e:
        import traceback, sys
        traceback.print_exc()
        # Write full traceback to a file for post-mortem
        try:
            with open('training_error.log', 'w', encoding='utf-8') as f:
                traceback.print_exc(file=f)
        except Exception:
            pass
        sys.exit(1)