evaluate.py (5517B)
1 import os 2 import csv 3 4 import torch 5 import torch.nn as nn 6 from torch.utils.data import DataLoader 7 import numpy as np 8 import matplotlib.pyplot as plt 9 import seaborn as sns 10 from sklearn.metrics import ( 11 accuracy_score, 12 f1_score, 13 classification_report, 14 confusion_matrix, 15 cohen_kappa_score, 16 ) 17 18 from utils import CLASS_NAMES 19 20 21 def evaluate_model(model: nn.Module, test_loader: DataLoader, device: torch.device) -> dict: 22 """Run inference on test_loader and return a dict of all metrics.""" 23 model.eval() 24 all_preds, all_labels = [], [] 25 26 with torch.no_grad(): 27 for images, labels in test_loader: 28 images = images.to(device, non_blocking=True) 29 outputs = model(images) 30 preds = outputs.argmax(dim=1) 31 all_preds.extend(preds.cpu().numpy()) 32 all_labels.extend(labels.numpy()) 33 34 all_preds = np.array(all_preds) 35 all_labels = np.array(all_labels) 36 37 return { 38 'accuracy': accuracy_score(all_labels, all_preds), 39 'weighted_f1': f1_score(all_labels, all_preds, average='weighted', zero_division=0), 40 'macro_f1': f1_score(all_labels, all_preds, average='macro', zero_division=0), 41 'classification_report': classification_report( 42 all_labels, all_preds, target_names=CLASS_NAMES, zero_division=0 43 ), 44 'confusion_matrix': confusion_matrix(all_labels, all_preds), 45 'cohen_kappa': cohen_kappa_score(all_labels, all_preds, weights='quadratic'), 46 } 47 48 49 def plot_confusion_matrix(cm: np.ndarray, save_path: str) -> None: 50 """Save a seaborn confusion matrix heatmap to save_path.""" 51 fig, ax = plt.subplots(figsize=(8, 6)) 52 sns.heatmap( 53 cm, annot=True, fmt='d', cmap='Blues', 54 xticklabels=CLASS_NAMES, yticklabels=CLASS_NAMES, 55 ax=ax 56 ) 57 ax.set_xlabel('Predicted') 58 ax.set_ylabel('True') 59 ax.set_title('Confusion Matrix') 60 plt.tight_layout() 61 fig.savefig(save_path, dpi=150) 62 plt.close(fig) 63 64 65 def plot_training_curves( 66 train_losses: list, val_losses: list, val_f1s: list, save_path: str 67 ) -> None: 68 """Save loss and F1 training curves to save_path.""" 69 epochs = range(1, len(train_losses) + 1) 70 fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4)) 71 72 ax1.plot(epochs, train_losses, label='Train loss') 73 ax1.plot(epochs, val_losses, label='Val loss') 74 ax1.set_xlabel('Epoch') 75 ax1.set_ylabel('Loss') 76 ax1.set_title('Training and Validation Loss') 77 ax1.legend() 78 79 ax2.plot(epochs, val_f1s, label='Val macro F1', color='green') 80 ax2.set_xlabel('Epoch') 81 ax2.set_ylabel('Macro F1') 82 ax2.set_title('Validation Macro F1') 83 ax2.legend() 84 85 plt.tight_layout() 86 fig.savefig(save_path, dpi=150) 87 plt.close(fig) 88 89 90 def generate_summary(all_results: dict, save_dir: str) -> None: 91 """Print comparison table, save as CSV, and generate bar chart.""" 92 os.makedirs(save_dir, exist_ok=True) 93 94 rows = [] 95 for key, r in all_results.items(): 96 model_name, clahe_str = key.rsplit('_clahe=', 1) 97 rows.append({ 98 'model': model_name, 99 'clahe': clahe_str, 100 'weighted_f1': r['weighted_f1'], 101 'macro_f1': r['macro_f1'], 102 'accuracy': r['accuracy'], 103 'cohen_kappa': r['cohen_kappa'], 104 'train_time_s': round(r['train_time'], 1), 105 'epochs_trained': r['epochs_trained'], 106 }) 107 108 # Print table 109 header = f"{'Model':<18} {'CLAHE':<6} {'W-F1':>6} {'M-F1':>6} {'Acc':>6} {'Kappa':>6} {'Time(s)':>8} {'Epochs':>7}" 110 print('\n' + '=' * len(header)) 111 print(header) 112 print('-' * len(header)) 113 for row in rows: 114 print( 115 f"{row['model']:<18} {row['clahe']:<6} " 116 f"{row['weighted_f1']:>6.4f} {row['macro_f1']:>6.4f} " 117 f"{row['accuracy']:>6.4f} {row['cohen_kappa']:>6.4f} " 118 f"{row['train_time_s']:>8} {row['epochs_trained']:>7}" 119 ) 120 print('=' * len(header)) 121 122 # Save CSV 123 csv_path = os.path.join(save_dir, 'summary.csv') 124 with open(csv_path, 'w', newline='') as f: 125 writer = csv.DictWriter(f, fieldnames=list(rows[0].keys())) 126 writer.writeheader() 127 writer.writerows(rows) 128 print(f"\nSummary saved to {csv_path}") 129 130 # Comparison bar chart — show both macro and weighted F1 131 x_labels = [f"{r['model']}\nCLAHE={r['clahe']}" for r in rows] 132 w_f1_vals = [r['weighted_f1'] for r in rows] 133 m_f1_vals = [r['macro_f1'] for r in rows] 134 135 x = np.arange(len(x_labels)) 136 bar_width = 0.35 137 fig, ax = plt.subplots(figsize=(12, 5)) 138 bars_w = ax.bar(x - bar_width / 2, w_f1_vals, bar_width, label='Weighted F1', color='#4C72B0') 139 bars_m = ax.bar(x + bar_width / 2, m_f1_vals, bar_width, label='Macro F1', color='#DD8452') 140 ax.set_ylabel('F1-Score') 141 ax.set_title('F1 Comparison Across Experiments') 142 ax.set_xticks(x) 143 ax.set_xticklabels(x_labels) 144 ax.set_ylim(0, 1.0) 145 ax.legend() 146 for bar, val in zip(bars_w, w_f1_vals): 147 ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01, 148 f'{val:.3f}', ha='center', va='bottom', fontsize=8) 149 for bar, val in zip(bars_m, m_f1_vals): 150 ax.text(bar.get_x() + bar.get_width() / 2, bar.get_height() + 0.01, 151 f'{val:.3f}', ha='center', va='bottom', fontsize=8) 152 plt.tight_layout() 153 chart_path = os.path.join(save_dir, 'comparison_chart.png') 154 fig.savefig(chart_path, dpi=150) 155 plt.close(fig) 156 print(f"Comparison chart saved to {chart_path}")