| from pathlib import PurePath
|
|
|
| import wandb
|
| import torch
|
| import torch.utils.data
|
| import torch.nn as nn
|
| from torch.amp import GradScaler
|
| from torchmetrics import AUROC, Accuracy, Precision, Specificity, ROC
|
|
|
| from auto_detect_breast_mri.config import get_config, resolve_path
|
| from auto_detect_breast_mri.data.metadata import get_uka_metatensor
|
| from auto_detect_breast_mri.data import loaders
|
| from auto_detect_breast_mri.models.resnets import model_names
|
| from auto_detect_breast_mri.data.transforms import pre_image_shape, minimalAugmentation
|
| from auto_detect_breast_mri.models.checkpoints import load_pretrained_model
|
| from auto_detect_breast_mri.training.cli import generate_parser
|
| from auto_detect_breast_mri.training.loops import eval_epoch_subjects
|
| CPU = "cpu"
|
| GPU = "cuda"
|
|
|
| torch.manual_seed(31)
|
|
|
| DEVICE = GPU
|
| if not torch.cuda.is_available():
|
| DEVICE = CPU
|
|
|
| transform = minimalAugmentation
|
| t = 'basic'
|
| batch_size = 32
|
| fraction = 1.0
|
| fold = 0
|
|
|
| parser = generate_parser()
|
| args = parser.parse_args()
|
| model_name = args.model_name
|
| path_base = resolve_path(args.data_path, "data_root", "root folder of the NIfTI data")
|
| feature_path = resolve_path(args.feature_path, "metadata_file", "metadata export")
|
| split_files_folder = resolve_path(args.split_files_folder, "split_root",
|
| "folder holding the split files")
|
| model_path = args.model_path
|
|
|
| split_files_folder_name = PurePath(split_files_folder.rstrip('/')).name
|
|
|
|
|
| split_file_path = split_files_folder + f"fold{fold}/stratified_test_set"
|
| feature_dataframe = get_uka_metatensor(0, feature_path)
|
| test_loader_abrv = loaders.get_subjects_dataloader(path_base,
|
| feature_dataframe,
|
| pre_image_shape,
|
| transform,
|
| 'abbreviated',
|
| split_file_path,
|
| batch_size,
|
| fraction,
|
| fold,
|
| subfold=0)
|
|
|
| test_loader_full = loaders.get_subjects_dataloader(path_base,
|
| feature_dataframe,
|
| pre_image_shape,
|
| transform,
|
| 'full',
|
| split_file_path,
|
| batch_size,
|
| fraction,
|
| fold,
|
| subfold=0)
|
|
|
| criterion = nn.BCEWithLogitsLoss()
|
| auc_roc = nn.ModuleDict({state: AUROC(**{"task": "binary"}).to(DEVICE) for state in [model_name + "_abrv", model_name + "_full"]})
|
| roc = nn.ModuleDict({state: ROC(**{"task": "binary"}).to(DEVICE) for state in [model_name + "_abrv", model_name + "_full"]})
|
| acc = nn.ModuleDict({state: Accuracy(**{"task": "binary"}).to(DEVICE) for state in [model_name + "_abrv", model_name + "_full"]})
|
| sens = nn.ModuleDict({state: Precision(**{"task": "binary"}).to(DEVICE) for state in [model_name + "_abrv", model_name + "_full"]})
|
| spec = nn.ModuleDict({state: Specificity(**{"task": "binary"}).to(DEVICE) for state in [model_name + "_abrv", model_name + "_full"]})
|
| values = (auc_roc, roc, acc, sens, spec)
|
|
|
| model_abrv = model_names.get(model_name + '_abrv')
|
| model_full = model_names.get(model_name + '_full')
|
|
|
| scaler = GradScaler()
|
| used_mixed_precision = True
|
|
|
| wandb.init(**get_config().wandb_init_kwargs(), config={
|
| "learning-rate": lr,
|
| "model": model_name,
|
| "number of test samples abrv": len(test_loader_abrv.dataset),
|
| "number of test samples full": len(test_loader_full.dataset),
|
| "batch size": batch_size,
|
| "Augmentation": t + str(transform),
|
| "Machine": "HPC",
|
| "Mixed Precision": used_mixed_precision,
|
| "Fold": fold,
|
| "splits": split_files_folder_name,
|
| }, name="{}_comparison_split={}".format(model_name, fold))
|
|
|
|
|
| model_abrv.to(DEVICE)
|
| model_full.to(DEVICE)
|
|
|
|
|
|
|
| abrv_model_path = model_path.format(model_name + "_abrv", fold)
|
| print("load " + abrv_model_path)
|
| checkpoint_abrv = load_pretrained_model(model_abrv, abrv_model_path, DEVICE)
|
| model_abrv = model_abrv.load_state_dict(checkpoint_abrv)
|
| model_abrv.eval()
|
|
|
| full_model_path = model_path.format(model_name + "_full", fold)
|
| print("load " + full_model_path)
|
| checkpoint_full = load_pretrained_model(model_full, full_model_path, DEVICE)
|
| model_full = model_abrv.load_state_dict(checkpoint_full)
|
| model_full.eval()
|
|
|
| abrv_loss, abrv_results_dict = eval_epoch_subjects(model_abrv, criterion, test_loader_abrv, DEVICE, values, model_name + '_abrv', plot_tpr_fpr=True, scaler=scaler, title_suffix="ABRV")
|
| full_loss, full_results_dict = eval_epoch_subjects(model_full, criterion, test_loader_full, DEVICE, values, model_name + '_full', plot_tpr_fpr=True, scaler=scaler, title_suffix="FULL")
|
|
|
| filename = f"AUC_ROC-{model_name}_{fold}.png"
|
| fpr_abrv, tpr_abrv, thresholds_abrv = roc[model_name + '_abrv'].to('cpu').compute()
|
| fpr_full, tpr_full, thresholds_full = roc[model_name + '_full'].to('cpu').compute()
|
| y_pred = {model_name + '_abrv': tpr_abrv, model_name + '_full': tpr_full}
|
| results['fpr'] = fpr
|
| results['thresholds'] = thresholds
|
| plot_roc_curve(modelname, filename) |
