hypertower/scripts/output_analysis/explainability/aggregate_gradcam.py

#!/usr/bin/env python3
"""
Aggregate raw GradCAM heatmaps across all folds for a run.

For each combination of (eye, class, correct/incorrect) computes:
  - mean heatmap
  - std heatmap
  - count

Also computes a scalar per patient: fraction of GradCAM attention mass that
falls within the expert-segmented optic disc region (from GT contour files),
using the manifest.csv to locate the contour for each patient/eye.

Outputs
-------
  {out_dir}/mean_heatmaps.npz
      Keys: {eye}_{class_name}_{correct|incorrect}_{mean|std|count}
      e.g.  OD_Glaucoma_correct_mean   shape (224, 224)

  {out_dir}/attention_stats.csv
      per-patient scalars: patient_id, fold, eye, true_name, pred_name,
      correct, confidence, disc_frac, entropy

Usage
-----
    python scripts/output_analysis/explainability/aggregate_gradcam.py \
        --run-dir   analysis_data/pipeline_nocrop \
        --eval-mode binary \
        --tower-mode single
"""
from __future__ import annotations

import argparse
from pathlib import Path

import numpy as np
import pandas as pd
from PIL import Image, ImageDraw


def load_disc_mask(contour_path: Path, orig_size: tuple[int, int],
                   cam_h: int, cam_w: int) -> np.ndarray | None:
    """
    Load a PAPILA disc contour TXT file, polygon-fill at original image
    dimensions, then resize to (cam_h, cam_w).  Returns a bool array or
    None if the contour cannot be loaded.
    """
    try:
        arr = np.loadtxt(str(contour_path), dtype=np.float32)
    except Exception:
        return None
    if arr.ndim == 1:
        arr = arr.reshape(-1, 2)
    if arr.shape[0] < 3 or arr.shape[1] < 2:
        return None

    # orig_size is (W, H) as PIL convention
    img = Image.new("L", orig_size, 0)
    draw = ImageDraw.Draw(img)
    draw.polygon([tuple(pt) for pt in arr[:, :2]], fill=1)
    mask = np.array(img.resize((cam_w, cam_h), Image.NEAREST), dtype=bool)
    return mask


def build_disc_lookup(manifest_path: Path) -> dict[tuple[int, str], tuple[Path, tuple[int, int]]]:
    """
    Returns {(patient_id_int, eye): (disc_contour_path, (img_W, img_H))}.
    Only PAPILA rows are included.
    """
    mf = pd.read_csv(manifest_path)
    lookup: dict[tuple[int, str], tuple[Path, tuple[int, int]]] = {}
    for _, row in mf.iterrows():
        sid = str(row["sample_id"])
        if not sid.startswith("papila_RET"):
            continue
        # sample_id: papila_RET002OD  or  papila_RET002OS
        suffix = sid[len("papila_RET"):]   # e.g. "002OD"
        eye    = suffix[-2:]                # "OD" or "OS"
        pid    = int(suffix[:-2])           # 2
        disc_path = Path(str(row["annotation_disc"]))
        img_path  = Path(str(row["image_path"]))
        if not disc_path.exists():
            continue
        # read original image size once
        try:
            with Image.open(img_path) as im:
                orig_size = im.size   # (W, H)
        except Exception:
            continue
        lookup[(pid, eye)] = (disc_path, orig_size)
    return lookup


def attention_entropy(cam: np.ndarray) -> float:
    flat = cam.flatten().astype(np.float64)
    flat = flat / (flat.sum() + 1e-12)
    return float(-np.sum(flat * np.log(flat + 1e-12)))


