scripts/evaluate_pet_clinical_probe.py

from __future__ import annotations

import argparse
import csv
from pathlib import Path

import numpy as np
import pandas as pd
import torch
from sklearn.linear_model import LogisticRegression, Ridge
from sklearn.metrics import (
    accuracy_score,
    balanced_accuracy_score,
    f1_score,
    mean_absolute_error,
    mean_squared_error,
    r2_score,
    roc_auc_score,
)
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import LabelEncoder, StandardScaler, label_binarize
from torch.utils.data import DataLoader

from pet_vlm_dataset import PETSUVRDataset, collate_pet_suvr
from train_pet_foundation import PETSUVRFoundationModel, build_encoder


CLASSIFICATION_TASKS = {
    "clinical_label_3way": {"column": "clinical_label", "labels": ["CN", "MCI", "AD"]},
    "ad_vs_cn": {"column": "clinical_label", "labels": ["CN", "AD"]},
    "pmci_vs_smci": {"column": "conversion_label", "labels": ["SMCI", "PMCI"]},
}

REGRESSION_TASKS = {
    "adas11": "adas11",
    "mmse": "mmse",
    "ravlt_immediate": "ravlt_immediate",
    "ldeltotal": "ldeltotal",
}


def extract_pet_embeddings(
    model: PETSUVRFoundationModel,
    manifest: Path,
    output_size: tuple[int, int, int],
    batch_size: int,
    num_workers: int,
    device: torch.device,
) -> tuple[pd.DataFrame, np.ndarray]:
    dataset = PETSUVRDataset(manifest, output_size=output_size)
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=num_workers, collate_fn=collate_pet_suvr)
    features = []
    model.eval()
    with torch.no_grad():
        for batch in loader:
            image = batch["image"].to(device, non_blocking=True)
            pet_feat = model.pet_encoder(image)
            pet_z = torch.nn.functional.normalize(model.pet_projector(pet_feat), dim=-1)
            features.append(pet_z.cpu().numpy())
    return pd.read_csv(manifest), np.concatenate(features, axis=0)


def subset_classification(df: pd.DataFrame, x: np.ndarray, column: str, labels: list[str]) -> tuple[np.ndarray, np.ndarray]:
    mask = df[column].isin(labels).to_numpy()
    return x[mask], df.loc[mask, column].astype(str).to_numpy()


def subset_regression(df: pd.DataFrame, x: np.ndarray, column: str) -> tuple[np.ndarray, np.ndarray]:
    y = pd.to_numeric(df[column], errors="coerce")
    mask = y.notna().to_numpy()
    return x[mask], y.loc[mask].astype(float).to_numpy()


def evaluate_classification(name: str, x_train: np.ndarray, y_train: np.ndarray, x_val: np.ndarray, y_val: np.ndarray, x_test: np.ndarray, y_test: np.ndarray) -> dict[str, object]:
    encoder = LabelEncoder()
    encoder.fit(np.concatenate([y_train, y_val, y_test]))
    y_train_i = encoder.transform(y_train)
    y_val_i = encoder.transform(y_val)
    y_test_i = encoder.transform(y_test)

    best_c = None
    best_score = -np.inf
    best_model = None
    for c in [0.01, 0.03, 0.1, 0.3, 1.0, 3.0, 10.0]:
        model = make_pipeline(StandardScaler(), LogisticRegression(C=c, max_iter=5000, class_weight="balanced"))
        model.fit(x_train, y_train_i)
        score = balanced_accuracy_score(y_val_i, model.predict(x_val))
        if score > best_score:
            best_c = c
            best_score = score
            best_model = model

    assert best_model is not None
    pred = best_model.predict(x_test)
    proba = best_model.predict_proba(x_test)
    result: dict[str, object] = {
        "task": name,
        "type": "classification",
        "n_train": len(y_train),
        "n_val": len(y_val),
        "n_test": len(y_test),
        "selected_param": best_c,
        "accuracy": accuracy_score(y_test_i, pred),
        "balanced_accuracy": balanced_accuracy_score(y_test_i, pred),
        "macro_f1": f1_score(y_test_i, pred, average="macro"),
    }
    if len(encoder.classes_) == 2:
        result["auroc"] = roc_auc_score(y_test_i, proba[:, 1])
    else:
        y_bin = label_binarize(y_test_i, classes=np.arange(len(encoder.classes_)))
        result["auroc"] = roc_auc_score(y_bin, proba, average="macro", multi_class="ovr")
    return result


def evaluate_regression(name: str, x_train: np.ndarray, y_train: np.ndarray, x_val: np.ndarray, y_val: np.ndarray, x_test: np.ndarray, y_test: np.ndarray) -> dict[str, object]:
    best_alpha = None
    best_rmse = np.inf
    best_model = None
    for alpha in [0.01, 0.03, 0.1, 0.3, 1.0, 3.0, 10.0, 30.0, 100.0]:
        model = make_pipeline(StandardScaler(), Ridge(alpha=alpha))
        model.fit(x_train, y_train)
        rmse = mean_squared_error(y_val, model.predict(x_val)) ** 0.5
        if rmse < best_rmse:
            best_alpha = alpha
            best_rmse = rmse
            best_model = model

    assert best_model is not None
    pred = best_model.predict(x_test)
    return {
        "task": name,
        "type": "regression",
        "n_train": len(y_train),
        "n_val": len(y_val),
        "n_test": len(y_test),
        "selected_param": best_alpha,
        "mae": mean_absolute_error(y_test, pred),
        "rmse": mean_squared_error(y_test, pred) ** 0.5,
        "r2": r2_score(y_test, pred),
        "pearson": float(np.corrcoef(y_test, pred)[0, 1]) if np.std(pred) > 1e-8 else np.nan,
    }


