eval.py

from __future__ import print_function

import numpy as np

import matplotlib.pyplot as plt
import argparse
import torch
import torch.nn as nn
import pdb
import os
import pandas as pd
from utils.utils import *
from math import floor
import matplotlib.pyplot as plt
from dataset_modules.dataset_generic import Generic_WSI_Classification_Dataset, Generic_MIL_Dataset, save_splits
import h5py
from utils.eval_utils import *

# Training settings
parser = argparse.ArgumentParser(description='CLAM Evaluation Script')
parser.add_argument('--data_root_dir', type=str, default=None,
                    help='data directory')
parser.add_argument('--results_dir', type=str, default='./results',
                    help='relative path to results folder, i.e. '+
                    'the directory containing models_exp_code relative to project root (default: ./results)')
parser.add_argument('--save_exp_code', type=str, default=None,
                    help='experiment code to save eval results')
parser.add_argument('--models_exp_code', type=str, default=None,
                    help='experiment code to load trained models (directory under results_dir containing model checkpoints')
parser.add_argument('--splits_dir', type=str, default=None,
                    help='splits directory, if using custom splits other than what matches the task (default: None)')
parser.add_argument('--model_size', type=str, choices=['small', 'big'], default='small', 
                    help='size of model (default: small)')
parser.add_argument('--model_type', type=str, choices=['clam_sb', 'clam_mb', 'mil'], default='clam_sb', 
                    help='type of model (default: clam_sb)')
parser.add_argument('--k', type=int, default=10, help='number of folds (default: 10)')
parser.add_argument('--k_start', type=int, default=-1, help='start fold (default: -1, last fold)')
parser.add_argument('--k_end', type=int, default=-1, help='end fold (default: -1, first fold)')
parser.add_argument('--fold', type=int, default=-1, help='single fold to evaluate')
parser.add_argument('--micro_average', action='store_true', default=False, 
                    help='use micro_average instead of macro_avearge for multiclass AUC')
parser.add_argument('--split', type=str, choices=['train', 'val', 'test', 'all'], default='test')
parser.add_argument('--task', type=str, choices=['bc_cls', 'task_1_tumor_vs_normal',  'task_2_tumor_subtyping'])
parser.add_argument('--drop_out', type=float, default=0.25, help='dropout')
parser.add_argument('--embed_dim', type=int, default=1024)
args = parser.parse_args()

device=torch.device("cuda" if torch.cuda.is_available() else "cpu")

args.save_dir = os.path.join('./eval_results', 'EVAL_' + str(args.save_exp_code))
args.models_dir = os.path.join(args.results_dir, str(args.models_exp_code))

os.makedirs(args.save_dir, exist_ok=True)

if args.splits_dir is None:
    args.splits_dir = args.models_dir

assert os.path.isdir(args.models_dir)
assert os.path.isdir(args.splits_dir)

settings = {'task': args.task,
            'split': args.split,
            'save_dir': args.save_dir, 
            'models_dir': args.models_dir,
            'model_type': args.model_type,
            'drop_out': args.drop_out,
            'model_size': args.model_size}

with open(args.save_dir + '/eval_experiment_{}.txt'.format(args.save_exp_code), 'w') as f:
    print(settings, file=f)
f.close()

print(settings)
if args.task == 'task_1_tumor_vs_normal':
    args.n_classes=2
    dataset = Generic_MIL_Dataset(csv_path = 'dataset_csv/tumor_vs_normal_dummy_clean.csv',
                            data_dir= os.path.join(args.data_root_dir, 'tumor_vs_normal_resnet_features'),
                            shuffle = False, 
                            print_info = True,
                            label_dict = {'normal_tissue':0, 'tumor_tissue':1},
                            patient_strat=False,
                            ignore=[])

elif args.task == 'task_2_tumor_subtyping':
    args.n_classes=3
    dataset = Generic_MIL_Dataset(csv_path = 'dataset_csv/tumor_subtyping_dummy_clean.csv',
                            data_dir= os.path.join(args.data_root_dir, 'tumor_subtyping_resnet_features'),
                            shuffle = False, 
                            print_info = True,
                            label_dict = {'subtype_1':0, 'subtype_2':1, 'subtype_3':2},
                            patient_strat= False,
                            ignore=[])
elif args.task == 'bc_cls':
    args.n_classes=2
    dataset = Generic_MIL_Dataset(csv_path = 'data/processed/all/process_list_autogen.csv',
                            data_dir= os.path.join(args.data_root_dir, 'resnet_features'),
                            shuffle = False,
                            print_info = True,
                            label_dict = {0:0, 1:1},
                            patient_strat=True,
                            ignore=[])

# elif args.task == 'tcga_kidney_cv':
#     args.n_classes=3
#     dataset = Generic_MIL_Dataset(csv_path = 'dataset_csv/tcga_kidney_clean.csv',
#                             data_dir= os.path.join(args.data_root_dir, 'tcga_kidney_20x_features'),
#                             shuffle = False, 
#                             print_info = True,
#                             label_dict = {'TCGA-KICH':0, 'TCGA-KIRC':1, 'TCGA-KIRP':2},
#                             patient_strat= False,
#                             ignore=['TCGA-SARC'])

else:
    raise NotImplementedError

if args.k_start == -1:
    start = 0
else:
    start = args.k_start
if args.k_end == -1:
    end = args.k
else:
    end = args.k_end

if args.fold == -1:
    folds = range(start, end)
else:
    folds = range(args.fold, args.fold+1)
ckpt_paths = [os.path.join(args.models_dir, 's_{}_checkpoint.pt'.format(fold)) for fold in folds]
datasets_id = {'train': 0, 'val': 1, 'test': 2, 'all': -1}

