ryan/hypertower
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

pre-restructure

Browse Source
This commit is contained in:
rpotter6298
2026-02-26 06:47:58 +01:00
parent 17a255daa8
commit 8980cf5f9b
10 changed files with 977 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1 -1
View File
@@ -5,7 +5,7 @@ Papila/*
REFUGE/*
old_**/
models/*
cache_data/
# model artifacts
models/refuge/
models/v2/refuge/
classes/refuge_classification.py
+30
View File
@@ -401,6 +401,7 @@ class RefugeClassification:
f"[classifier] Building datasets from {len(candidates)} labelled samples (train/val)"
)
skipped: List[str] = []
for sample in tqdm(
candidates,
desc="Preparing records",
@@ -410,6 +411,7 @@ class RefugeClassification:
try:
geom, disc_mask, cup_mask = self._resolve_geometry(sample, crop_scale)
except RuntimeError:
skipped.append(sample.sample_id)
continue
record = RefugeClassificationRecord(
sample=sample,
@@ -424,6 +426,12 @@ class RefugeClassification:
else:
val_records.append(record)
if skipped:
print(
f"[classifier] WARNING: {len(skipped)}/{len(candidates)} samples skipped "
f"due to empty segmentation mask: {skipped}"
)
if (not val_records or self.use_all_labeled) and train_records and self.auto_val_ratio > 0.0:
rng = random.Random(42)
label_groups: Dict[int, List[RefugeClassificationRecord]] = {}
@@ -578,17 +586,20 @@ class RefugeClassification:
) -> List[RefugeClassificationRecord]:
scale = crop_scale if crop_scale is not None else self.crop_scale
records: List[RefugeClassificationRecord] = []
skipped: List[str] = []
iterator: Iterable[RefugeSample]
if progress_prefix is not None:
iterator = tqdm(samples, desc=progress_prefix, unit="sample", leave=False)
else:
iterator = samples
labeled = [s for s in samples if s.label is not None]
for sample in iterator:
if sample.label is None:
continue
try:
geom, disc_mask, cup_mask = self._resolve_geometry(sample, scale)
except RuntimeError:
skipped.append(sample.sample_id)
continue
records.append(
RefugeClassificationRecord(
@@ -598,8 +609,27 @@ class RefugeClassification:
cup_mask=cup_mask,
)
)
prefix = f"[{progress_prefix}]" if progress_prefix else "[classifier]"
if skipped:
print(
f"{prefix} WARNING: {len(skipped)}/{len(labeled)} samples skipped "
f"due to empty segmentation mask: {skipped}"
)
else:
print(f"{prefix} All {len(labeled)} samples processed successfully.")
return records
def clear_disk_cache(self) -> None:
"""Delete all cached geometry/mask .npz files in cache_dir."""
if self.cache_dir is None or not self.cache_dir.exists():
return
removed = 0
for f in self.cache_dir.glob("*.npz"):
f.unlink()
removed += 1
self.geometry_cache.clear()
print(f"[classifier] Cleared {removed} cached geometry files from {self.cache_dir}")
def _cache_key(self, sample_id: str, scale: float) -> str:
scale_tag = int(round(scale * 100))
return f"{sample_id}_s{scale_tag}"
classes/unet_segmenter.py
+7 -4
View File
@@ -448,10 +448,13 @@ class UNetSegmenter:
vals = np.where(counts > 0)[0]
if vals.size < 1:
raise ValueError(f"Mask {mask_path} does not contain discernible labels")
vals = vals[np.argsort(-counts[vals])]
disc_val = int(vals[0])
cup_val = int(vals[1]) if vals.size > 1 else None
disc_mask = (arr == disc_val).astype(np.uint8)
# Disc = ALL non-background pixels (full optic disc: rim + cup combined).
# Previously this was rim-only, which caused the cup structural prior
# (cup & disc) to produce empty cup masks since cup and rim don't overlap.
disc_mask = (arr != bg_val).astype(np.uint8)
# Cup = the darkest non-background value (0 in REFUGE = inner cup region).
# Using min-value rather than frequency avoids swapping when cup area > rim area.
cup_val = int(np.min(vals)) if vals.size > 1 else None
cup_mask = (arr == cup_val).astype(np.uint8) if cup_val is not None else np.zeros_like(disc_mask, dtype=np.uint8)
return disc_mask, cup_mask if cup_mask.any() else None, (w, h)
classes/v2/v2_hypertower.py
+179 -18
View File
@@ -143,6 +143,12 @@ class UNetImageCropper:
stem = image_path.stem
return self.cache_dir / f"{stem}_s{int(self.scale * 100)}.npz"
def clear_cache(self) -> None:
if self.cache_dir is None or not self.cache_dir.exists():
return
removed = sum(1 for f in self.cache_dir.glob("*.npz") if f.unlink() or True)
print(f"[UNetImageCropper] Cleared {removed} cached crop files from {self.cache_dir}")
def _infer_masks(self, image: Image.Image) -> Optional[Tuple[np.ndarray, np.ndarray]]:
resized = self.segmenter.preprocess_image(image)
tensor = self.to_tensor(resized).unsqueeze(0).to(self.segmenter.device)
@@ -281,6 +287,12 @@ class ManifestImageCropper:
return None
return self.cache_dir / f"{image_path.stem}_s{int(self.scale * 100)}.npz"
def clear_cache(self) -> None:
if self.cache_dir is None or not self.cache_dir.exists():
return
removed = sum(1 for f in self.cache_dir.glob("*.npz") if f.unlink() or True)
print(f"[ManifestImageCropper] Cleared {removed} cached crop files from {self.cache_dir}")
@staticmethod
def _load_contour(path: Path) -> np.ndarray:
coords = np.loadtxt(path)
@@ -512,8 +524,9 @@ class HyperTower:
crop_weights = getattr(args, "img_crop_weights", None)
use_gt = getattr(args, "img_crop_gt", False)
if crop_manifest:
crop_cache = getattr(args, "img_crop_cache", Path("analysis_data/hypertower_crops"))
crop_cache = getattr(args, "img_crop_cache", Path("cache_data/hypertower_crops"))
crop_cache = Path(crop_cache)
persist_cache = bool(getattr(args, "persist_img_crop_cache", False))
if use_gt:
self.image_preprocessor = ManifestImageCropper(
manifest_path=Path(crop_manifest),
@@ -521,6 +534,8 @@ class HyperTower:
target_size=getattr(args, "img_crop_size", 224),
cache_dir=crop_cache,
)
if not persist_cache:
self.image_preprocessor.clear_cache()
print(f"[HyperTower] GT disc cropper enabled → cache at {crop_cache}")
elif crop_weights:
self.image_preprocessor = UNetImageCropper(
@@ -533,6 +548,8 @@ class HyperTower:
target_size=getattr(args, "img_crop_size", 224),
cache_dir=crop_cache,
)
if not persist_cache:
self.image_preprocessor.clear_cache()
print(f"[HyperTower] UNet disc cropper enabled → cache at {crop_cache}")
else:
print("[HyperTower] img_crop_manifest provided but no weights/gt flag; skipping cropping")
@@ -2254,7 +2271,8 @@ def build_image_preprocessor_from_args(args):
use_gt = bool(getattr(args, "img_crop_gt", False))
if not crop_manifest:
return None
crop_cache = Path(getattr(args, "img_crop_cache", Path("analysis_data/hypertower_crops")))
crop_cache = Path(getattr(args, "img_crop_cache", Path("cache_data/hypertower_crops")))
persist_cache = bool(getattr(args, "persist_img_crop_cache", False))
if use_gt:
pre = ManifestImageCropper(
manifest_path=Path(crop_manifest),
@@ -2262,6 +2280,8 @@ def build_image_preprocessor_from_args(args):
target_size=getattr(args, "img_crop_size", 224),
cache_dir=crop_cache,
)
if not persist_cache:
pre.clear_cache()
print(f"[V2 modes] GT disc cropper enabled -> cache at {crop_cache}", flush=True)
return pre
if crop_weights:
@@ -2275,6 +2295,8 @@ def build_image_preprocessor_from_args(args):
target_size=getattr(args, "img_crop_size", 224),
cache_dir=crop_cache,
)
if not persist_cache:
pre.clear_cache()
print(f"[V2 modes] UNet disc cropper enabled -> cache at {crop_cache}", flush=True)
return pre
print(
@@ -2686,6 +2708,14 @@ class FoldResult:
bilat_val_threshold: float
bilat_val_bias: Optional[str]
bilat_val_n: int
# Holdout metrics (evaluated at the same epoch as best val; nan if no holdout)
classic_holdout_auc: float
classic_holdout_acc: float
ensemble_holdout_auc: float
ensemble_holdout_acc: float
bilat_holdout_auc: float
bilat_holdout_acc: float
holdout_n: int # number of holdout bilateral samples
# Training sample counts
single_train_n: int
bilat_train_n: int
@@ -2800,6 +2830,10 @@ def run_fold(
bilat_val_mcc=nan, bilat_val_f1=nan, bilat_val_recall=None,
bilat_val_ece=nan, bilat_val_threshold=nan, bilat_val_bias=None,
bilat_val_n=0,
classic_holdout_auc=nan, classic_holdout_acc=nan,
ensemble_holdout_auc=nan, ensemble_holdout_acc=nan,
bilat_holdout_auc=nan, bilat_holdout_acc=nan,
holdout_n=0,
single_train_n=len(eye_train), bilat_train_n=len(bilat_train),
)
return empty, FoldArtifacts(
@@ -2861,6 +2895,23 @@ def run_fold(
**loader_kw,
)
# ---- holdout loader (if holdout patients are available) ------------------
holdout_bilat: list = []
holdout_loader = None
if split.holdout is not None and not split.holdout.empty:
holdout_bilat = filter_bilateral_samples(
profile_patient.build_samples(df=split.holdout, clinical=data)
)
if holdout_bilat:
holdout_loader = make_loader(
holdout_bilat, slots_patient,
image_transform=eval_transform,
image_preprocessor=image_preprocessor,
shuffle=False,
**loader_kw,
)
print(f" [fold {fold+1}] holdout_n={len(holdout_bilat)} (bilateral patients)", flush=True)
opt_single = torch.optim.Adam(single.parameters(), lr=args.lr) if run_single else None
opt_bilateral = torch.optim.Adam(bilateral.parameters(), lr=args.lr) if run_bilat else None
@@ -2871,11 +2922,15 @@ def run_fold(
"main_epoch_single", "main_epoch_bilat",
"single_active", "bilat_active",
"single_train_loss", "single_train_acc",
"classic_val_auc", "classic_val_acc", "classic_val_n",
"ensemble_val_auc", "ensemble_val_acc", "ensemble_val_n",
"bilat_train_loss", "bilat_train_acc",
"bilat_val_auc", "bilat_val_acc", "bilat_val_n",
"classic_val_auc", "classic_val_acc", "classic_val_n",
"ensemble_val_auc", "ensemble_val_acc", "ensemble_val_n",
"bilat_train_loss", "bilat_train_acc",
"bilat_val_auc", "bilat_val_acc", "bilat_val_n",
"classic_holdout_auc", "classic_holdout_acc",
"ensemble_holdout_auc", "ensemble_holdout_acc",
"bilat_holdout_auc", "bilat_holdout_acc",
"is_best_single", "is_best_bilat",
"is_best_holdout_single", "is_best_holdout_bilat",
]
fold_logger = HypertowerLogger(run_dir=fold_dir)
@@ -2889,6 +2944,16 @@ def run_fold(
snap_classic: dict = {}
snap_ensemble: dict = {}
snap_bilat: dict = {}
# holdout snaps (metrics captured at the same epoch as best val)
snap_holdout_single: dict = {}
snap_holdout_bilat: dict = {}
# separate best-holdout trackers (for checkpointing)
best_holdout_single_auc = -1.0
best_holdout_bilat_auc = -1.0
best_epoch_holdout_single = 0
best_epoch_holdout_bilat = 0
best_holdout_single_state: Optional[dict] = None
best_holdout_bilat_state: Optional[dict] = None
if run_single:
print(
@@ -3011,8 +3076,36 @@ def run_fold(
bi_n = 0
bi_acc_img = bi_acc_md = bi_auc_img = bi_auc_md = nan
# --- holdout evaluation -----------------------------------------------
if holdout_loader is not None:
if run_single and tower_mode == "single":
y_cl_h, p_cl_h, _, _ = collect_probs_single_components(
single, holdout_loader, device, aggregate_patient=False
)
_, cl_auc_h, _ = _score_arrays(y_cl_h, p_cl_h, num_classes)
cl_acc_h = float((p_cl_h.argmax(1) == y_cl_h).mean()) if y_cl_h.size else nan
en_auc_h = en_acc_h = nan
elif run_single and tower_mode == "ensemble":
y_en_h, p_en_h, _, _ = collect_probs_single_components(
single, holdout_loader, device, aggregate_patient=True
)
_, en_auc_h, _ = _score_arrays(y_en_h, p_en_h, num_classes)
en_acc_h = float((p_en_h.argmax(1) == y_en_h).mean()) if y_en_h.size else nan
cl_auc_h = cl_acc_h = nan
else:
cl_auc_h = cl_acc_h = en_auc_h = en_acc_h = nan
if run_bilat:
y_bi_h, p_bi_h, _, _ = collect_probs_bilateral_components(bilateral, holdout_loader, device)
_, bi_auc_h, _ = _score_arrays(y_bi_h, p_bi_h, num_classes)
bi_acc_h = float((p_bi_h.argmax(1) == y_bi_h).mean()) if y_bi_h.size else nan
else:
bi_auc_h = bi_acc_h = nan
else:
cl_auc_h = cl_acc_h = en_auc_h = en_acc_h = bi_auc_h = bi_acc_h = nan
# Best-epoch checks: checkpointing is restricted to the main phase only.
target_single_auc = cl_auc if tower_mode == "single" else en_auc
target_holdout_single_auc = cl_auc_h if tower_mode == "single" else en_auc_h
single_ckpt_eligible = run_single and (phase_single == "main")
is_best_single = (
single_ckpt_eligible
@@ -3029,6 +3122,19 @@ def run_fold(
else:
snap_en, _, _, _ = _tune_and_snap(y_en, p_en, en_acc, num_classes, args, args.ece_bins)
snap_ensemble = snap_en
# capture holdout metrics at this val-best epoch
snap_holdout_single = {"auc": float(target_holdout_single_auc), "acc": float(cl_acc_h if tower_mode == "single" else en_acc_h)}
is_best_holdout_single = (
holdout_loader is not None
and single_ckpt_eligible
and (not np.isnan(target_holdout_single_auc))
and (target_holdout_single_auc > best_holdout_single_auc)
)
if is_best_holdout_single:
best_holdout_single_auc = target_holdout_single_auc
best_epoch_holdout_single = epoch + 1
best_holdout_single_state = copy.deepcopy(single.state_dict())
bilat_ckpt_eligible = run_bilat and (phase_bilat == "main")
is_best_bilat = (
@@ -3042,6 +3148,19 @@ def run_fold(
best_bilat_state = copy.deepcopy(bilateral.state_dict())
snap_bi, _, _, _ = _tune_and_snap(y_bi, p_bi, bi_acc, num_classes, args, args.ece_bins)
snap_bilat = snap_bi
# capture holdout metrics at this val-best epoch
snap_holdout_bilat = {"auc": float(bi_auc_h), "acc": float(bi_acc_h)}
is_best_holdout_bilat = (
holdout_loader is not None
and bilat_ckpt_eligible
and (not np.isnan(bi_auc_h))
and (bi_auc_h > best_holdout_bilat_auc)
)
if is_best_holdout_bilat:
best_holdout_bilat_auc = bi_auc_h
best_epoch_holdout_bilat = epoch + 1
best_holdout_bilat_state = copy.deepcopy(bilateral.state_dict())
fold_logger.write_epoch_row({
"fold": fold, "epoch": epoch + 1,
@@ -3052,15 +3171,22 @@ def run_fold(
"single_active": int(single_active),
"bilat_active": int(bilat_active),
"single_train_loss": _f(sl_loss), "single_train_acc": _f(sl_acc),
"classic_val_auc": _f(cl_auc), "classic_val_acc": _f(cl_acc), "classic_val_n": cl_n,
"ensemble_val_auc": _f(en_auc), "ensemble_val_acc": _f(en_acc), "ensemble_val_n": en_n,
"bilat_train_loss": _f(bl_loss), "bilat_train_acc": _f(bl_acc),
"bilat_val_auc": _f(bi_auc), "bilat_val_acc": _f(bi_acc), "bilat_val_n": bi_n,
"is_best_single": int(is_best_single),
"is_best_bilat": int(is_best_bilat),
"classic_val_auc": _f(cl_auc), "classic_val_acc": _f(cl_acc), "classic_val_n": cl_n,
"ensemble_val_auc": _f(en_auc), "ensemble_val_acc": _f(en_acc), "ensemble_val_n": en_n,
"bilat_train_loss": _f(bl_loss), "bilat_train_acc": _f(bl_acc),
"bilat_val_auc": _f(bi_auc), "bilat_val_acc": _f(bi_acc), "bilat_val_n": bi_n,
"classic_holdout_auc": _f(cl_auc_h), "classic_holdout_acc": _f(cl_acc_h),
"ensemble_holdout_auc": _f(en_auc_h), "ensemble_holdout_acc": _f(en_acc_h),
"bilat_holdout_auc": _f(bi_auc_h), "bilat_holdout_acc": _f(bi_acc_h),
"is_best_single": int(is_best_single),
"is_best_bilat": int(is_best_bilat),
"is_best_holdout_single": int(is_best_holdout_single),
"is_best_holdout_bilat": int(is_best_holdout_bilat),
}, optional_cols=epoch_fields)
if args.log_every > 0 and (epoch + 1) % args.log_every == 0:
hld_auc = target_holdout_single_auc if run_single else bi_auc_h
hld_suffix = f" hld_auc={hld_auc:.4f}" if holdout_loader is not None else ""
if run_single:
if tower_mode == "single":
msg = (
@@ -3070,6 +3196,7 @@ def run_fold(
f"img(acc={cl_acc_img:.4f},auc={cl_auc_img:.4f}) "
f"md(acc={cl_acc_md:.4f},auc={cl_auc_md:.4f}) "
f"(best_fused={best_single_auc:.4f} @ep{best_epoch_single})"
f"{hld_suffix}"
)
else:
msg = (
@@ -3079,6 +3206,7 @@ def run_fold(
f"img(acc={en_acc_img:.4f},auc={en_auc_img:.4f}) "
f"md(acc={en_acc_md:.4f},auc={en_auc_md:.4f}) "
f"(best_fused={best_single_auc:.4f} @ep{best_epoch_single})"
f"{hld_suffix}"
)
else:
msg = (
@@ -3088,6 +3216,7 @@ def run_fold(
f"img(acc={bi_acc_img:.4f},auc={bi_auc_img:.4f}) "
f"md(acc={bi_acc_md:.4f},auc={bi_auc_md:.4f}) "
f"(best_bilat={best_bilat_auc:.4f} @ep{best_epoch_bilat})"
f"{hld_suffix}"
)
print(msg, flush=True)
fold_logger.info(msg)
@@ -3099,6 +3228,10 @@ def run_fold(
torch.save(best_single_state, fold_dir / "best_single.pt")
if best_bilat_state is not None:
torch.save(best_bilat_state, fold_dir / "best_bilateral.pt")
if best_holdout_single_state is not None:
torch.save(best_holdout_single_state, fold_dir / "best_holdout_single.pt")
if best_holdout_bilat_state is not None:
torch.save(best_holdout_bilat_state, fold_dir / "best_holdout_bilateral.pt")
if run_single:
if tower_mode == "single":
@@ -3181,6 +3314,13 @@ def run_fold(
bilat_val_threshold=snap_bilat.get("threshold", nan),
bilat_val_bias=_svf(snap_bilat.get("bias")),
bilat_val_n=snap_bilat.get("n", 0),
classic_holdout_auc=snap_holdout_single.get("auc", nan) if tower_mode == "single" else nan,
classic_holdout_acc=snap_holdout_single.get("acc", nan) if tower_mode == "single" else nan,
ensemble_holdout_auc=snap_holdout_single.get("auc", nan) if tower_mode == "ensemble" else nan,
ensemble_holdout_acc=snap_holdout_single.get("acc", nan) if tower_mode == "ensemble" else nan,
bilat_holdout_auc=snap_holdout_bilat.get("auc", nan),
bilat_holdout_acc=snap_holdout_bilat.get("acc", nan),
holdout_n=len(holdout_bilat),
single_train_n=len(eye_train),
bilat_train_n=len(bilat_train),
), FoldArtifacts(
@@ -3217,6 +3357,20 @@ def _summary(results: list[FoldResult]) -> dict:
sub[f"{m}_std"] = std
out[label] = sub
for label, prefix in [
("classic_holdout", "classic_holdout"),
("ensemble_holdout", "ensemble_holdout"),
("bilat_holdout", "bilat_holdout"),
]:
sub = {}
for m in ["auc", "acc"]:
vals = [getattr(r, f"{prefix}_{m}") for r in results]
mean, std = _ms(vals)
sub[f"{m}_mean"] = mean
if m == "auc":
sub[f"{m}_std"] = std
out[label] = sub
# Deltas: ensemble − classic (eval strategy effect, same model)
# bilateral − ensemble (bilateral training effect)
for delta_label, prefix_a, prefix_b in [
@@ -3310,8 +3464,12 @@ def build_parser() -> argparse.ArgumentParser:
default="single",
help="Train/evaluate a single tower mode.",
)
ap.add_argument("--n-splits", type=int, default=5)
ap.add_argument("--fold-seed", type=int, default=42)
ap.add_argument("--n-splits", type=int, default=5)
ap.add_argument("--fold-seed", type=int, default=42)
ap.add_argument("--holdout-per-class", type=int, default=5,
help="Patients per class reserved for holdout before train/test split (0 disables)")
ap.add_argument("--holdout-seed", type=int, default=123,
help="Random seed for holdout sampling")
ap.add_argument(
"--folds",
type=int,
@@ -3365,8 +3523,10 @@ def build_parser() -> argparse.ArgumentParser:
help="Disc-radius multiplier for square crop.")
ap.add_argument("--img-crop-size", type=int, default=224,
help="Output ROI size before tower transforms.")
ap.add_argument("--img-crop-cache", type=str, default="analysis_data/hypertower_crops",
ap.add_argument("--img-crop-cache", type=str, default="cache_data/hypertower_crops",
help="Cache directory for cropped images and geometry sidecars.")
ap.add_argument("--persist-img-crop-cache", action="store_true",
help="Keep existing cached crop .npz files instead of clearing at run start.")
# Architecture
ap.add_argument("--md-hidden-dim", type=int, default=128,
help="MDTower hidden dimension.")
@@ -3407,7 +3567,8 @@ def build_parser() -> argparse.ArgumentParser:
)
ap.add_argument("--ece-bins", type=int, default=10)
ap.add_argument("--log-every", type=int, default=1)
ap.add_argument("--save-checkpoints", action="store_true")
ap.add_argument("--save-checkpoints", action=argparse.BooleanOptionalAction, default=True,
help="Save best_single.pt / best_holdout_single.pt per fold (use --no-save-checkpoints to disable)")
return ap
@@ -3498,8 +3659,8 @@ def run_mode(args) -> Path:
)
split_args = SimpleNamespace(
eval_mode=mode,
holdout_per_class=0,
holdout_seed=123,
holdout_per_class=args.holdout_per_class,
holdout_seed=args.holdout_seed,
n_splits=args.n_splits,
fold_seed=args.fold_seed,
)
models/v2/refuge/segmentation/per_image_refuge_build/best.pt
BIN
View File
Binary file not shown.
scripts/basic_analysis/compare_hypertower_modes.py
+5
View File
@@ -36,12 +36,17 @@ def parse_args():
def main():
seq_args, remaining = parse_args()
base_parser = build_parser()
first_run = True
for eval_mode in seq_args.eval_modes:
for tower_mode in seq_args.tower_modes:
tower_mode = "single" if tower_mode == "classic" else tower_mode
cli = list(remaining) + ["--eval-mode", eval_mode, "--tower-mode", tower_mode]
# Clear cache only on the first run; reuse it for all subsequent runs.
if not first_run:
cli.append("--persist-img-crop-cache")
args = base_parser.parse_args(cli)
run_mode(args)
first_run = False
if __name__ == "__main__":
scripts/main/refuge/refuge_build.py
+30 -40
View File
@@ -2,7 +2,6 @@
Usage examples (after activating .venv_refuge):
python refuge_build.py --train-seg
python refuge_build.py --train-clf
python refuge_build.py --eval --with-ttt
@@ -297,6 +296,8 @@ def build_papila_records(
papila_clf.backbone.to(args.device)
papila_clf.classifier_head.to(args.device)
papila_clf.rotation_head.to(args.device)
if getattr(args, "clear_clf_cache", False):
papila_clf.clear_disk_cache()
records = papila_clf.build_records_for_samples(
samples, crop_scale=args.crop_scale, progress_prefix="papila"
@@ -305,24 +306,6 @@ def build_papila_records(
return records, papila_clf
def train_segmentation(args: argparse.Namespace) -> None:
pre = ensure_preprocessing()
seg = RefugeSegmentation(pre)
seg.build_datasets(
image_size=args.seg_image_size,
batch_size=args.seg_batch_size,
num_workers=args.num_workers,
)
history = seg.train(
epochs=args.seg_epochs,
lr=args.seg_lr,
weight_decay=args.seg_weight_decay,
checkpoint_dir=SEG_CKPT.parent,
device=args.device,
)
print("Segmentation training complete. Best Dice:", history.get("best_dice"))
def train_unet_segmenter(args: argparse.Namespace) -> None:
manifest_path = args.seg_manifest or Path("manifest.csv")
mask_cache_dir = None if args.in_memory_cache else args.mask_cache_dir
@@ -402,6 +385,8 @@ def train_classifier(args: argparse.Namespace) -> None:
use_all_labeled=args.clf_use_all,
auto_val_ratio=args.clf_auto_val_ratio,
)
if args.clear_clf_cache:
clf.clear_disk_cache()
clf.build_datasets(
crop_scale=args.crop_scale,
crop_size=args.crop_size,
@@ -531,6 +516,8 @@ def evaluate(args: argparse.Namespace) -> None:
pre = ensure_preprocessing()
seg = _load_segmentation(pre, args)
clf, clf_ckpt = _load_classifier(pre, seg, args)
if args.clear_clf_cache:
clf.clear_disk_cache()
def evaluate_subset(
clf_obj: RefugeClassification,
@@ -637,23 +624,27 @@ def evaluate_segmentation(args: argparse.Namespace) -> None:
in_memory_cache=args.in_memory_cache,
loader_workers=args.loader_workers,
)
if args.seg_weights is None:
raise SystemExit(
"--seg-weights must be specified for --eval-seg; "
"e.g. --seg-weights models/v2/refuge/segmentation/per_image_refuge_build/best.pt"
)
ckpt = args.seg_weights
if not ckpt.exists():
raise FileNotFoundError(f"Segmentation weights not found at {ckpt}")
state = torch.load(ckpt, map_location=segmenter.device)
state_dict = state.get("model", state)
segmenter.model.load_state_dict(state_dict, strict=False)
print(f"[seg-eval] Loaded weights from {ckpt}")
if args.in_memory_cache:
segmenter.prebuild_in_memory_cache(
cache_workers=max(0, int(args.cache_workers)),
include_train=False,
include_val=bool(args.eval_seg_splits is None or "val" in args.eval_seg_splits),
include_holdout=bool(args.eval_seg_splits is None or "holdout" in args.eval_seg_splits),
include_val="val" in args.eval_seg_splits,
include_holdout="holdout" in args.eval_seg_splits,
)
ckpt = resolve_unet_weights(args.seg_weights)
if ckpt.exists():
state = torch.load(ckpt, map_location=segmenter.device)
state_dict = state.get("model", state)
segmenter.model.load_state_dict(state_dict, strict=False)
print(f"[seg-eval] Loaded weights from {ckpt}")
else:
raise FileNotFoundError(f"Segmentation weights not found at {ckpt}")
dataset_filter = args.eval_seg_datasets
split_filter = args.eval_seg_splits
output_dir = args.eval_seg_output or Path("analysis_data/segmenter_eval")
@@ -672,9 +663,6 @@ def evaluate_segmentation(args: argparse.Namespace) -> None:
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="REFUGE pipeline helper")
parser.add_argument(
"--train-seg", action="store_true", help="Train the segmentation model"
)
parser.add_argument(
"--train-unet-seg",
action="store_true",
@@ -809,9 +797,14 @@ def parse_args() -> argparse.Namespace:
parser.add_argument(
"--clf-cache-dir",
type=Path,
default=Path("analysis_data/classifier_cache"),
default=Path("cache_data/classifier_cache"),
help="Directory to cache classifier preprocessing artifacts",
)
parser.add_argument(
"--clear-clf-cache",
action="store_true",
help="Delete all cached geometry/mask files before running (use when segmenter weights have changed)",
)
parser.add_argument(
"--clf-use-all",
action="store_true",
@@ -850,7 +843,8 @@ def parse_args() -> argparse.Namespace:
"--eval-seg-splits",
nargs="+",
choices=["train", "val", "holdout"],
help="Segmentation splits to evaluate (default: val)",
default=["holdout"],
help="Segmentation splits to evaluate (default: holdout)",
)
parser.add_argument(
"--eval-seg-output",
@@ -935,7 +929,6 @@ def main() -> None:
if not any(
[
args.train_seg,
args.train_unet_seg,
args.train_clf,
args.eval,
@@ -944,12 +937,9 @@ def main() -> None:
]
):
raise SystemExit(
"Specify at least one action: --train-seg, --train-unet-seg, --train-clf, --eval, --eval-seg, or --export-backbone"
"Specify at least one action: --train-unet-seg, --train-clf, --eval, --eval-seg, or --export-backbone"
)
if args.train_seg:
train_segmentation(args)
if args.train_unet_seg:
train_unet_segmenter(args)
scripts/main/v2/run_modes_roi_overnight.sh
Executable
+34
View File
@@ -0,0 +1,34 @@
#!/usr/bin/env bash
set -euo pipefail
# Runs two back-to-back Hypertower mode comparisons with ROI cropping:
# 1) GT masks
# 2) UNet masks
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
cd "$ROOT_DIR"
COMMON_ARGS=(
--eval-modes binary multiclass
--tower-modes single ensemble bilateral
--epochs 40
--n-splits 5
--batch-size 8
--backbone refugelike
--img-crop-manifest manifest.csv
)
echo "[1/2] Starting GT ROI run..."
python3 scripts/basic_analysis/compare_hypertower_modes.py \
"${COMMON_ARGS[@]}" \
--img-crop-gt \
--run-name v2_modes_full_40ep_5fold_roi_gt_holdout
echo "[2/2] Starting UNet ROI run..."
python3 scripts/basic_analysis/compare_hypertower_modes.py \
"${COMMON_ARGS[@]}" \
--img-crop-weights models/v2/refuge/segmentation/per_image_refuge_build/best.pt \
--img-crop-normalize per_image \
--run-name v2_modes_full_40ep_5fold_roi_unet_perimage_refugebuild_holdout
echo "All runs complete."
scripts/main/v2/run_modes_roi_smoke.sh
Executable
+42
View File
@@ -0,0 +1,42 @@
#!/usr/bin/env bash
set -euo pipefail
# Quick smoke test for ROI mode runs:
# 1) GT masks
# 2) UNet masks
# Uses 1 epoch and 1 fold for fast validation.
ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
cd "$ROOT_DIR"
COMMON_ARGS=(
--eval-modes binary multiclass
--tower-modes single ensemble bilateral
--epochs 1
--n-splits 2
--folds 1
--batch-size 8
--backbone refugelike
--img-crop-manifest manifest.csv
--warmup-tower-epochs 0
--warmup-fused-epochs 0
--single-warmup-tower-epochs 0
--single-warmup-fused-epochs 0
--bilat-warmup-tower-epochs 0
--bilat-warmup-fused-epochs 0
)
echo "[smoke 1/2] Starting GT ROI run..."
python3 scripts/basic_analysis/compare_hypertower_modes.py \
"${COMMON_ARGS[@]}" \
--img-crop-gt \
--run-name smoke_v2_modes_roi_gt
echo "[smoke 2/2] Starting UNet ROI run..."
python3 scripts/basic_analysis/compare_hypertower_modes.py \
"${COMMON_ARGS[@]}" \
--img-crop-weights models/v2/refuge/segmentation/per_image_refuge_build/best.pt \
--img-crop-normalize per_image \
--run-name smoke_v2_modes_roi_unet_perimage
echo "Smoke runs complete."
scripts/output_analysis/explainability/explain_fold.py
+649
View File
@@ -0,0 +1,649 @@
#!/usr/bin/env python3
"""
Post-hoc explainability for a single saved fold.
Phase 1 — MD permutation feature importance (bar chart + CSV).
Phase 2 — GradCAM overlays on all holdout (or val) patients.
Usage:
python scripts/output_analysis/explainability/explain_fold.py \
--fold-dir analysis_data/.../binary/single/fold0 \
[--checkpoint best_single.pt | best_holdout_single.pt] \
[--split holdout] # falls back to val if no holdout
[--image-dir Papila/FundusImages] \
[--clinical-dir Papila/ClinicalData] \
[--n-permutations 30] \
[--seed 0] \
[--alpha 0.45]
"""
from __future__ import annotations
import argparse
import json
import sys
from pathlib import Path
from types import SimpleNamespace
import matplotlib
matplotlib.use("Agg")
import matplotlib.cm as cm
import matplotlib.patches as mpatches
import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from sklearn.metrics import roc_auc_score
REPO_ROOT = Path(__file__).resolve().parents[3]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from classes.v2.data_bundle import DataBundle
from classes.v2.papila_builders import build_papila_data
from classes.v2.profiles.papila import build_papila_profile
from classes.v2.split_manager import PatientFirstSplitManager
from classes.v2.v2_hypertower import (
SingleEyeHT,
_score_arrays,
build_eval_transform,
filter_bilateral_samples,
make_loader,
)
# ---------------------------------------------------------------------------
# Label display helpers
# ---------------------------------------------------------------------------
BINARY_LABELS = {0: "Normal", 1: "Glaucoma"}
MULTICLASS_LABELS = {0: "Normal", 1: "Glaucoma", 2: "Suspect"}
def label_name(label: int, eval_mode: str) -> str:
mapping = BINARY_LABELS if eval_mode == "binary" else MULTICLASS_LABELS
return mapping.get(int(label), str(label))
# ---------------------------------------------------------------------------
# GradCAM
# ---------------------------------------------------------------------------
class GradCAM:
"""Minimal GradCAM using forward/backward hooks. No extra dependencies."""
def __init__(self, target_layer: torch.nn.Module) -> None:
self._acts: torch.Tensor | None = None
self._grads: torch.Tensor | None = None
self._h1 = target_layer.register_forward_hook(self._save_acts)
self._h2 = target_layer.register_full_backward_hook(self._save_grads)
def _save_acts(self, _m, _i, output):
self._acts = output.detach()
def _save_grads(self, _m, _gi, grad_output):
self._grads = grad_output[0].detach()
def compute(
self,
img: torch.Tensor,
meta: torch.Tensor,
model: torch.nn.Module,
target_class: int | None = None,
) -> tuple[np.ndarray, int]:
"""Return (cam [H,W] in [0,1], predicted_class_index)."""
model.eval()
with torch.enable_grad():
out = model(img, meta)
pred = int(out.argmax(1).item())
tc = pred if target_class is None else target_class
model.zero_grad()
out[0, tc].backward()
if self._acts is None or self._grads is None:
raise RuntimeError("GradCAM hooks did not fire — check target_layer.")
weights = self._grads.mean(dim=(2, 3), keepdim=True) # [1,C,1,1]
cam = F.relu((weights * self._acts).sum(dim=1, keepdim=True)) # [1,1,h,w]
cam = F.interpolate(cam, img.shape[-2:], mode="bilinear", align_corners=False)
cam_np = cam.squeeze().cpu().numpy()
lo, hi = cam_np.min(), cam_np.max()
cam_np = (cam_np - lo) / (hi - lo + 1e-8)
return cam_np, pred
def remove(self) -> None:
self._h1.remove()
self._h2.remove()
def get_gradcam_layer(model: SingleEyeHT, backbone: str) -> torch.nn.Module:
"""Return the final spatial feature map layer for GradCAM."""
bb = model.img_tower.backbone
key = backbone.lower()
if key in ("refugelike",) or "resnet" in key:
return bb.layer4[-1]
if "efficientnet" in key or "refuge_efficient" in key:
return bb.features[-1]
if "densenet" in key or key == "refuge_densenet":
return bb.features.denseblock4
if "mobilenet" in key:
return bb.features[-1]
if "vgg" in key:
return bb.features[-1]
raise ValueError(f"Unknown backbone for GradCAM target layer: {backbone!r}")
def overlay_gradcam(
original_pil: Image.Image, cam: np.ndarray, alpha: float = 0.45
) -> Image.Image:
"""Blend a jet-coloured GradCAM map onto the original image."""
cam_u8 = (cam * 255).astype(np.uint8)
cam_resized = (
np.array(Image.fromarray(cam_u8).resize(original_pil.size, Image.BILINEAR))
/ 255.0
)
colored = (cm.jet(cam_resized)[:, :, :3] * 255).astype(np.uint8)
return Image.blend(original_pil.convert("RGB"), Image.fromarray(colored), alpha)
# ---------------------------------------------------------------------------
# Feature index map
# ---------------------------------------------------------------------------
def build_feature_index_map(data: DataBundle) -> dict[str, dict]:
"""
Return a mapping feature_name → {"value_dims": [...], "missing_dims": [...]}
that covers every input dimension of the MD tower vector.
Layout (from DataBundle.vectorize_row):
[scalar_0..scalar_n-1 | cat_onehot | scalar_missing_0..scalar_missing_n-1]
"""
n_scalar = len(data.scalar_cols)
cat_expanded = sum(len(m) for m in data.cat_maps.values())
feature_map: dict[str, dict] = {}
idx = 0
# Scalar features: value_dim + corresponding missing flag
for i, col in enumerate(data.scalar_cols):
missing_dim = n_scalar + cat_expanded + i
feature_map[col] = {"value_dims": [i], "missing_dims": [missing_dim]}
idx += 1
# Categorical features: permute the entire one-hot block
cat_offset = n_scalar
for col in data.cat_cols:
n_cats = len(data.cat_maps[col])
dims = list(range(cat_offset, cat_offset + n_cats))
feature_map[col] = {"value_dims": dims, "missing_dims": []}
cat_offset += n_cats
return feature_map
# ---------------------------------------------------------------------------
# Phase 1 — MD permutation importance
# ---------------------------------------------------------------------------
def run_permutation_importance(
model: SingleEyeHT,
loader,
data: DataBundle,
num_classes: int,
device: torch.device,
n_permutations: int,
seed: int,
out_dir: Path,
) -> None:
print("\n[Phase 1] MD permutation importance ...", flush=True)
# ---- cache image embeddings + collect meta tensors + labels ----
img_feats_list, md_list, label_list = [], [], []
model.eval()
with torch.no_grad():
for batch in loader:
imgs = batch["image_1"].to(device)
meta = batch["matrix_1"].to(device)
labels = batch["label_1"]
img_feats_list.append(model.img_tower(imgs))
md_list.append(meta)
label_list.append(labels)
img_feats = torch.cat(img_feats_list) # [N, img_dim]
md_tensor = torch.cat(md_list) # [N, feature_dim]
y_true = torch.cat(label_list).numpy()
N = len(y_true)
if N == 0:
print(" [Phase 1] No samples — skipping.", flush=True)
return
# ---- baseline AUC ----
with torch.no_grad():
md_feats = model.md_tower(md_tensor)
fused, _, _ = model.bridge(img_feats, md_feats)
probs_baseline = torch.softmax(fused, dim=1).cpu().numpy()
_, baseline_auc, _ = _score_arrays(y_true, probs_baseline, num_classes)
print(f" Baseline AUC: {baseline_auc:.4f} (N={N})", flush=True)
# ---- feature index map ----
feat_map = build_feature_index_map(data)
rng = np.random.default_rng(seed)
results = []
for feat_name, dims in feat_map.items():
all_dims = dims["value_dims"] + dims["missing_dims"]
drops = []
for _ in range(n_permutations):
perm = md_tensor.clone()
perm_idx = torch.from_numpy(rng.permutation(N)).to(device)
perm[:, all_dims] = perm[perm_idx][:, all_dims]
with torch.no_grad():
md_p = model.md_tower(perm)
fused_p, _, _ = model.bridge(img_feats, md_p)
probs_p = torch.softmax(fused_p, dim=1).cpu().numpy()
_, auc_p, _ = _score_arrays(y_true, probs_p, num_classes)
drops.append(baseline_auc - auc_p)
mean_drop = float(np.mean(drops))
std_drop = float(np.std(drops))
results.append({"feature": feat_name, "importance": mean_drop, "std": std_drop})
print(
f" {feat_name:30s} Δ AUC = {mean_drop:+.4f} ± {std_drop:.4f}", flush=True
)
results.sort(key=lambda r: r["importance"], reverse=True)
# ---- save CSV ----
import csv
csv_path = out_dir / "md_permutation_importance.csv"
with csv_path.open("w", newline="") as f:
writer = csv.DictWriter(f, fieldnames=["feature", "importance", "std"])
writer.writeheader()
writer.writerows(results)
# ---- bar chart ----
names = [r["feature"] for r in results]
imps = [r["importance"] for r in results]
stds = [r["std"] for r in results]
colors = ["#e05c5c" if v >= 0 else "#5c9ee0" for v in imps]
fig, ax = plt.subplots(figsize=(9, max(4, len(names) * 0.45)))
y_pos = np.arange(len(names))
bars = 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)
ax.set_title(
f"MD Tower — Permutation Feature Importance\n"
f"baseline AUC={baseline_auc:.4f} N={N} repeats={n_permutations}",
fontsize=11,
)
fig.tight_layout()
fig.savefig(out_dir / "md_permutation_importance.png", dpi=150)
plt.close(fig)
print(f" Saved → {out_dir / 'md_permutation_importance.png'}", flush=True)
# ---------------------------------------------------------------------------
# Phase 2 — GradCAM overlays
# ---------------------------------------------------------------------------
def run_gradcam(
model: SingleEyeHT,
loader,
data: DataBundle,
eval_df,
eval_mode: str,
backbone: str,
device: torch.device,
alpha: float,
out_dir: Path,
) -> None:
print("\n[Phase 2] GradCAM overlays ...", flush=True)
gradcam_dir = out_dir / "gradcam"
gradcam_dir.mkdir(exist_ok=True)
target_layer = get_gradcam_layer(model, backbone)
gcam = GradCAM(target_layer)
num_classes = model.bridge.classifier_fused[-1].out_features
overlay_grid_items: list[
tuple[Image.Image | None, Image.Image | None, str, bool]
] = []
model.eval()
for batch in loader:
img_od = batch["image_1"].to(device) # [1, 3, H, W]
img_os = batch["image_2"].to(device) # [1, 3, H, W]
meta_od = batch["matrix_1"].to(device) # [1, feature_dim]
meta_os = batch["matrix_2"].to(device)
label = int(batch["label_1"][0].item())
pid = batch["id_1"][0]
# GradCAM for each eye (OD drives the prediction label)
cam_od, pred = gcam.compute(img_od, meta_od, model)
cam_os, _ = gcam.compute(img_os, meta_os, model)
# Confidence of predicted class
with torch.no_grad():
out_od = model(img_od, meta_od)
conf = float(torch.softmax(out_od, dim=1)[0, pred].item())
# Load original (un-normalised) images from disk
row_od = eval_df[
(eval_df["Patient ID"] == int(pid)) & (eval_df["eyeID"] == "OD")
]
row_os = eval_df[
(eval_df["Patient ID"] == int(pid)) & (eval_df["eyeID"] == "OS")
]
orig_od = (
Image.open(data.get_image_path(row_od.iloc[0])).convert("RGB")
if len(row_od)
else None
)
orig_os = (
Image.open(data.get_image_path(row_os.iloc[0])).convert("RGB")
if len(row_os)
else None
)
true_name = label_name(label, eval_mode)
pred_name = label_name(pred, eval_mode)
correct = label == pred
title = (
f"Patient {pid} | True: {true_name} | Pred: {pred_name} "
f"| conf={conf:.2f} {'✓' if correct else '✗'}"
)
# ---- per-patient 2×2 figure (OD raw | OD overlay / OS raw | OS overlay) ----
fig, axes = plt.subplots(2, 2, figsize=(10, 9))
fig.suptitle(
title, fontsize=11, fontweight="bold", color="green" if correct else "red"
)
# Row 0: OD
if orig_od is not None:
axes[0, 0].imshow(orig_od)
axes[0, 0].set_title("OD — original", fontsize=9)
axes[0, 1].imshow(overlay_gradcam(orig_od, cam_od, alpha))
axes[0, 1].set_title("OD — GradCAM", fontsize=9)
else:
axes[0, 0].set_title("OD — (missing)", fontsize=9)
axes[0, 0].axis("off")
axes[0, 1].axis("off")
# Row 1: OS
if orig_os is not None:
axes[1, 0].imshow(orig_os)
axes[1, 0].set_title("OS — original", fontsize=9)
axes[1, 1].imshow(overlay_gradcam(orig_os, cam_os, alpha))
axes[1, 1].set_title("OS — GradCAM", fontsize=9)
else:
axes[1, 0].set_title("OS — (missing)", fontsize=9)
axes[1, 0].axis("off")
axes[1, 1].axis("off")
fig.tight_layout()
out_path = gradcam_dir / f"patient_{pid}_OD_OS.png"
fig.savefig(out_path, dpi=120)
plt.close(fig)
print(
f" Patient {pid}: {true_name} → {pred_name} ({conf:.2f}) → {out_path.name}",
flush=True,
)
# Accumulate for summary grid
od_overlay = overlay_gradcam(orig_od, cam_od, alpha) if orig_od else None
os_overlay = overlay_gradcam(orig_os, cam_os, alpha) if orig_os else None
short_lbl = f"P{pid} {true_name[:3]}→{pred_name[:3]} {'✓' if correct else '✗'}"
overlay_grid_items.append((od_overlay, os_overlay, short_lbl, correct))
gcam.remove()
# ---- summary grid: N_patients rows × 2 cols (OD overlay | OS overlay) ----
n = len(overlay_grid_items)
if n == 0:
print(" [Phase 2] No patients to visualise.", flush=True)
return
fig, axes = plt.subplots(n, 2, figsize=(8, n * 3.2 + 0.8))
if n == 1:
axes = axes[np.newaxis, :]
fig.suptitle("GradCAM Summary Grid — all holdout patients", fontsize=12)
for i, (od_ov, os_ov, lbl, correct) in enumerate(overlay_grid_items):
color = "green" if correct else "red"
for j in range(2):
axes[i, j].axis("off")
if od_ov is not None:
axes[i, 0].imshow(od_ov)
axes[i, 0].set_title(f"{lbl}\nOD", fontsize=7, color=color)
if os_ov is not None:
axes[i, 1].imshow(os_ov)
axes[i, 1].set_title(f"{lbl}\nOS", fontsize=7, color=color)
fig.tight_layout()
grid_path = out_dir / "gradcam_summary_grid.png"
fig.savefig(grid_path, dpi=120)
plt.close(fig)
print(f" Summary grid → {grid_path}", flush=True)
# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def parse_args():
ap = argparse.ArgumentParser(
description="Post-hoc explainability for a saved fold."
)
ap.add_argument(
"--fold-dir",
type=Path,
required=True,
help="Path to fold directory, e.g. analysis_data/.../binary/single/fold0",
)
ap.add_argument(
"--checkpoint",
default="best_single.pt",
help="Checkpoint filename inside fold_dir (default: best_single.pt; "
"use best_holdout_single.pt for holdout-selected model)",
)
ap.add_argument(
"--split",
choices=["holdout", "val"],
default="holdout",
help="Which patient set to analyse (default: holdout, falls back to val)",
)
ap.add_argument("--image-dir", default="Papila/FundusImages")
ap.add_argument("--clinical-dir", default="Papila/ClinicalData")
ap.add_argument("--label-col", default="Diagnosis")
ap.add_argument("--cat-cols", nargs="*", default=["Gender", "Phakic/Pseudophakic"])
ap.add_argument("--fold-seed", type=int, default=42)
ap.add_argument("--holdout-seed", type=int, default=123)
ap.add_argument("--holdout-per-class", type=int, default=5)
ap.add_argument("--n-splits", type=int, default=5)
ap.add_argument(
"--n-permutations",
type=int,
default=30,
help="Repetitions per feature for permutation importance (default: 30)",
)
ap.add_argument("--seed", type=int, default=0)
ap.add_argument(
"--alpha",
type=float,
default=0.45,
help="GradCAM overlay opacity (default: 0.45)",
)
ap.add_argument("--batch-size", type=int, default=1)
ap.add_argument("--no-phase1", action="store_true", help="Skip MD importance")
ap.add_argument("--no-phase2", action="store_true", help="Skip GradCAM")
return ap.parse_args()
def main():
args = parse_args()
fold_dir = args.fold_dir.resolve()
if not fold_dir.is_dir():
sys.exit(f"[ERROR] fold_dir does not exist: {fold_dir}")
ckpt_path = fold_dir / args.checkpoint
if not ckpt_path.exists():
sys.exit(
f"[ERROR] Checkpoint not found: {ckpt_path}\n"
f" Run training with --save-checkpoints (now the default) to produce checkpoints."
)
# ---- read config from summary.json in parent (tower-mode) dir ----
summary_path = fold_dir.parent / "summary.json"
if not summary_path.exists():
sys.exit(f"[ERROR] summary.json not found: {summary_path}")
summary = json.loads(summary_path.read_text())
backbone = summary["backbone"]
eval_mode = summary["eval_mode"]
tower_mode = summary.get("tower_mode", "single")
fold_idx = int(fold_dir.name.replace("fold", ""))
print(
f"[explain_fold] fold={fold_idx} backbone={backbone} eval_mode={eval_mode} tower_mode={tower_mode}"
)
if tower_mode not in ("single", "ensemble"):
sys.exit(
f"[ERROR] explain_fold currently supports single/ensemble tower modes, got: {tower_mode!r}"
)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f"[explain_fold] device={device} checkpoint={args.checkpoint}")
# ---- build DataBundle ----
print("[explain_fold] Loading clinical data ...", flush=True)
data = build_papila_data(
image_dir=args.image_dir,
clinical_dir=args.clinical_dir,
label_col=args.label_col,
cat_cols=args.cat_cols,
n_splits=args.n_splits,
random_seed=args.fold_seed,
)
df_mode = data.df.copy()
if eval_mode == "binary":
df_mode = df_mode[df_mode[args.label_col].isin([0, 1])].reset_index(drop=True)
num_classes = 2 if eval_mode == "binary" else int(df_mode[args.label_col].nunique())
# ---- reconstruct the exact same split ----
print("[explain_fold] Reconstructing split ...", flush=True)
splitter = PatientFirstSplitManager(
patient_col="Patient ID", label_col=args.label_col
)
split_args = SimpleNamespace(
eval_mode=eval_mode,
holdout_per_class=args.holdout_per_class,
holdout_seed=args.holdout_seed,
n_splits=args.n_splits,
fold_seed=args.fold_seed,
)
clinical_ns = SimpleNamespace(df=df_mode, label_col=args.label_col)
plans = splitter.build_plans(clinical=clinical_ns, args=split_args, profile=None)
if fold_idx >= len(plans):
sys.exit(f"[ERROR] fold_idx={fold_idx} but only {len(plans)} plans built.")
split = plans[fold_idx]
if (
args.split == "holdout"
and split.holdout is not None
and not split.holdout.empty
):
eval_df = split.holdout
split_name = "holdout"
else:
if args.split == "holdout":
print(" [WARN] No holdout set available; falling back to val.", flush=True)
eval_df = split.val
split_name = "val"
print(
f" Using {split_name} set: {eval_df['Patient ID'].nunique()} patients",
flush=True,
)
# ---- build loader ----
profile_patient = build_papila_profile(
patient_col="Patient ID", label_col=args.label_col, sample_mode="patient"
)
samples = filter_bilateral_samples(
profile_patient.build_samples(df=eval_df, clinical=data)
)
if not samples:
sys.exit("[ERROR] No bilateral samples found in the eval set.")
loader = make_loader(
samples,
profile_patient.slot_descriptors(),
image_transform=build_eval_transform(backbone),
image_preprocessor=None,
batch_size=args.batch_size,
shuffle=False,
num_workers=0,
)
# ---- load model ----
print(f"[explain_fold] Loading model from {ckpt_path} ...", flush=True)
model = SingleEyeHT(
backbone=backbone,
freeze_ratio=0.0,
augment=False,
clinical_data=data,
num_classes=num_classes,
).to(device)
state = torch.load(ckpt_path, map_location=device)
model.load_state_dict(state)
model.eval()
# ---- output directory ----
out_dir = fold_dir / "explainability"
out_dir.mkdir(exist_ok=True)
print(f"[explain_fold] Output → {out_dir}", flush=True)
# ---- Phase 1 ----
if not args.no_phase1:
run_permutation_importance(
model=model,
loader=loader,
data=data,
num_classes=num_classes,
device=device,
n_permutations=args.n_permutations,
seed=args.seed,
out_dir=out_dir,
)
# ---- Phase 2 ----
if not args.no_phase2:
run_gradcam(
model=model,
loader=loader,
data=data,
eval_df=eval_df,
eval_mode=eval_mode,
backbone=backbone,
device=device,
alpha=args.alpha,
out_dir=out_dir,
)
print("\n[explain_fold] Done.", flush=True)
if __name__ == "__main__":
main()