alnn_rewrite/evaluate_models.py

# This program evaluates every model on the combined validation and test set, then saves the results to a netcdf file.

import json
import pathlib as pl

import numpy as np
import pandas as pd
import torch
import xarray as xr
from torch.utils.data import DataLoader

# Custom modules
from data.dataset import (
    ADNIDataset,
    divide_dataset,
    initalize_dataloaders,
    load_adni_data_from_file,
)

# Config
from model.cnn import CNN3D
from utils.config import config

mri_files = pl.Path(config["data"]["mri_files_path"]).glob("*.nii")
xls_file = pl.Path(config["data"]["xls_file_path"])


def xls_pre(df: pd.DataFrame) -> pd.DataFrame:
    """
    Preprocess the Excel DataFrame.
    This function can be customized to filter or modify the DataFrame as needed.
    """

    data = df[["Image Data ID", "Sex", "Age (current)"]]
    data["Sex"] = data["Sex"].str.strip()  # type: ignore
    data = data.replace({"M": 0, "F": 1})  # type: ignore
    data.set_index("Image Data ID")  # type: ignore

    return data


dataset = load_adni_data_from_file(
    mri_files, xls_file, device=config["training"]["device"], xls_preprocessor=xls_pre
)

# Divide the dataset into training and validation sets, using the same seed as training
with open(pl.Path(config["output"]["path"]) / "config.json") as f:
    training_config = json.load(f)
    try:
        loaded_seed = int(training_config["data"]["seed"])
    except (ValueError, KeyError) as e:
        print(
            f"Warning: No previous seed found for dataset division, using seed from config. Error: {e}"
        )
        loaded_seed = config["data"]["seed"]


datasets = divide_dataset(dataset, config["data"]["data_splits"], seed=loaded_seed)


# Initialize the dataloadersx
train_loader, val_loader, test_loader = initalize_dataloaders(
    datasets, batch_size=config["training"]["batch_size"]
)


# Combine validation and test sets for final evaluation
combined_loader: DataLoader[ADNIDataset] = torch.utils.data.DataLoader(
    torch.utils.data.ConcatDataset([val_loader.dataset, test_loader.dataset]),
    batch_size=1,
    shuffle=False,
)


# 50 models are too large to load into memory at once, so we will load and evaluate them one at a time
model_dir = pl.Path(config["output"]["path"])
model_files = sorted(model_dir.glob("model_run_*.pt"))

placeholder = np.zeros(
    (len(model_files), len(combined_loader), config["data"]["num_classes"]),
    dtype=np.float32,
)  # Placeholder for results

# Get the total list of image_ids
img_ids = [img_id for _, _, _, img_id in combined_loader.dataset]


placeholder[:] = np.nan  # Fill with NaNs for easier identification of missing data
dimensions = ["model", "img_id", "img_class"]
coords = {
    "model": [int(mf.stem.split("_")[2]) for mf in model_files],
    "img_id": img_ids,
    "img_class": list(range(config["data"]["num_classes"])),
}
results = xr.DataArray(placeholder, coords=coords, dims=dimensions)

# Now initialize an additional dataarray to hold the labels per image
labels_placeholder = np.zeros(
    (len(combined_loader), config["data"]["num_classes"]), dtype=np.float32
)
labels_placeholder[:] = np.nan
labels_coords = {
    "img_id": img_ids,
    "label": list(range(config["data"]["num_classes"])),
}  # type: ignore

labels = xr.DataArray(
    labels_placeholder, coords=labels_coords, dims=["img_id", "label"]
)


for model_file in model_files:
    model_num = int(model_file.stem.split("_")[2])
    print(f"Evaluating model {model_num}...")

    # Load the model state
    model = (
        CNN3D(
            image_channels=config["data"]["image_channels"],
            clin_data_channels=config["data"]["clin_data_channels"],
            num_classes=config["data"]["num_classes"],
            droprate=config["training"]["droprate"],
        )
        .float()
        .to(config["training"]["device"])
    )

    model.load_state_dict(
        torch.load(model_file, map_location=config["training"]["device"]), strict=False
    )
    model.eval()

    with torch.no_grad():
        for batch_idx, (mri, xls, label, img_id) in enumerate(combined_loader):
            outputs = model((mri.float(), xls.float()))
            probabilities = outputs.cpu().numpy()[0, :]  # type: ignore

            results.loc[model_num, img_id, :] = probabilities  # type: ignore
            labels.loc[int(img_id.cpu()), :] = label.cpu().numpy()[0, :]  # type: ignore

# Combine results and labels into a single Dataset
output_set = xr.Dataset({"predictions": results, "labels": labels})

# Save results to netcdf file
output_path = pl.Path(config["output"]["path"]) / "model_evaluation_results.nc"
output_set.to_netcdf(output_path, mode="w")  # type: ignore
print(f"Results saved to {output_path}")