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()