diff --git a/.gitignore b/.gitignore index a816ee8..d955830 100644 --- a/.gitignore +++ b/.gitignore @@ -5,7 +5,7 @@ Papila/* REFUGE/* old_**/ models/* - +cache_data/ # model artifacts models/refuge/ models/v2/refuge/ diff --git a/classes/refuge_classification.py b/classes/refuge_classification.py index efbac25..b36d85c 100755 --- a/classes/refuge_classification.py +++ b/classes/refuge_classification.py @@ -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}" diff --git a/classes/unet_segmenter.py b/classes/unet_segmenter.py index 82c850e..407f8c7 100755 --- a/classes/unet_segmenter.py +++ b/classes/unet_segmenter.py @@ -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) diff --git a/classes/v2/v2_hypertower.py b/classes/v2/v2_hypertower.py index 92e5340..7dc1a55 100644 --- a/classes/v2/v2_hypertower.py +++ b/classes/v2/v2_hypertower.py @@ -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, ) diff --git a/models/v2/refuge/segmentation/per_image_refuge_build/best.pt b/models/v2/refuge/segmentation/per_image_refuge_build/best.pt index f33fe5d..d20edf6 100644 Binary files a/models/v2/refuge/segmentation/per_image_refuge_build/best.pt and b/models/v2/refuge/segmentation/per_image_refuge_build/best.pt differ diff --git a/scripts/basic_analysis/compare_hypertower_modes.py b/scripts/basic_analysis/compare_hypertower_modes.py index e2ef054..9b14b6a 100644 --- a/scripts/basic_analysis/compare_hypertower_modes.py +++ b/scripts/basic_analysis/compare_hypertower_modes.py @@ -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__": diff --git a/scripts/main/refuge/refuge_build.py b/scripts/main/refuge/refuge_build.py index db7bbf6..3e21d5a 100644 --- a/scripts/main/refuge/refuge_build.py +++ b/scripts/main/refuge/refuge_build.py @@ -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) diff --git a/scripts/main/v2/run_modes_roi_overnight.sh b/scripts/main/v2/run_modes_roi_overnight.sh new file mode 100755 index 0000000..c2fb52a --- /dev/null +++ b/scripts/main/v2/run_modes_roi_overnight.sh @@ -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." diff --git a/scripts/main/v2/run_modes_roi_smoke.sh b/scripts/main/v2/run_modes_roi_smoke.sh new file mode 100755 index 0000000..ca64b4c --- /dev/null +++ b/scripts/main/v2/run_modes_roi_smoke.sh @@ -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." diff --git a/scripts/output_analysis/explainability/explain_fold.py b/scripts/output_analysis/explainability/explain_fold.py new file mode 100644 index 0000000..49a2600 --- /dev/null +++ b/scripts/output_analysis/explainability/explain_fold.py @@ -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()