evaluate_model.py

#!/usr/bin/env python3
# filepath: /home/omer/Masaüstü/code-base/AUTSL_transformer_trainning/evaluate_model.py
import os
import time
import torch
import torch.nn.functional as F
import pandas as pd
import numpy as np
import argparse
from torch.utils.data import DataLoader
import pytorch_lightning as pl
from train import SignLanguageTransformer, SignDataset
import yaml

def parse_args():
    parser = argparse.ArgumentParser(description='Test Sign Language Transformer Model')
    parser.add_argument('--version', type=int, required=True,
                        help='Version number of the model in the logs/my_model directory')
    parser.add_argument('--test_csv', type=str, default='test_labels.csv',
                        help='Path to test CSV file')
    parser.add_argument('--test_dir', type=str,
                        default='/home/omer/Masaüstü/datasets/AUTSL_medipipe_landmarks/test',
                        help='Directory containing test data files')
    parser.add_argument('--batch_size', type=int, default=64,
                        help='Batch size for testing')
    return parser.parse_args()

def _macro_f1_from_preds_targets(preds: torch.Tensor, targets: torch.Tensor, num_classes: int) -> float:
    """
    Macro F1 (average over classes appearing in targets).
    preds/targets: shape [N], int64
    """
    eps = 1e-12
    preds = preds.to(torch.int64).cpu()
    targets = targets.to(torch.int64).cpu()

    # Confusion matrix: rows=targets, cols=preds
    idx = targets * num_classes + preds
    cm = torch.bincount(idx, minlength=num_classes * num_classes).reshape(num_classes, num_classes)

    tp = torch.diag(cm).to(torch.float32)
    fp = cm.sum(dim=0).to(torch.float32) - tp
    fn = cm.sum(dim=1).to(torch.float32) - tp

    precision = tp / (tp + fp + eps)
    recall = tp / (tp + fn + eps)
    f1 = 2.0 * precision * recall / (precision + recall + eps)

    support = cm.sum(dim=1).to(torch.float32)  # target count per class
    valid = support > 0
    if valid.any():
        return f1[valid].mean().item()
    return 0.0

def evaluate_model(model, test_loader):
    """
    Evaluates the model on the test dataset and calculates desired metrics.
    """
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = model.to(device)
    model.eval()

    all_probs = []
    all_preds = []
    all_targets = []

    total_inference_seconds = 0.0
    total_samples = 0

    print(f"Device: {device}")
    print("Evaluating test dataset...")

    with torch.no_grad():
        for batch in test_loader:
            x, y = batch
            x = x.to(device, non_blocking=True)
            y = y.to(device, non_blocking=True)

            # Measure forward (inference) time only
            if device.type == "cuda":
                torch.cuda.synchronize()
            t0 = time.perf_counter()

            logits = model(x)

            if device.type == "cuda":
                torch.cuda.synchronize()
            t1 = time.perf_counter()

            total_inference_seconds += (t1 - t0)
            total_samples += x.shape[0]

            probs = F.softmax(logits, dim=1)
            preds = torch.argmax(logits, dim=1)

            all_probs.append(probs.detach().cpu())
            all_preds.append(preds.detach().cpu())
            all_targets.append(y.detach().cpu())

    # Concatenate all results
    all_probs = torch.cat(all_probs, dim=0)
    all_preds = torch.cat(all_preds, dim=0)
    all_targets = torch.cat(all_targets, dim=0)

    # Calculate accuracy
    correct = all_preds == all_targets
    accuracy = correct.float().mean().item()

    # Probability values of predicted classes (top-1 values)
    top1_scores = torch.gather(all_probs, 1, all_preds.unsqueeze(1)).squeeze(1)

    # Probability values of true classes
    true_class_scores = torch.gather(all_probs, 1, all_targets.unsqueeze(1)).squeeze(1)

    # Separate correctly and incorrectly classified samples
    correct_mask = correct
    incorrect_mask = ~correct

    # Calculate desired metrics
    avg_top1_all = top1_scores.mean().item()
    avg_top1_correct = top1_scores[correct_mask].mean().item() if correct_mask.sum() > 0 else 0
    avg_top1_incorrect = top1_scores[incorrect_mask].mean().item() if incorrect_mask.sum() > 0 else 0
    avg_true_class_incorrect = true_class_scores[incorrect_mask].mean().item() if incorrect_mask.sum() > 0 else 0

    # F1 (macro)
    num_classes = all_probs.shape[1]
    macro_f1 = _macro_f1_from_preds_targets(all_preds, all_targets, num_classes=num_classes)

    # Average inference time (per sample)
    avg_inference_seconds_per_sample = (total_inference_seconds / max(total_samples, 1))
    avg_inference_ms_per_sample = avg_inference_seconds_per_sample * 1000.0

    # Print results
    print("\n" + "="*50)
    print(f"Test Accuracy: {accuracy:.4f}")
    print(f"Macro F1: {macro_f1:.4f}")
    print(f"Average inference time: {avg_inference_ms_per_sample:.3f} ms/sample")
    print(f"Total samples: {len(all_targets)}")
    print(f"Correct predictions: {correct_mask.sum().item()}")
    print(f"Incorrect predictions: {incorrect_mask.sum().item()}")
    print("="*50)
    print(f"Average top-1 score for all samples: {avg_top1_all:.4f}")
    print(f"Average top-1 score for correct predictions: {avg_top1_correct:.4f}")
    print(f"Average top-1 score for incorrect predictions: {avg_top1_incorrect:.4f}")
    print(f"Average true class score for incorrect predictions: {avg_true_class_incorrect:.4f}")
    print("="*50)

    return {
        "accuracy": accuracy,
        "macro_f1": macro_f1,
        "avg_inference_seconds_per_sample": avg_inference_seconds_per_sample,
        "avg_inference_ms_per_sample": avg_inference_ms_per_sample,
        "avg_top1_all": avg_top1_all,
        "avg_top1_correct": avg_top1_correct,
        "avg_top1_incorrect": avg_top1_incorrect,
        "avg_true_class_incorrect": avg_true_class_incorrect,
        "total_samples": len(all_targets),
        "correct_predictions": correct_mask.sum().item(),
        "incorrect_predictions": incorrect_mask.sum().item(),
        "total_inference_seconds": total_inference_seconds
    }

