alnn_rewrite/tasks/load_data.py

import pathlib as pl
from typing import Any, Dict, List, Tuple

import pandas as pd

import model.dataset as ds
from util.progress import ProgressTracker
from util.seeding import seed_everything
from util.ui_logger import PipelineLogger


def seed_pipeline(config: Dict[str, Any], log: PipelineLogger) -> int:
    """Resolve the base seed (default 0) and seed all RNGs once. Returns it."""
    if config["data"].get("seed") is None:
        log.info("Seed is not defined - using default seed of 0")
        config["data"]["seed"] = 0
    base_seed = int(config["data"]["seed"])
    deterministic = bool(config.get("training", {}).get("deterministic", False))
    applied = seed_everything(base_seed, deterministic=deterministic)
    log.info(
        f"Seeded RNGs with base seed {applied}"
        + (" (deterministic cuDNN)" if deterministic else "")
    )
    return base_seed


def load_full_dataset(
    config: Dict[str, Any], log: PipelineLogger
) -> Tuple["ds.ADNIDataset", List[Tuple[int, str]], Dict[int, str]]:
    """Load the full ADNI dataset and its patient mapping (no split).

    Shared by the single-split loader and the k-fold driver, which partition the
    returned dataset differently.
    """
    log.info("Loading Files")
    mri_files = pl.Path(config["data"]["mri_files_path"]).glob("*.nii")
    xls_file = pl.Path(config["data"]["xls_file_path"])

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

    ptid_df = pd.read_csv(xls_file)
    ptid_df.columns = ptid_df.columns.str.strip()
    ptid_df = ptid_df[["Image Data ID", "PTID"]].dropna(  # type: ignore
        subset=["Image Data ID", "PTID"]
    )
    ptid_df["Image Data ID"] = ptid_df["Image Data ID"].astype(int)
    ptid_df["PTID"] = ptid_df["PTID"].astype(str).str.strip()
    ptid_df = ptid_df[ptid_df["PTID"] != ""]

    ptids = list(zip(ptid_df["Image Data ID"].tolist(), ptid_df["PTID"].tolist()))
    image_to_ptid = {int(iid): str(pid) for iid, pid in ptids}
    return dataset, ptids, image_to_ptid


def load_data_task(
    track: ProgressTracker,
    log: PipelineLogger,
    config: Dict[str, Any],
    state: Dict[str, Any],
):
    # Seed everything once at the start of the pipeline so data loading, weight
    # init, and shuffling are reproducible. Per-member seeds are derived from
    # this base inside the training tasks.
    base_seed = seed_pipeline(config, log)

    dataset, ptids, image_to_ptid = load_full_dataset(config, log)

    # Mapping used later by evaluation to attach patient ids to the sample axis.
    state["image_to_ptid"] = image_to_ptid

    # Split is grouped by PTID to prevent patient-level leakage across partitions.
    datasets = ds.divide_dataset_by_patient_id(
        dataset,
        ptids,
        config["data"]["data_splits"],
        seed=config["data"]["seed"],
    )

    # Initialize the dataloaders. The train shuffle generator is seeded from the
    # base seed so shuffling is reproducible; val/test are left unshuffled.
    train_loader, val_loader, test_loader = ds.initalize_dataloaders(
        datasets,
        batch_size=config["training"]["batch_size"],
        seed=base_seed,
    )

    log.info("Dataloaders initalized")
    state["train_loader"] = train_loader
    state["val_loader"] = val_loader
    state["test_loader"] = test_loader