if __name__ == "__main__":
    all_results = []
    all_auc = []
    all_acc = []
    best_results_per_fold = []        # Will store best metrics for each fold
    best_thresholds_per_fold = []     # Best threshold for each fold
    best_avg_per_fold = []            # Best avg(sens, spec) per fold

    # Define threshold search range
    thresholds_to_try = np.arange(0.40, 0.60, 0.01)  # You can adjust this range
    
    optimal_thresholds = [0.500, 0.580, 0.490, 0.460, 0.500] # found best thresholds

    print(">>> Searching for optimal threshold individually per fold...\n")

    for ckpt_idx in range(len(ckpt_paths)):
        print(f'>>>>>> Processing Fold {folds[ckpt_idx]}')

        # Load the correct split for this fold
        if datasets_id[args.split] < 0:
            split_dataset = dataset
        else:
            csv_path = '{}/splits_{}.csv'.format(args.splits_dir, folds[ckpt_idx])
            datasets = dataset.return_splits(from_id=False, csv_path=csv_path)
            split_dataset = datasets[datasets_id[args.split]]

        # Variables to track best threshold for this fold
        fold_best_threshold = 0.5
        fold_best_avg = -1
        fold_best_min = -1
        fold_best_metrics = None

        if optimal_thresholds is not None:
            thresholds_to_try = [optimal_thresholds[ckpt_idx]]

        for threshold in thresholds_to_try:
            print(f'   Trying threshold: {threshold:.3f}', end='\r')

            model, patient_results, test_error, auc, df, metrics = eval(
                split_dataset, args, ckpt_paths[ckpt_idx], threshold=threshold
            )

            sens = metrics['sensitivity']
            spec = metrics['specificity']
            current_avg = (sens + spec) / 2
            current_min = min(sens, spec)

            # Selection criteria: maximize average of sens & spec
            # With optional guard: don't pick high avg if one metric is too low
            if (current_avg > fold_best_avg) or (fold_best_metrics is None):
                if not (current_avg < 0.90 and current_min > 0.90):  # avoids degenerate cases
                    fold_best_threshold = threshold
                    fold_best_avg = current_avg
                    fold_best_min = current_min
                    fold_best_metrics = metrics
                    fold_best_auc = auc
                    fold_best_acc = 1 - test_error
                    fold_best_df = df

        print(f"\n   Best threshold for fold {folds[ckpt_idx]}: {fold_best_threshold:.3f}")
        print(f"      Sensitivity: {100*fold_best_metrics['sensitivity']:.2f}%")
        print(f"      Specificity: {100*fold_best_metrics['specificity']:.2f}%")
        print(f"      Accuracy:    {100*fold_best_acc:.2f}%")
        print(f"      AUC:         {fold_best_auc:.4f}\n")

        # Save results with the best threshold for this fold
        fold_best_df = fold_best_df.rename(columns={'slide_id': 'ws_id', 'Y': 'label', 'Y_hat': 'pred', 'p_1': 'prob'})
        fold_best_df = fold_best_df[['ws_id', 'label', 'pred', 'prob']]
        output_fullname = os.path.join(args.save_dir, f'results_fold_{folds[ckpt_idx]}_threshold_{threshold}.csv')
        print(f'Fold [{ckpt_idx}] Saving predictions to {output_fullname}')
        fold_best_df.to_csv(
            output_fullname,
            index=False
        )

        # Save confusion matrices
        cm = fold_best_metrics['confusion_matrix']
        cm_norm = fold_best_metrics['confusion_matrix_normalized'] * 100

        fig = ConfusionMatrixDisplay(cm, display_labels=["Vu lanh", "K vu"])
        fig.plot()
        plt.title(f'Fold {folds[ckpt_idx]} - Raw CM (Th={fold_best_threshold:.3f})')
        plt.savefig(os.path.join(args.save_dir, f'confusion_matrix_fold{folds[ckpt_idx]}.png'), dpi=300)
        plt.close()

        fig_norm = ConfusionMatrixDisplay(cm_norm, display_labels=["Vu lanh", "K vu"])
        fig_norm.plot(values_format=".1f")
        plt.title(f'Fold {folds[ckpt_idx]} - Normalized CM (%)')
        plt.savefig(os.path.join(args.save_dir, f'confusion_matrix_normalized_fold{folds[ckpt_idx]}.png'), dpi=300)
        plt.close()

        # Store best per-fold results
        best_results_per_fold.append(fold_best_metrics)
        best_thresholds_per_fold.append(fold_best_threshold)
        best_avg_per_fold.append(fold_best_avg)
        all_auc.append(fold_best_auc)
        all_acc.append(fold_best_acc)

    # After all folds: summarize across folds (using individually optimized thresholds)
    print("\n" + "="*60)
    print("FINAL SUMMARY (Each fold uses its own optimal threshold)")
    print("="*60)

    mean_sens = np.mean([r['sensitivity'] for r in best_results_per_fold])
    std_sens = np.std([r['sensitivity'] for r in best_results_per_fold])
    mean_spec = np.mean([r['specificity'] for r in best_results_per_fold])
    std_spec = np.std([r['specificity'] for r in best_results_per_fold])
    mean_acc = np.mean(all_acc)
    std_acc = np.std(all_acc)
    mean_auc = np.mean(all_auc)

    print(f"Mean Sensitivity:  {100*mean_sens:.2f} ± {100*std_sens:.2f}%")
    print(f"Mean Specificity:  {100*mean_spec:.2f} ± {100*std_spec:.2f}%")
    print(f"Mean Accuracy:     {100*mean_acc:.2f} ± {100*std_acc:.2f}%")
    print(f"Mean AUC:          {mean_auc:.4f}")
    print(f"Best thresholds per fold: {[f'{t:.3f}' for t in best_thresholds_per_fold]}")