def main():
    # Parse command line arguments
    args = parse_args()

    # Set log and checkpoint paths
    log_dir = f"/home/omer/Masaüstü/code-base/AUTSL_transformer_trainning/logs/yeni_rk4_learnable/version_{args.version}"
    checkpoints_dir = os.path.join(log_dir, "checkpoints")

    checkpoints_to_test = {
        "max_val_acc": os.path.join(checkpoints_dir, "max_val_acc.ckpt"),
        "last": os.path.join(checkpoints_dir, "last.ckpt")
    }

    all_results = {}
    hparams_to_save = None

    for name, checkpoint_path in checkpoints_to_test.items():
        if not os.path.exists(checkpoint_path):
            print(f"\nWarning: Checkpoint not found, skipping: {checkpoint_path}")
            continue

        print("\n" + "#"*60)
        print(f"TEST STARTED: Checkpoint '{name}' ({checkpoint_path})")
        print("#"*60)

        # Load model from checkpoint
        print(f"Loading model: {checkpoint_path}")
        model = SignLanguageTransformer.load_from_checkpoint(checkpoint_path)

        # Retrieve hyperparameters from the first model once
        if hparams_to_save is None:
            hparams_to_save = {
                'd_model': model.hparams.d_model,
                'dropout': model.hparams.dropout,
                'enc_calculate_num': model.hparams.enc_calculate_num,
                'encoder_history_type': model.hparams.encoder_history_type,
                'num_layers': model.hparams.num_layers,
                'pre_norm': model.hparams.pre_norm,
                'rk_type': model.hparams.rk_type
            }

        # Display model parameters
        print(f"\nModel parameters:")
        print(f"  num_frames: {model.hparams.num_frames}")
        print(f"  d_model: {model.hparams.d_model}")
        print(f"  nhead: {model.hparams.nhead}")
        print(f"  num_layers: {model.hparams.num_layers}")
        print(f"  dropout: {model.hparams.dropout}")
        print(f"  learning rate: {model.hparams.lr}")

        # Create test dataset
        test_dataset = SignDataset(
            csv_path=args.test_csv,
            data_dir=args.test_dir,
            num_frames=model.hparams.num_frames,
            reduce_features=False,
            drop_visibility=False
        )

        test_loader = DataLoader(
            test_dataset,
            batch_size=args.batch_size,
            shuffle=False,
            num_workers=4,
            pin_memory=torch.cuda.is_available()
        )

        print(f"\nTest dataset: {len(test_dataset)} samples")

        # Evaluate the model
        metrics = evaluate_model(model, test_loader)
        all_results[name] = metrics

    # Save results to YAML file
    if all_results:
        output_dir = "test_logs_f1-2"
        os.makedirs(output_dir, exist_ok=True)
        output_path = os.path.join(output_dir, f"test_{args.version}.yaml")

        # Create final output data
        final_output = {}
        if hparams_to_save:
            final_output['hyperparameters'] = hparams_to_save
        final_output['test_results'] = all_results

        with open(output_path, 'w') as f:
            yaml.dump(final_output, f, default_flow_style=False, sort_keys=False)

        print(f"\nTest results saved to: {output_path}")
    else:
        print("\nNo valid checkpoint found for evaluation.")

if __name__ == "__main__":
    main()