src/evaluate.py

"""
Evaluation script for AI Image Detection
"""

import os
import json
import yaml
import numpy as np
from pathlib import Path

import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from sklearn.metrics import (
    confusion_matrix, classification_report, roc_auc_score, roc_curve,
    precision_recall_curve, f1_score, accuracy_score
)
import matplotlib.pyplot as plt
import seaborn as sns
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 evaluate_model(model, test_loader, device, class_names=['Natural', 'Synthetic']):
    """
    Evaluate model on test set and compute metrics
    
    Returns:
        dict: Dictionary containing all metrics
    """
    model.eval()
    all_preds = []
    all_labels = []
    all_probs = []
    
    with torch.no_grad():
        for images, labels in tqdm(test_loader, desc="Evaluating"):
            images = images.to(device)
            labels = labels.to(device)
            
            outputs = model(images)
            probs = torch.softmax(outputs, dim=1)
            _, predicted = torch.max(outputs.data, 1)
            
            all_preds.extend(predicted.cpu().numpy())
            all_labels.extend(labels.cpu().numpy())
            all_probs.extend(probs.cpu().numpy())
    
    all_preds = np.array(all_preds)
    all_labels = np.array(all_labels)
    all_probs = np.array(all_probs)
    
    # Calculate metrics
    metrics = {
        'accuracy': accuracy_score(all_labels, all_preds),
        'f1_score': f1_score(all_labels, all_preds),
        'roc_auc': roc_auc_score(all_labels, all_probs[:, 1]),
    }
    
    # Confusion matrix
    cm = confusion_matrix(all_labels, all_preds)
    metrics['confusion_matrix'] = cm.tolist()
    
    # Classification report
    report = classification_report(all_labels, all_preds, target_names=class_names, output_dict=True)
    metrics['classification_report'] = report
    
    # ROC curve data
    fpr, tpr, _ = roc_curve(all_labels, all_probs[:, 1])
    metrics['roc_curve'] = {'fpr': fpr.tolist(), 'tpr': tpr.tolist()}
    
    # Precision-Recall curve
    precision, recall, _ = precision_recall_curve(all_labels, all_probs[:, 1])
    metrics['pr_curve'] = {'precision': precision.tolist(), 'recall': recall.tolist()}
    
    return metrics, all_preds, all_labels, all_probs, cm


def plot_confusion_matrix(cm, class_names, save_path=None):
    """Plot and save confusion matrix"""
    plt.figure(figsize=(8, 6))
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
                xticklabels=class_names, yticklabels=class_names)
    plt.title('Confusion Matrix')
    plt.ylabel('True Label')
    plt.xlabel('Predicted Label')
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"Confusion matrix saved to {save_path}")
    plt.close()


def plot_roc_curve(fpr, tpr, roc_auc, save_path=None):
    """Plot and save ROC curve"""
    plt.figure(figsize=(8, 6))
    plt.plot(fpr, tpr, color='darkorange', lw=2, label=f'ROC curve (AUC = {roc_auc:.2f})')
    plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel('False Positive Rate')
    plt.ylabel('True Positive Rate')
    plt.title('ROC Curve')
    plt.legend(loc="lower right")
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"ROC curve saved to {save_path}")
    plt.close()


def plot_pr_curve(precision, recall, save_path=None):
    """Plot and save Precision-Recall curve"""
    plt.figure(figsize=(8, 6))
    plt.plot(recall, precision, color='blue', lw=2)
    plt.xlabel('Recall')
    plt.ylabel('Precision')
    plt.title('Precision-Recall Curve')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.tight_layout()
    
    if save_path:
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        print(f"PR curve saved to {save_path}")
    plt.close()


def evaluate(config=None, config_path='config.yaml', model_path=None, results_output_path=None):
    """Main evaluation function"""
    
    if config is None:
        config = load_config(config_path)
    
    if results_output_path is None:
        results_output_path = config['output']['results_path']
    
    os.makedirs(results_output_path, exist_ok=True)

    print("=== AI Image Detection - Evaluation ===")
    
    device = torch.device('cuda' if torch.cuda.is_available() and config['device'] == 'cuda' else 'cpu')
    print(f"Device: {device}")
    
    # Load model
    model = AIImageClassifier(
        model_name=config['model']['name'],
        num_classes=config['model']['num_classes'],
        pretrained=False,
        dropout=config['model']['dropout']
    )
    
    if model_path is None:
        model_path = f"{config['output']['checkpoint_path']}/best_model.pth"
    
    checkpoint = torch.load(model_path, map_location=device)
    if isinstance(checkpoint, dict) and 'model_state_dict' in checkpoint:
        model.load_state_dict(checkpoint['model_state_dict'])
    else:
        model.load_state_dict(checkpoint)
    model = model.to(device)
    print(f"Loaded model from {model_path}")
    
    # Load test data
    test_transform = get_transforms(
        image_size=config['preprocessing']['image_size'],
        mode='val',
        normalize_mean=config['preprocessing']['normalize_mean'],
        normalize_std=config['preprocessing']['normalize_std']
    )
    
    test_dataset = AIImageDataset(config['data']['test_path'], transform=test_transform)
    test_loader = DataLoader(
        test_dataset,
        batch_size=config['training']['batch_size'],
        shuffle=False,
        num_workers=config['num_workers']
    )
    
    # Evaluate
    metrics, preds, labels, probs, cm = evaluate_model(model, test_loader, device)
    
    # Print results
    print("\n=== Metrics ===")
    print(f"Accuracy: {metrics['accuracy']:.4f}")
    print(f"F1-Score: {metrics['f1_score']:.4f}")
    print(f"ROC-AUC: {metrics['roc_auc']:.4f}")
    print("\n" + classification_report(labels, preds, target_names=['Natural', 'Synthetic']))
    
    # Save metrics
    results_json_path = f"{results_output_path}/metrics.json"
    with open(results_json_path, 'w') as f:
        json.dump({
            'accuracy': float(metrics['accuracy']),
            'f1_score': float(metrics['f1_score']),
            'roc_auc': float(metrics['roc_auc']),
            'confusion_matrix': metrics['confusion_matrix'],
            'classification_report': metrics['classification_report']
        }, f, indent=2)
    print(f"\nMetrics saved to {results_json_path}")
    
    # Plot results
    plot_confusion_matrix(cm, ['Natural', 'Synthetic'], 
                         f"{results_output_path}/confusion_matrix.png")
    plot_roc_curve(metrics['roc_curve']['fpr'], metrics['roc_curve']['tpr'], 
                   metrics['roc_auc'], f"{results_output_path}/roc_curve.png")
    plot_pr_curve(metrics['pr_curve']['precision'], metrics['pr_curve']['recall'],
                  f"{results_output_path}/pr_curve.png")
    
    print("\n=== Evaluation Complete ===")


if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description='Evaluate AI Image Detection Model')
    parser.add_argument('--config', type=str, default='config.yaml', help='Path to config file')
    parser.add_argument('--model_path', type=str, default=None, help='Path to model checkpoint')
    parser.add_argument('--results_path', type=str, default=None, help='Path to save results')
    args = parser.parse_args()
    
    evaluate(config_path=args.config, model_path=args.model_path, results_output_path=args.results_path)