def load_fold(gradcam_dir: Path):
    idx_path = gradcam_dir / "gradcam_index.csv"
    if not idx_path.exists():
        return None
    idx = pd.read_csv(idx_path)
    records = []
    for _, row in idx.iterrows():
        pid = row["patient_id"]
        for eye in ("OD", "OS"):
            npy = gradcam_dir / f"patient_{pid}_{eye}_cam.npy"
            if not npy.exists():
                continue
            cam = np.load(npy)
            records.append({
                "patient_id": pid,
                "eye":        eye,
                "true_label": int(row["true_label"]),
                "true_name":  row["true_name"],
                "pred_label": int(row["pred_label"]),
                "pred_name":  row["pred_name"],
                "confidence": float(row["confidence"]),
                "correct":    bool(row["correct"]),
                "cam":        cam,
            })
    return records


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--run-dir",    default="analysis_data/pipeline_nocrop")
    ap.add_argument("--eval-mode",  default="binary")
    ap.add_argument("--tower-mode", default="single")
    ap.add_argument("--manifest",   default="manifest.csv")
    ap.add_argument("--out",        default=None)
    args = ap.parse_args()

    run_dir  = Path(args.run_dir)
    mode_dir = run_dir / args.eval_mode / args.tower_mode
    out_dir  = mode_dir / "gradcam_aggregate"
    if args.out:
        out_dir = Path(args.out)
    out_dir.mkdir(parents=True, exist_ok=True)

    # ---- build disc mask lookup ----
    manifest_path = Path(args.manifest)
    disc_lookup = build_disc_lookup(manifest_path)
    print(f"Disc mask lookup: {len(disc_lookup)} entries from {manifest_path}")

    # ---- collect all records ----
    all_records = []
    stat_rows   = []
    fold_dirs   = sorted(
        [d for d in mode_dir.iterdir() if d.is_dir() and d.name.startswith("fold")],
        key=lambda p: int(p.name.replace("fold", "")),
    )
    if not fold_dirs:
        print(f"No fold dirs found under {mode_dir}")
        return

    for fd in fold_dirs:
        gcam_dir = fd / "explainability" / "gradcam"
        records  = load_fold(gcam_dir)
        if records is None:
            print(f"  [skip] {fd.name}: no gradcam_index.csv")
            continue
        print(f"  {fd.name}: {len(records)} eye records")
        for r in records:
            r["fold"] = fd.name
            all_records.append(r)

    if not all_records:
        print("No records found — re-run explain_fold.py first.")
        return

    print(f"\nTotal eye records: {len(all_records)}")

    h, w = all_records[0]["cam"].shape

    # ---- per-record stats ----
    n_missing = 0
    for r in all_records:
        cam   = r["cam"]
        total = cam.sum() + 1e-12
        pid   = int(r["patient_id"])
        eye   = r["eye"]

        disc_mask = None
        key = (pid, eye)
        if key in disc_lookup:
            disc_path, orig_size = disc_lookup[key]
            disc_mask = load_disc_mask(disc_path, orig_size, h, w)
        if disc_mask is None:
            n_missing += 1
            disc_frac = float("nan")
        else:
            disc_frac = float(cam[disc_mask].sum() / total)

        stat_rows.append({
            "patient_id": r["patient_id"],
            "fold":       r["fold"],
            "eye":        r["eye"],
            "true_name":  r["true_name"],
            "pred_name":  r["pred_name"],
            "correct":    r["correct"],
            "confidence": r["confidence"],
            "disc_frac":  disc_frac,
            "entropy":    attention_entropy(cam),
        })

    if n_missing:
        print(f"  Warning: {n_missing} records had no disc mask (disc_frac=NaN)")

    stats_df = pd.DataFrame(stat_rows)
    stats_path = out_dir / "attention_stats.csv"
    stats_df.to_csv(stats_path, index=False)
    print(f"Saved attention stats → {stats_path}")

    # ---- mean heatmaps ----
    npz_arrays = {}
    groups: dict[tuple, list[np.ndarray]] = {}
    for r in all_records:
        key = (r["eye"], r["true_name"], "correct" if r["correct"] else "incorrect")
        groups.setdefault(key, []).append(r["cam"])
    for r in all_records:
        key = (r["eye"], r["true_name"], "all")
        groups.setdefault(key, []).append(r["cam"])

    for (eye, cls, split), cams in groups.items():
        stack = np.stack(cams, axis=0)
        key_base = f"{eye}_{cls}_{split}"
        npz_arrays[f"{key_base}_mean"]  = stack.mean(axis=0).astype(np.float32)
        npz_arrays[f"{key_base}_std"]   = stack.std(axis=0).astype(np.float32)
        npz_arrays[f"{key_base}_count"] = np.array(len(cams))
        print(f"  {key_base}: N={len(cams)}")

    npz_path = out_dir / "mean_heatmaps.npz"
    np.savez_compressed(npz_path, **npz_arrays)
    print(f"Saved mean heatmaps → {npz_path}")


if __name__ == "__main__":
    main()