""" Two-layer MLP probe for ReMAP-PET 3-way classification. Compares linear vs non-linear probing on the same embeddings. """ from __future__ import annotations import argparse, json from pathlib import Path import numpy as np import pandas as pd import torch from sklearn.neural_network import MLPClassifier from sklearn.linear_model import LogisticRegression from sklearn.preprocessing import LabelEncoder, StandardScaler from sklearn.metrics import balanced_accuracy_score, roc_auc_score, accuracy_score from sklearn.pipeline import make_pipeline from torch.utils.data import DataLoader from pet_vlm_dataset import PETSUVRDataset, collate_pet_suvr from train_pet_foundation import PETSUVRFoundationModel, build_encoder def extract_embeddings(checkpoint, split_csv, output_size, device): ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False) saved_args = argparse.Namespace(**ckpt.get("args", {})) for name in ("backbone", "embed_dim", "freeze_encoder"): if not hasattr(saved_args, name) or getattr(saved_args, name) is None: setattr(saved_args, name, ckpt.get("args", {}).get(name)) embed_dim = getattr(saved_args, "embed_dim", None) or 256 backbone = getattr(saved_args, "backbone", "medicalnet") freeze = getattr(saved_args, "freeze_encoder", True) manifest = pd.read_csv(split_csv) dataset = PETSUVRDataset(split_csv, output_size=output_size) loader = DataLoader(dataset, batch_size=4, shuffle=False, num_workers=2, collate_fn=collate_pet_suvr) encoder = build_encoder(saved_args) model = PETSUVRFoundationModel(encoder, int(dataset[0]["suvr"].numel()), embed_dim, freeze).to(device) model.load_state_dict(ckpt["model"], strict=True) model.eval() pet_embs, suvr_preds = [], [] with torch.no_grad(): for batch in loader: image = batch["image"].to(device) suvr = batch["suvr"].to(device) pet_feat = model.pet_encoder(image) pet_z = torch.nn.functional.normalize(model.pet_projector(pet_feat), dim=-1) pet_embs.append(pet_z.cpu().numpy()) suvr_preds.append(model(image, suvr)["pred_suvr"].cpu().numpy()) pet_z = np.concatenate(pet_embs, axis=0) pred_suvr = np.concatenate(suvr_preds, axis=0) return manifest, pet_z, pred_suvr def main(): parser = argparse.ArgumentParser() 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("--output-size", type=int, nargs=3, default=(96, 96, 96)) args = parser.parse_args() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Loading checkpoint: {args.checkpoint}") train_df, train_pet, train_suvr = extract_embeddings(args.checkpoint, args.train, tuple(args.output_size), device) val_df, val_pet, val_suvr = extract_embeddings(args.checkpoint, args.val, tuple(args.output_size), device) test_df, test_pet, test_suvr = extract_embeddings(args.checkpoint, args.test, tuple(args.output_size), device) # Extract labels for 3-way classification encoder = LabelEncoder() all_labels = pd.concat([train_df["clinical_label"], val_df["clinical_label"], test_df["clinical_label"]]) encoder.fit(all_labels.astype(str)) y_train = encoder.transform(train_df["clinical_label"].astype(str)) y_val = encoder.transform(val_df["clinical_label"].astype(str)) y_test = encoder.transform(test_df["clinical_label"].astype(str)) n_classes = len(encoder.classes_) # ---- Linear probe (baseline) ---- print("\n=== Linear Probe (Logistic Regression) ===") best_bal, best_c = 0, 0.01 for c in [0.01, 0.03, 0.1, 0.3, 1.0, 3.0, 10.0, 30.0]: pipe = make_pipeline(StandardScaler(), LogisticRegression(C=c, max_iter=5000, solver="lbfgs")) pipe.fit(train_pet, y_train) val_pred = pipe.predict(val_pet) bal = balanced_accuracy_score(y_val, val_pred) print(f" C={c:.2f} val_bal={bal:.4f}") if bal > best_bal: best_bal, best_c = bal, c pipe = make_pipeline(StandardScaler(), LogisticRegression(C=best_c, max_iter=5000, solver="lbfgs")) pipe.fit(train_pet, y_train) test_pred = pipe.predict(test_pet) test_prob = pipe.predict_proba(test_pet) linear_bal = balanced_accuracy_score(y_test, test_pred) linear_auroc = roc_auc_score(y_test, test_prob, multi_class="ovr", average="macro") print(f" Best C={best_c}, Test BalAcc={linear_bal:.4f}, AUROC={linear_auroc:.4f}") # ---- MLP probe ---- print("\n=== MLP Probe (2-layer) ===") best_bal2, best_alpha = 0, 0.01 for alpha in [0.0001, 0.001, 0.01, 0.1, 1.0]: for hidden in [64, 128, 256]: pipe2 = make_pipeline( StandardScaler(), MLPClassifier(hidden_layer_sizes=(hidden, hidden // 2), alpha=alpha, max_iter=2000, early_stopping=True, validation_fraction=0.1, random_state=42) ) pipe2.fit(train_pet, y_train) val_pred2 = pipe2.predict(val_pet) bal2 = balanced_accuracy_score(y_val, val_pred2) if bal2 > best_bal2: best_bal2, best_alpha, best_hidden = bal2, alpha, hidden print(f" Best alpha={best_alpha}, hidden={best_hidden}, val_bal={best_bal2:.4f}") pipe2 = make_pipeline( StandardScaler(), MLPClassifier(hidden_layer_sizes=(best_hidden, best_hidden // 2), alpha=best_alpha, max_iter=5000, random_state=42) ) pipe2.fit(train_pet, y_train) test_pred2 = pipe2.predict(test_pet) test_prob2 = pipe2.predict_proba(test_pet) mlp_bal = balanced_accuracy_score(y_test, test_pred2) mlp_auroc = roc_auc_score(y_test, test_prob2, multi_class="ovr", average="macro") print(f" Test BalAcc={mlp_bal:.4f}, AUROC={mlp_auroc:.4f}") # ---- Combined features probe ---- print("\n=== Linear Probe (PET + Predicted SUVR) ===") train_comb = np.concatenate([train_pet, train_suvr], axis=1) val_comb = np.concatenate([val_pet, val_suvr], axis=1) test_comb = np.concatenate([test_pet, test_suvr], axis=1) best_bal3, best_c3 = 0, 0.01 for c in [0.01, 0.03, 0.1, 0.3, 1.0, 3.0, 10.0, 30.0]: pipe3 = make_pipeline(StandardScaler(), LogisticRegression(C=c, max_iter=5000, solver="lbfgs")) pipe3.fit(train_comb, y_train) val_pred3 = pipe3.predict(val_comb) bal3 = balanced_accuracy_score(y_val, val_pred3) if bal3 > best_bal3: best_bal3, best_c3 = bal3, c pipe3 = make_pipeline(StandardScaler(), LogisticRegression(C=best_c3, max_iter=5000, solver="lbfgs")) pipe3.fit(train_comb, y_train) test_pred3 = pipe3.predict(test_comb) test_prob3 = pipe3.predict_proba(test_comb) comb_bal = balanced_accuracy_score(y_test, test_pred3) comb_auroc = roc_auc_score(y_test, test_prob3, multi_class="ovr", average="macro") print(f" Best C={best_c3}, Test BalAcc={comb_bal:.4f}, AUROC={comb_auroc:.4f}") # Summary print(f"\n=== Summary (3-way test AUROC) ===") print(f" Linear probe (original): {linear_auroc:.4f}") print(f" MLP probe (2-layer): {mlp_auroc:.4f}") print(f" Linear probe (PET+SUVR): {comb_auroc:.4f}") print(f" MedicalNet frozen baseline: 0.7567") if __name__ == "__main__": main()