def main() -> None:
    parser = argparse.ArgumentParser(description="Run clinical downstream probes from PET checkpoint embeddings.")
    parser.add_argument("--checkpoint", type=Path, required=True)
    parser.add_argument("--train", type=Path, default=Path("data/metadata/splits/train_clinical.csv"))
    parser.add_argument("--val", type=Path, default=Path("data/metadata/splits/val_clinical.csv"))
    parser.add_argument("--test", type=Path, default=Path("data/metadata/splits/test_clinical.csv"))
    parser.add_argument("--backbone", choices=["small_cnn", "medicalnet", "brainiac", "brainfm", "swinunetr", "sam_med3d"], default=None)
    parser.add_argument("--medicalnet-weights", type=Path, default=Path("pretrained/medicalnet/resnet_50_23dataset.pth"))
    parser.add_argument("--brainiac-weights", type=Path, default=Path("pretrained/brainiac/backbone.safetensors"))
    parser.add_argument("--brainfm-weights", type=Path, default=Path("pretrained/brainfm/assets/brainfm_pretrained.pth"))
    parser.add_argument("--brainfm-code-root", type=Path, default=Path("pretrained/brainfm"))
    parser.add_argument("--swinunetr-weights", type=Path, default=Path("pretrained/swinunetr/model_swinvit.pt"))
    parser.add_argument("--sam-med3d-weights", type=Path, default=Path("pretrained/sam-med3d/sam_med3d_turbo.pth"))
    parser.add_argument("--batch-size", type=int, default=4)
    parser.add_argument("--num-workers", type=int, default=2)
    parser.add_argument("--output-size", type=int, nargs=3, default=None)
    parser.add_argument("--embed-dim", type=int, default=None)
    parser.add_argument("--freeze-encoder", action=argparse.BooleanOptionalAction, default=None)
    parser.add_argument("--out", type=Path, default=Path("runs/clinical/pet_clinical_probe.csv"))
    args = parser.parse_args()

    ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False)
    saved_args = ckpt.get("args", {})
    for name in ("backbone", "embed_dim", "freeze_encoder"):
        if getattr(args, name, None) is None and name in saved_args:
            setattr(args, name, saved_args[name])
    if args.output_size is None:
        args.output_size = tuple(saved_args.get("output_size", (96, 96, 96)))

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    sample_dataset = PETSUVRDataset(args.train, output_size=tuple(args.output_size))
    n_regions = int(sample_dataset[0]["suvr"].numel())
    encoder = build_encoder(args)
    model = PETSUVRFoundationModel(encoder, n_regions, args.embed_dim or 256, bool(args.freeze_encoder)).to(device)
    model.load_state_dict(ckpt["model"], strict=True)

    train_df, x_train_all = extract_pet_embeddings(model, args.train, tuple(args.output_size), args.batch_size, args.num_workers, device)
    val_df, x_val_all = extract_pet_embeddings(model, args.val, tuple(args.output_size), args.batch_size, args.num_workers, device)
    test_df, x_test_all = extract_pet_embeddings(model, args.test, tuple(args.output_size), args.batch_size, args.num_workers, device)

    results = []
    for name, spec in CLASSIFICATION_TASKS.items():
        x_train, y_train = subset_classification(train_df, x_train_all, spec["column"], spec["labels"])
        x_val, y_val = subset_classification(val_df, x_val_all, spec["column"], spec["labels"])
        x_test, y_test = subset_classification(test_df, x_test_all, spec["column"], spec["labels"])
        if min(len(y_train), len(y_val), len(y_test)) > 0:
            results.append(evaluate_classification(name, x_train, y_train, x_val, y_val, x_test, y_test))

    for name, column in REGRESSION_TASKS.items():
        x_train, y_train = subset_regression(train_df, x_train_all, column)
        x_val, y_val = subset_regression(val_df, x_val_all, column)
        x_test, y_test = subset_regression(test_df, x_test_all, column)
        if min(len(y_train), len(y_val), len(y_test)) > 0:
            results.append(evaluate_regression(name, x_train, y_train, x_val, y_val, x_test, y_test))

    args.out.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = ["checkpoint", "backbone", *sorted({key for row in results for key in row.keys()})]
    with args.out.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()
        for row in results:
            writer.writerow({"checkpoint": str(args.checkpoint), "backbone": args.backbone, **row})

    print(f"checkpoint={args.checkpoint}")
    print(f"wrote={args.out}")
    for row in results:
        print(row)


if __name__ == "__main__":
    main()