rpotter6298
280060db82
- Introduced multiple regression experiment configurations targeting vf_md, including: - cd_solo_reg_set.json: CD tower only regression setup. - img_solo_reg_set.json: Image tower only regression setup. - reg_head_epoch_sweep.json: Baseline regression sweeps at different epochs (50, 75, 100). - reg_head_set.json: Various regression setups including baseline and OrthoBridge configurations. - single_eye_reg.json: Single-eye regression setup for worst-eye aggregation analysis. - Added ensemble configurations for OrthoBridge with different inner bridges: - ortho_alts_ensemble.json: Ensemble tests with ConcatBridge, PairwiseAdditiveBridge, and GatedAdditiveBridge. - ortho_alts_tritower.json: Tritower tests with the same inner bridges. - Created V2-M specific configurations: - baseline_reg_nt50.json: Regression baseline with V2-M backbone. - geom_vec_gt.json and geom_vec_unet.json: Geometry vector injection experiments with V2-M. - single_l1_bridges.json: Single-eye ensemble experiments with various bridge types. - tritower_geom_gt.json: Tritower setup with GT contour-rasterized masks. - Promoted existing experiments to higher repetitions for robustness.
305 lines
13 KiB
Python
305 lines
13 KiB
Python
"""Generic explainability primitives — GradCAM + region-of-interest analysis
|
||
+ column-shuffling permutation importance.
|
||
|
||
This module is dataset-agnostic. All PAPILA-specific knowledge (where the disc
|
||
contour lives, eye orientation conventions, clinical feature names, etc.) is
|
||
supplied by the profile through optional methods consumed via getattr — see
|
||
v4.classes.profiles.v4papila.ImageDataView and ClinicalDataView for examples.
|
||
|
||
Public surface
|
||
--------------
|
||
GradCAM — hooks-based class-activation maps via fwd/bwd hooks
|
||
overlay_gradcam — render a heatmap blended onto a PIL image
|
||
cam_region_stats — quantify CAM attention inside/outside a binary ROI
|
||
cam_region_patch — crop a square patch centered on the ROI centroid
|
||
permutation_importance — per-feature drop-in-score under column shuffling
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
from typing import Callable, Sequence
|
||
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn as nn
|
||
import torch.nn.functional as F
|
||
from matplotlib import cm
|
||
from PIL import Image
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# GradCAM
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
|
||
class GradCAM:
|
||
"""Minimal GradCAM via forward/backward hooks.
|
||
|
||
Architecture-agnostic: the caller supplies a forward_fn that returns a
|
||
(1, num_classes) logits tensor, so this class doesn't need to know whether
|
||
the model takes (img,), (img, meta), or composed tower/bridge/head inputs.
|
||
|
||
Usage
|
||
-----
|
||
cam_module = GradCAM(target_layer)
|
||
cam, pred = cam_module.compute(
|
||
forward_fn=lambda: head(bridge([tower_img(x), tower_cd(c)])),
|
||
output_shape=(224, 224),
|
||
)
|
||
cam_module.remove()
|
||
"""
|
||
|
||
def __init__(self, target_layer: 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, _module, _inputs, output):
|
||
self._acts = output.detach()
|
||
|
||
def _save_grads(self, _module, _grad_in, grad_out):
|
||
self._grads = grad_out[0].detach()
|
||
|
||
def compute(
|
||
self,
|
||
forward_fn: Callable[[], torch.Tensor],
|
||
*,
|
||
output_shape: tuple[int, int],
|
||
target_class: int | None = None,
|
||
) -> tuple[np.ndarray, int]:
|
||
"""Run forward via forward_fn, backward against target_class, return CAM.
|
||
|
||
Parameters
|
||
----------
|
||
forward_fn
|
||
Zero-arg callable returning a (1, C) logits tensor. Caller is
|
||
responsible for placing model in eval mode and zeroing grads.
|
||
output_shape
|
||
(H, W) at which the CAM is bilinearly upsampled. Typically the
|
||
eval-crop image size, e.g. (224, 224).
|
||
target_class
|
||
Class index to backprop against. If None, uses the predicted class.
|
||
|
||
Returns
|
||
-------
|
||
(cam, pred)
|
||
cam : float32 ndarray of shape output_shape, normalised to [0, 1]
|
||
pred : int, the predicted class index
|
||
"""
|
||
with torch.enable_grad():
|
||
out = forward_fn()
|
||
pred = int(out.argmax(1).item())
|
||
tc = pred if target_class is None else target_class
|
||
out[0, tc].backward()
|
||
|
||
assert self._grads is not None and self._acts is not None, (
|
||
"GradCAM hooks did not fire — check that target_layer is in the "
|
||
"forward path of forward_fn()."
|
||
)
|
||
weights = self._grads.mean(dim=(2, 3), keepdim=True)
|
||
cam = F.relu((weights * self._acts).sum(dim=1, keepdim=True))
|
||
cam = F.interpolate(cam, output_shape, mode="bilinear", align_corners=False)
|
||
cam_np = cam.squeeze().cpu().numpy()
|
||
lo, hi = cam_np.min(), cam_np.max()
|
||
return ((cam_np - lo) / (hi - lo + 1e-8)).astype(np.float32), pred
|
||
|
||
def remove(self) -> None:
|
||
"""Detach the forward/backward hooks. Call when done with this instance."""
|
||
self._h1.remove()
|
||
self._h2.remove()
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# Visualisation
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
|
||
def overlay_gradcam(pil: Image.Image, cam: np.ndarray, alpha: float = 0.45) -> Image.Image:
|
||
"""Blend a jet-colored heatmap of `cam` onto `pil` at alpha blending weight."""
|
||
cam_u8 = (np.clip(cam, 0.0, 1.0) * 255).astype(np.uint8)
|
||
cam_r = np.array(Image.fromarray(cam_u8).resize(pil.size, Image.BILINEAR)) / 255.0
|
||
colored = (cm.jet(cam_r)[:, :, :3] * 255).astype(np.uint8)
|
||
return Image.blend(pil.convert("RGB"), Image.fromarray(colored), alpha)
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# Region-of-interest analysis (generic — no PAPILA knowledge)
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
|
||
def cam_region_stats(cam: np.ndarray, roi_mask: np.ndarray) -> dict[str, float]:
|
||
"""Quantify CAM attention restricted to / outside a binary ROI mask.
|
||
|
||
The CAM and the mask must be at the same resolution.
|
||
|
||
Returns
|
||
-------
|
||
dict with:
|
||
inside_mean : mean CAM intensity inside the ROI
|
||
outside_mean : mean CAM intensity outside the ROI
|
||
inside_fraction : sum(CAM inside) / sum(CAM total); fraction of total
|
||
attention that falls inside the ROI
|
||
roi_area_frac : ROI area / total area; how much of the image the ROI
|
||
covers — useful for normalising inside_fraction
|
||
"""
|
||
if cam.shape != roi_mask.shape:
|
||
raise ValueError(
|
||
f"cam shape {cam.shape} != roi_mask shape {roi_mask.shape}"
|
||
)
|
||
mask = roi_mask.astype(bool)
|
||
cam_sum = float(cam.sum())
|
||
inside_sum = float(cam[mask].sum()) if mask.any() else 0.0
|
||
outside_sum = cam_sum - inside_sum
|
||
n_inside = int(mask.sum())
|
||
n_outside = int(mask.size - n_inside)
|
||
return {
|
||
"inside_mean": inside_sum / max(1, n_inside),
|
||
"outside_mean": outside_sum / max(1, n_outside),
|
||
"inside_fraction": inside_sum / cam_sum if cam_sum > 0 else 0.0,
|
||
"roi_area_frac": n_inside / mask.size,
|
||
}
|
||
|
||
|
||
def cam_region_patch(
|
||
cam: np.ndarray,
|
||
roi_mask: np.ndarray,
|
||
*,
|
||
span: float = 5.0,
|
||
patch_size: int = 96,
|
||
) -> tuple[np.ndarray | None, float | None]:
|
||
"""Crop a square patch centered on the ROI centroid.
|
||
|
||
The crop side is `span × roi_radius` (in CAM pixels), where roi_radius is
|
||
derived from the mask area assuming a circular ROI. The output is resized
|
||
to (patch_size, patch_size).
|
||
|
||
Returns
|
||
-------
|
||
(patch, roi_radius_in_patch_pixels)
|
||
patch : float32 ndarray of shape (patch_size, patch_size), values in
|
||
[0, 1]. Padded with zeros if the centered crop window extends
|
||
outside the CAM.
|
||
roi_radius_in_patch_pixels : float, the ROI radius re-expressed in
|
||
output-patch pixels (useful for drawing ROI circles in figures).
|
||
Returns (None, None) if the mask is empty.
|
||
"""
|
||
mask = roi_mask.astype(bool)
|
||
if not mask.any():
|
||
return None, None
|
||
|
||
ys, xs = np.where(mask)
|
||
cy, cx = ys.mean(), xs.mean()
|
||
roi_r = float(np.sqrt(mask.sum() / np.pi))
|
||
half = max(1, int(round(span * roi_r / 2)))
|
||
h, w = cam.shape
|
||
|
||
y0, y1 = int(round(cy)) - half, int(round(cy)) + half
|
||
x0, x1 = int(round(cx)) - half, int(round(cx)) + half
|
||
pt = max(0, -y0)
|
||
pb = max(0, y1 - h)
|
||
pl = max(0, -x0)
|
||
pr = max(0, x1 - w)
|
||
cam_pad = np.pad(cam, ((pt, pb), (pl, pr)), constant_values=0.0)
|
||
window = cam_pad[y0 + pt: y1 + pt, x0 + pl: x1 + pl]
|
||
|
||
patch = np.array(
|
||
Image.fromarray((np.clip(window, 0, 1) * 255).astype(np.uint8))
|
||
.resize((patch_size, patch_size), Image.BILINEAR)
|
||
) / 255.0
|
||
roi_r_out = patch_size * roi_r / (2 * half)
|
||
return patch.astype(np.float32), roi_r_out
|
||
|
||
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
# Permutation importance (column-shuffling)
|
||
# ─────────────────────────────────────────────────────────────────────────────
|
||
|
||
def permutation_importance(
|
||
*,
|
||
score_fn: Callable[[np.ndarray], float],
|
||
X: np.ndarray,
|
||
groups: Sequence[Sequence[int]] | None = None,
|
||
n_permutations: int = 30,
|
||
seed: int = 0,
|
||
feature_names: Sequence[str] | None = None,
|
||
) -> dict:
|
||
"""Per-feature drop-in-score from column shuffling.
|
||
|
||
Generic and feature-agnostic — the caller wires the model (or any
|
||
arbitrary scoring pipeline) into ``score_fn``. Higher score = better;
|
||
drop = baseline_score − score(X_with_column_permuted).
|
||
|
||
Parameters
|
||
----------
|
||
score_fn
|
||
Callable taking the input matrix ``X`` (shape (N, F)) and returning
|
||
a scalar metric. Caller is responsible for any constants the metric
|
||
depends on (e.g. labels, held-out image embeddings).
|
||
X
|
||
Float ndarray of shape (N, F). The base matrix to permute.
|
||
groups
|
||
Optional list of column-index groups. If provided, each group is
|
||
permuted together (useful when several model dims encode one
|
||
conceptual feature, e.g. one-hot encodings of a categorical column).
|
||
Defaults to one group per column.
|
||
n_permutations
|
||
Number of independent shuffles per group; mean and std of the drop
|
||
are reported across these.
|
||
seed
|
||
Seed for the shuffling RNG.
|
||
feature_names
|
||
Optional human-readable names, one per group. If not provided, names
|
||
default to f"group_{i}" or f"col_{j}" when groups is None.
|
||
|
||
Returns
|
||
-------
|
||
dict with keys:
|
||
baseline : float, the score on unmodified X
|
||
mean_drop : ndarray of shape (G,), mean (baseline − shuffled) per group
|
||
std_drop : ndarray of shape (G,), std of per-permutation drops
|
||
feature_names : list[str] of length G
|
||
groups : list[list[int]] of resolved column groups
|
||
"""
|
||
if X.ndim != 2:
|
||
raise ValueError(f"X must be 2-D (N, F); got shape {X.shape}")
|
||
n, f_dim = X.shape
|
||
|
||
if groups is None:
|
||
resolved_groups: list[list[int]] = [[j] for j in range(f_dim)]
|
||
else:
|
||
resolved_groups = [list(g) for g in groups]
|
||
all_cols = [c for g in resolved_groups for c in g]
|
||
if any(c < 0 or c >= f_dim for c in all_cols):
|
||
raise ValueError(
|
||
f"groups contain column indices outside [0, {f_dim})"
|
||
)
|
||
|
||
n_groups = len(resolved_groups)
|
||
if feature_names is None:
|
||
names = (
|
||
[f"col_{g[0]}" for g in resolved_groups] if groups is None
|
||
else [f"group_{i}" for i in range(n_groups)]
|
||
)
|
||
else:
|
||
names = list(feature_names)
|
||
if len(names) != n_groups:
|
||
raise ValueError(
|
||
f"feature_names has length {len(names)} but there are "
|
||
f"{n_groups} groups"
|
||
)
|
||
|
||
rng = np.random.default_rng(seed)
|
||
baseline = float(score_fn(X))
|
||
drops = np.empty((n_groups, n_permutations), dtype=float)
|
||
for gi, cols in enumerate(resolved_groups):
|
||
for p in range(n_permutations):
|
||
X_perm = X.copy()
|
||
perm = rng.permutation(n)
|
||
X_perm[:, cols] = X_perm[perm][:, cols]
|
||
drops[gi, p] = baseline - float(score_fn(X_perm))
|
||
|
||
return {
|
||
"baseline": baseline,
|
||
"mean_drop": drops.mean(axis=1),
|
||
"std_drop": drops.std(axis=1),
|
||
"feature_names": names,
|
||
"groups": resolved_groups,
|
||
}
|