import torch import numpy as np import pandas as pd import torch from torch.utils.data import DataLoader from tqdm import tqdm from dataloading.dataloader2D import NiftiSegmentationDataset import yaml from pathlib import Path import matplotlib.pyplot as plt from sklearn.calibration import calibration_curve from sklearn.metrics import roc_auc_score, roc_curve from sklearn.preprocessing import label_binarize from odelia_breast_mri.scripts.main_predict import evaluate import torchvision.transforms.functional as TF from dataloading.collate_function import custom_collate from models.swinunetr import SwinUNETRMultiTask def tta_transforms(x): """ Apply test-time augmentations to a single batch tensor. Input: x (B, C, H, W) Returns a list of augmented versions. """ return [ x, # original TF.hflip(x), TF.vflip(x), TF.rotate(x, 90), TF.rotate(x, 180), TF.rotate(x, 270), TF.hflip(TF.vflip(x)), # h+v flip ] def plot_multiclass_reliability(probs, labels, class_names=None, n_bins=10): """ Plot calibration curves for each class in a multiclass problem. Args: probs: numpy array of shape (n_samples, n_classes) with predicted probabilities. labels: numpy array of shape (n_samples,) with integer class labels. class_names: list of class names (optional). n_bins: number of bins for calibration curve. """ n_classes = probs.shape[1] if class_names is None: class_names = [f"Class {i}" for i in range(n_classes)] plt.figure(figsize=(8, 8)) for i in range(n_classes): # Binary labels: is this class or not binarized_labels = (labels == i).astype(int) fraction_of_positives, mean_predicted_value = calibration_curve( binarized_labels, probs[:, i], n_bins=n_bins, strategy='uniform' ) plt.plot( mean_predicted_value, fraction_of_positives, "s-", label=f"{class_names[i]}" ) plt.plot([0, 1], [0, 1], "k:", label="Perfectly calibrated") plt.title("Multiclass Reliability Diagram") plt.xlabel("Mean Predicted Probability") plt.ylabel("Fraction of Positives") plt.legend() plt.grid() # plt.show() plt.savefig('/workspace/ClassifierSegmenter/results/multiclass_reliability_diagram.png') # ---------------------- # Load Config & Set Up # ---------------------- with open("/workspace/ClassifierSegmenter/config2d.yaml", "r") as f: config = yaml.safe_load(f) device = torch.device(config["device"] if torch.cuda.is_available() else "cpu") # ---------------------- # Load Model # ---------------------- checkpoint_paths = [ "/workspace/Classifier/checkpoints/final/model1/best_score_model.pth", "/workspace/Classifier/checkpoints/final/model2/best_score_model.pth", "/workspace/Classifier/checkpoints/final/model3/best_score_model.pth", "/workspace/Classifier/checkpoints/final/model4/best_score_model.pth", "/workspace/Classifier/checkpoints/final/model5/best_score_model.pth", "/workspace/Classifier/checkpoints/final/model6/best_score_model.pth", "/workspace/Classifier/checkpoints/final/model7/best_score_model.pth", "/workspace/Classifier/checkpoints/final/model8/best_score_model.pth", "/workspace/Classifier/checkpoints/final/model9/best_score_model.pth", "/workspace/Classifier/checkpoints/final/model10/best_score_model.pth" ] models = [] for i, ckpt_path in enumerate(checkpoint_paths, start=1): model = SwinUNETRMultiTask( img_size=(256, 256), in_channels=4, out_seg_channels=2, out_cls_classes=3 ).to(device) if ckpt_path is not None: state = torch.load(ckpt_path, map_location=device, weights_only=True) model.load_state_dict(state) model.to(device) model.eval() models.append(model) # ---------------------- # Load Test Dataset # ---------------------- test_dataset = NiftiSegmentationDataset(config["csv_file_test"], channel_keys=config["channel_keys"], augment=False) test_loader = DataLoader(test_dataset, batch_size=config["batch_size"], collate_fn=custom_collate, shuffle=False) metrics = {} # ---------------------- # Inference + Metrics # ---------------------- all_preds, all_probs, all_targets = [], [], [] all_patient_ids = [] with torch.no_grad(): if not config["tta"]: for batch in tqdm(test_loader, desc="Running Test Inference"): x = batch['image'].to(device) y = batch['cls_label'].to(device) mask = batch['mask'].to(device) if batch['mask'] is not None else None has_mask = batch['has_mask'].to(device) if batch['has_mask'] is not None else None with torch.autocast(device_type='cuda'): # Collect logits from all models logits_list = [] for model in models: _, out, _ = model(x) logits_list.append(out) # Average probabilities across models probs_list = [torch.softmax(logits, dim=1) for logits in logits_list] probs = torch.stack(probs_list, dim=0).mean(dim=0) preds = torch.argmax(probs, dim=1) all_preds.append(preds.cpu()) all_probs.append(probs.cpu()) all_targets.append(y.cpu()) all_patient_ids.extend(batch['patient_id']) else: for batch in tqdm(test_loader, desc="Running Test Inference w/ TTA"): x = batch['image'].to(device) y = batch['cls_label'].to(device) mask = batch['mask'].to(device) if batch['mask'] is not None else None has_mask = batch['has_mask'].to(device) if batch['has_mask'] is not None else None tta_versions = tta_transforms(x) # list of augmented batches # Combine TTA with ensemble: collect logits for each model and each augmentation model_logits = [] for model in models: all_logits = [] for aug_x in tta_versions: with torch.autocast(device_type='cuda'): _, logits, _ = model(aug_x) all_logits.append(logits) stacked_logits = torch.stack(all_logits, dim=0).mean(dim=0) model_logits.append(stacked_logits) # Average across models mean_logits = torch.stack(model_logits, dim=0).mean(dim=0) probs = torch.softmax(mean_logits, dim=1) preds = torch.argmax(probs, dim=1) all_preds.append(preds.cpu()) all_probs.append(probs.cpu()) all_targets.append(y.cpu()) all_patient_ids.extend(batch['patient_id']) # Concatenate everything all_preds = torch.cat(all_preds) all_probs = torch.cat(all_probs) all_targets = torch.cat(all_targets) # ---------------------- # Evaluate MACRO metrics # ---------------------- accuracy = (all_preds == all_targets).sum().item() / len(all_targets) auc, sensitivity, specificity = evaluate( all_targets, all_preds, all_probs, path_out=Path('/workspace/ClassifierSegmenter/results') ) print("\nāœ… MACRO Results") print(f"Accuracy: {accuracy:.4f}") print(f"AUC: {auc:.4f}") print(f"Sensitivity: {sensitivity:.4f}") print(f"Specificity: {specificity:.4f}") plot_multiclass_reliability( all_probs.numpy(), all_targets.numpy(), class_names=['no lesion', 'benign', 'malignant'], n_bins=10 ) # ---------------------- # Evaluate MICRO metrics # ---------------------- y_true_hot = label_binarize(all_targets.numpy().astype(str), classes=['0', '1', '2']) fpr, tpr, thresholds = roc_curve(y_true_hot.ravel(), all_probs.numpy().ravel(), drop_intermediate=False) roc_auc = roc_auc_score(y_true_hot, all_probs.numpy(), average="micro") roc_auc_macro = roc_auc_score(y_true_hot, all_probs.numpy(), average="macro") # Sensitivity at 90% specificity specificity_threshold = 0.90 fpr_threshold = 1 - specificity_threshold sensitivity_at_90_specificity = np.interp(fpr_threshold, fpr, tpr) # Specificity at 90% sensitivity sensitivity_threshold = 0.90 fpr_at_90_sensitivity = np.interp(sensitivity_threshold, tpr, fpr) specificity_at_90_sensitivity = 1 - fpr_at_90_sensitivity amalgamated_results = [roc_auc, specificity_at_90_sensitivity, sensitivity_at_90_specificity] averaged_results = np.mean(amalgamated_results) metrics["results"] = { "AUC": roc_auc, "AUC macro": roc_auc_macro, "Specificity": specificity_at_90_sensitivity, "Sensitivity": sensitivity_at_90_specificity, "Score": averaged_results } print("\nāœ… MICRO Results") print(f"AUC: {roc_auc:.4f}") print(f"Sensitivity: {sensitivity_at_90_specificity:.4f}") print(f"Specificity: {specificity_at_90_sensitivity:.4f}") print(f"Score: {averaged_results:.4f}") # Save results results_df = pd.DataFrame({ "PatientID": all_patient_ids, "TrueLabel": all_targets.numpy(), "PredLabel": all_preds.numpy(), **{f"Prob_Class_{i}": all_probs[:, i].numpy() for i in range(3)} }) results_df.to_csv("/workspace/ClassifierSegmenter/results/test_predictions.csv", index=False) print("\nšŸ“ Saved test predictions to 'results/test_predictions.csv'")