| """ |
| 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) |
|
|
| |
| 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_) |
|
|
| |
| 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}") |
|
|
| |
| 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}") |
|
|
| |
| 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}") |
|
|
| |
| 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() |
|
|
