test.py

import argparse
import torch
import torch.nn.functional as F
import numpy as np
from torch.utils.data import DataLoader
from sklearn.metrics import f1_score, roc_auc_score, accuracy_score

from data.load_data import CropData
from models.cnn.resnet import get_resnet18
from models.vit.deit import get_deit3
from engine.trainer import get_optimal_device


def evaluate(model, test_loader, device):
    model.eval()
    all_targets = []
    all_preds = []
    all_probs = []

    with torch.no_grad():
        for images, targets in test_loader:
            images, targets = images.to(device), targets.to(device)
            logits = model(images)
            probs = F.softmax(logits, dim=1)
            _, preds = torch.max(logits, 1)

            all_probs.append(probs.cpu().numpy())
            all_preds.extend(preds.cpu().numpy())
            all_targets.extend(targets.cpu().numpy())

    all_targets = np.array(all_targets)
    all_preds = np.array(all_preds)
    all_probs = np.concatenate(all_probs, axis=0)

    accuracy = accuracy_score(all_targets, all_preds)
    macro_f1 = f1_score(all_targets, all_preds, average='macro')
    macro_auc = roc_auc_score(all_targets, all_probs, multi_class='ovr', average='macro')

    return accuracy, macro_f1, macro_auc


def main():
    parser = argparse.ArgumentParser(description="Test Script for COL780 Assignment 3")
    parser.add_argument("--checkpoint", type=str, required=True, help="Path to the trained model checkpoint")
    parser.add_argument("--img_dir", type=str, default="data/A3_Dataset", help="Path to image directory")
    parser.add_argument("--test_csv", type=str, default="data/A3_Dataset/test.csv", help="Path to test CSV labels")
    parser.add_argument("--batch_size", type=int, default=32)
    parser.add_argument("--num_classes", type=int, default=10)
    
    parser.add_argument("--model_family", type=str, choices=["cnn", "transformer"], required=True, help="Family of the model to evaluate")
    parser.add_argument("--use_se", action="store_true", help="If provided, uses ResNet-18 with SE Blocks")
    parser.add_argument("--use_dyt", action="store_true", help="If provided, uses DeiT-3 with Dynamic Tanh (DyT)")
    
    args = parser.parse_args()

    device = get_optimal_device()

    print(f"Loading Test Dataset from {args.test_csv}...")
    test_dataset = CropData(img_dir_path=args.img_dir, csv_file_path=args.test_csv)
    test_loader = DataLoader(test_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0)

    print(f"Initializing {args.model_family.upper()} Model...")
    if args.model_family == "cnn":
        model = get_resnet18(num_classes=args.num_classes, use_se=args.use_se)
    elif args.model_family == "transformer":
        model = get_deit3(num_classes=args.num_classes, use_dyt=args.use_dyt)
    else:
        raise ValueError("Invalid model family")

    print(f"Loading weights from {args.checkpoint}...")
    model.load_state_dict(torch.load(args.checkpoint, map_location=device))
    model.to(device)

    print("Evaluating...")
    accuracy, macro_f1, macro_auc = evaluate(model, test_loader, device)

    print("\n--- TEST METRICS ---")
    print(f"Accuracy:       {accuracy:.4f}")
    print(f"Macro F1-Score: {macro_f1:.4f}")
    print(f"Macro ROC-AUC:  {macro_auc:.4f}")


if __name__ == "__main__":
    main()