pre-refactor 041426
This commit is contained in:
rpotter6298
@@ -0,0 +1,309 @@
|
||||
"""
|
||||
MD permutation feature importance for Phase 5 — logit_mlp_head checkpointed run.
|
||||
|
||||
For each of the 5 fold checkpoints:
|
||||
- loads test images + clinical metadata
|
||||
- caches image features (no grad)
|
||||
- permutes each clinical feature N times and measures AUC drop
|
||||
|
||||
Produces (in figures/explainability/):
|
||||
md_importance_phase5.png — aggregated bar chart across 5 folds
|
||||
md_importance_phase5.csv — mean/std per feature
|
||||
|
||||
Usage:
|
||||
python -m v3.scripts.output_analysis.explainability.permutation_importance_phase5
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import csv
|
||||
from pathlib import Path
|
||||
from types import SimpleNamespace
|
||||
|
||||
import matplotlib
|
||||
matplotlib.use("Agg")
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from sklearn.metrics import roc_auc_score
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||
CKPT_RUN = REPO_ROOT / "v3" / "results" / "phase5" / "logit_mlp_head_ckpt"
|
||||
FIGURES_ROOT = REPO_ROOT / "v3" / "figures" / "explainability"
|
||||
CLINICAL_DIR = REPO_ROOT / "Papila" / "ClinicalData"
|
||||
IMAGE_DIR = REPO_ROOT / "Papila" / "FundusImages"
|
||||
|
||||
BACKBONE = "resnet50"
|
||||
NUM_CLASSES = 2
|
||||
CD_HIDDEN = 128
|
||||
FUSION_DIM = 256
|
||||
N_PERMUTATIONS = 30
|
||||
SEED = 0
|
||||
|
||||
|
||||
# ── Model ─────────────────────────────────────────────────────────────────────
|
||||
|
||||
def build_model(ckpt_path: Path, device: torch.device):
|
||||
from v3.classes.models import SingleEyeHT
|
||||
sd = torch.load(ckpt_path, map_location="cpu")
|
||||
cd_in = sd["cd_tower.block0.0.weight"].shape[1]
|
||||
model = SingleEyeHT(
|
||||
backbone=BACKBONE, freeze_ratio=0.0, augment=False,
|
||||
clinical_data=SimpleNamespace(feature_dim=cd_in),
|
||||
num_classes=NUM_CLASSES, cd_hidden_dim=CD_HIDDEN, fusion_dim=FUSION_DIM,
|
||||
)
|
||||
model.load_state_dict(sd)
|
||||
model.to(device).eval()
|
||||
return model
|
||||
|
||||
|
||||
# ── Data ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
def build_data_bundle():
|
||||
from v3.classes.papila_builders import build_papila_data
|
||||
# Infer settings from available checkpoint to stay compatible.
|
||||
# Once the 10x5 run (--iop-drop-raw --exclude-cols Axial_Length) completes,
|
||||
# these will automatically match (feature_dim will drop from 25 → 21).
|
||||
ckpt = next(CKPT_RUN.glob("rep*/binary/ensemble/fold*/best_single.pt"), None)
|
||||
import torch as _t
|
||||
cd_in = _t.load(ckpt, map_location="cpu")["cd_tower.block0.0.weight"].shape[1] if ckpt else 25
|
||||
# cd_in=25 → old run (no iop_drop_raw, no excl); cd_in=21 → new run
|
||||
drop_raw = cd_in <= 21
|
||||
excl = ["Axial_Length"] if cd_in in (21, 23) else []
|
||||
return build_papila_data(
|
||||
image_dir=str(IMAGE_DIR), clinical_dir=str(CLINICAL_DIR),
|
||||
label_col="Diagnosis", cat_cols=["Gender", "Phakic/Pseudophakic"],
|
||||
iop_corr_method="ratio", iop_drop_raw=drop_raw, exclude_cols=excl,
|
||||
)
|
||||
|
||||
|
||||
def get_eval_transform():
|
||||
from torchvision import transforms
|
||||
return transforms.Compose([
|
||||
transforms.Resize(256), transforms.CenterCrop(224),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
|
||||
])
|
||||
|
||||
|
||||
def get_image_path(pid: int, eye: str) -> Path:
|
||||
return IMAGE_DIR / f"RET{pid:03d}{eye}.jpg"
|
||||
|
||||
|
||||
# ── Feature index map ─────────────────────────────────────────────────────────
|
||||
|
||||
def build_feature_index_map(data) -> dict[str, list[int]]:
|
||||
"""
|
||||
Map feature name → list of dimension indices in the vectorize_row output.
|
||||
Layout: [scalars (min-max scaled)] + [cat one-hots] + [scalar missing flags]
|
||||
"""
|
||||
n_scalar = len(data.scalar_cols)
|
||||
cat_expanded = sum(len(m) for m in data.cat_maps.values())
|
||||
feat_map: dict[str, list[int]] = {}
|
||||
|
||||
# Scalar: value dim + missing flag dim
|
||||
for i, col in enumerate(data.scalar_cols):
|
||||
feat_map[col] = [i, n_scalar + cat_expanded + i]
|
||||
|
||||
# Categorical: whole one-hot block
|
||||
cat_offset = n_scalar
|
||||
for col in data.cat_cols:
|
||||
n = len(data.cat_maps[col])
|
||||
feat_map[col] = list(range(cat_offset, cat_offset + n))
|
||||
cat_offset += n
|
||||
|
||||
return feat_map
|
||||
|
||||
|
||||
# ── Per-fold importance ───────────────────────────────────────────────────────
|
||||
|
||||
def run_fold(rep_idx: int, fold_idx: int, model, data, device: torch.device,
|
||||
n_permutations: int, seed: int) -> dict[str, tuple[float, float]]:
|
||||
"""
|
||||
Returns {feature_name: (mean_auc_drop, std_auc_drop)}.
|
||||
"""
|
||||
from v3.scripts.output_analysis.explainability.fold_patient_ids import (
|
||||
get_test_patient_ids,
|
||||
)
|
||||
transform = get_eval_transform()
|
||||
clinical = data.df
|
||||
|
||||
pids = get_test_patient_ids(rep_idx, fold_idx, clinical_dir=CLINICAL_DIR)
|
||||
|
||||
# Cache image features + build meta tensors + labels
|
||||
img_feats_list, meta_list, label_list = [], [], []
|
||||
model.eval()
|
||||
with torch.no_grad():
|
||||
for pid in pids:
|
||||
for eye in ("OD", "OS"):
|
||||
img_path = get_image_path(pid, eye)
|
||||
if not img_path.exists():
|
||||
continue
|
||||
row = clinical[
|
||||
(clinical["Patient ID"] == pid) & (clinical["eyeID"] == eye)
|
||||
]
|
||||
if len(row) == 0:
|
||||
continue
|
||||
row = row.iloc[0]
|
||||
label = int(row["Diagnosis"])
|
||||
|
||||
from PIL import Image
|
||||
pil = Image.open(img_path).convert("RGB")
|
||||
img_t = transform(pil).unsqueeze(0).to(device)
|
||||
feats = model.img_tower(img_t) # [1, img_dim]
|
||||
meta_vec = torch.tensor(data.vectorize_row(row),
|
||||
dtype=torch.float32).unsqueeze(0)
|
||||
|
||||
img_feats_list.append(feats.cpu())
|
||||
meta_list.append(meta_vec)
|
||||
label_list.append(label)
|
||||
|
||||
if not label_list or len(set(label_list)) < 2:
|
||||
print(f" fold{fold_idx}: insufficient data, skipping.")
|
||||
return {}
|
||||
|
||||
img_feats = torch.cat(img_feats_list).to(device) # [N, img_dim]
|
||||
meta_all = torch.cat(meta_list) # [N, feat_dim] on CPU
|
||||
y_true = np.array(label_list)
|
||||
|
||||
# Baseline AUC
|
||||
with torch.no_grad():
|
||||
md_feats = model.cd_tower(meta_all.to(device))
|
||||
out_f, _, _ = model.bridge(img_feats, md_feats)
|
||||
probs_base = F.softmax(out_f, dim=1)[:, 1].cpu().numpy()
|
||||
baseline_auc = roc_auc_score(y_true, probs_base)
|
||||
print(f" fold{fold_idx}: baseline AUC={baseline_auc:.4f} N={len(y_true)}")
|
||||
|
||||
feat_map = build_feature_index_map(data)
|
||||
rng = np.random.default_rng(seed + fold_idx)
|
||||
results: dict[str, tuple[float, float]] = {}
|
||||
|
||||
for feat_name, dims in feat_map.items():
|
||||
drops = []
|
||||
for _ in range(n_permutations):
|
||||
meta_perm = meta_all.clone()
|
||||
perm_idx = rng.permutation(len(meta_perm))
|
||||
meta_perm[:, dims] = meta_perm[perm_idx][:, dims]
|
||||
with torch.no_grad():
|
||||
md_p = model.cd_tower(meta_perm.to(device))
|
||||
out_p, _, _ = model.bridge(img_feats, md_p)
|
||||
probs_p = F.softmax(out_p, dim=1)[:, 1].cpu().numpy()
|
||||
try:
|
||||
drops.append(baseline_auc - roc_auc_score(y_true, probs_p))
|
||||
except Exception:
|
||||
pass
|
||||
if drops:
|
||||
results[feat_name] = (float(np.mean(drops)), float(np.std(drops)))
|
||||
|
||||
return results
|
||||
|
||||
|
||||
# ── Aggregate and plot ────────────────────────────────────────────────────────
|
||||
|
||||
def plot_importance(all_results: list[dict], out_png: Path, out_csv: Path) -> None:
|
||||
# Aggregate across folds
|
||||
all_feats = sorted({f for r in all_results for f in r})
|
||||
agg = {}
|
||||
for feat in all_feats:
|
||||
vals = [r[feat][0] for r in all_results if feat in r]
|
||||
if vals:
|
||||
agg[feat] = (float(np.mean(vals)), float(np.std(vals)))
|
||||
|
||||
# Sort by mean importance descending
|
||||
sorted_feats = sorted(agg, key=lambda f: agg[f][0], reverse=True)
|
||||
names = sorted_feats
|
||||
imps = [agg[f][0] for f in names]
|
||||
stds = [agg[f][1] for f in names]
|
||||
colors = ["#e05c5c" if v >= 0 else "#5c9ee0" for v in imps]
|
||||
|
||||
fig, ax = plt.subplots(figsize=(9, max(4, len(names) * 0.45 + 1.5)))
|
||||
y_pos = np.arange(len(names))
|
||||
ax.barh(y_pos, imps, xerr=stds, color=colors, ecolor="grey", capsize=3, height=0.6)
|
||||
ax.set_yticks(y_pos)
|
||||
ax.set_yticklabels(names, fontsize=9)
|
||||
ax.invert_yaxis()
|
||||
ax.axvline(0, color="black", linewidth=0.8)
|
||||
ax.set_xlabel("Mean AUC drop (baseline − permuted)", fontsize=10)
|
||||
n_reps = len(set(r for r in range(len(all_results)))) # placeholder
|
||||
ax.set_title(
|
||||
f"MD Tower — Permutation Feature Importance\n"
|
||||
f"Phase 5 logit_mlp_head_ckpt ({len(all_results)} folds, "
|
||||
f"error bars = std across folds)",
|
||||
fontsize=11,
|
||||
)
|
||||
fig.tight_layout()
|
||||
out_png.parent.mkdir(parents=True, exist_ok=True)
|
||||
fig.savefig(out_png, dpi=150, bbox_inches="tight")
|
||||
plt.close(fig)
|
||||
print(f"Saved: {out_png}")
|
||||
|
||||
with open(out_csv, "w", newline="") as f:
|
||||
w = csv.DictWriter(f, fieldnames=["feature", "mean_importance", "std_importance"])
|
||||
w.writeheader()
|
||||
for feat in sorted_feats:
|
||||
w.writerow({"feature": feat,
|
||||
"mean_importance": agg[feat][0],
|
||||
"std_importance": agg[feat][1]})
|
||||
print(f"Saved: {out_csv}")
|
||||
|
||||
|
||||
# ── Main ──────────────────────────────────────────────────────────────────────
|
||||
|
||||
def _discover_checkpoints(ckpt_run: Path) -> list[tuple[int, int, Path]]:
|
||||
found = []
|
||||
for rep_dir in sorted(ckpt_run.glob("rep*")):
|
||||
try:
|
||||
rep_idx = int(rep_dir.name.replace("rep", ""))
|
||||
except ValueError:
|
||||
continue
|
||||
for fold_dir in sorted((rep_dir / "binary" / "ensemble").glob("fold[0-9]")):
|
||||
ckpt = fold_dir / "best_single.pt"
|
||||
if ckpt.exists():
|
||||
found.append((rep_idx, int(fold_dir.name.replace("fold", "")), ckpt))
|
||||
return found
|
||||
|
||||
|
||||
def main(n_permutations: int = N_PERMUTATIONS, seed: int = SEED):
|
||||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||||
print(f"Device: {device}")
|
||||
|
||||
print("Building DataBundle ...")
|
||||
data = build_data_bundle()
|
||||
print(f" feature_dim={data.feature_dim}")
|
||||
|
||||
checkpoints = _discover_checkpoints(CKPT_RUN)
|
||||
print(f"Found {len(checkpoints)} checkpoint(s) across "
|
||||
f"{len(set(r for r,f,_ in checkpoints))} rep(s)")
|
||||
|
||||
if not checkpoints:
|
||||
print("No checkpoints found.")
|
||||
return
|
||||
|
||||
all_results = []
|
||||
for rep_idx, fold_idx, ckpt in checkpoints:
|
||||
print(f"\n── rep{rep_idx:02d} fold{fold_idx} ──")
|
||||
model = build_model(ckpt, device)
|
||||
result = run_fold(rep_idx, fold_idx, model, data, device, n_permutations, seed)
|
||||
if result:
|
||||
all_results.append(result)
|
||||
del model
|
||||
|
||||
if not all_results:
|
||||
print("No results — nothing to plot.")
|
||||
return
|
||||
|
||||
plot_importance(
|
||||
all_results,
|
||||
FIGURES_ROOT / "md_importance_phase5.png",
|
||||
FIGURES_ROOT / "md_importance_phase5.csv",
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
ap = argparse.ArgumentParser(description=__doc__,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--n-permutations", type=int, default=N_PERMUTATIONS)
|
||||
ap.add_argument("--seed", type=int, default=SEED)
|
||||
args = ap.parse_args()
|
||||
main(n_permutations=args.n_permutations, seed=args.seed)
|
||||
Reference in New Issue
Block a user