scripts/rule_study.py

#!/usr/bin/env python3
"""Study, on the TRAINING images of every cross-validation fold, which per-nucleus
score separates predictions that match an annotated nucleus from predictions that
match nothing, and pick a rejection threshold from that training-set fit.

    python scripts/rule_study.py --data-dir DATA --runs-dir runs/cv_v2 --out runs/cv_v3/rule_study

For fold k the model ``fold_k.pt`` (trained without image k) predicts every image;
the predictions on its eight training images form the pool on which the rule is
chosen (``selection = "train_fit"``). The prediction on the held-out image k is
stored as well, but only to apply the chosen rule later (evaluate.py) and to
report, after the fact, how the two pools differ. A second, nested selection
(``selection = "other_folds_heldout"``: the held-out predictions of the other
eight folds, i.e. never image k) is reported for comparison.

Per predicted nucleus the candidate scores are: the peak and mean of the predicted
distance map, the mean foreground and interior probability, the area, and the mean
blue intensity relative to the median predicted nucleus of the same image.
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time

import numpy as np
from scipy import ndimage as ndi

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
from cellquant.annotations import discover, load_image, rasterize  # noqa: E402
from cellquant.infer import Predictor  # noqa: E402
from cellquant.metrics import iou_matrix, match_by_centroid_in_mask, match_by_iou, prf  # noqa: E402
from cellquant.postprocess import PostConfig, analyze  # noqa: E402

FEATURES = ("max_dist", "p90_dist", "mean_dist", "mean_pfg", "mean_pint", "area", "blue_rel")


def instance_table(img: np.ndarray, probs: dict, inst: np.ndarray, R) -> dict:
    """Per predicted nucleus: candidate scores and whether it matches an annotated nucleus."""
    n = int(inst.max())
    if n == 0:
        return {k: [] for k in FEATURES + ("gt_cover", "best_iou", "matched50", "matched_centroid", "blue")} | {"n_gt": R.n_cells}
    ids = np.arange(1, n + 1)
    p_fg = probs["cell"][1] + probs["cell"][2]
    dist = probs.get("dist")
    if dist is None:
        dist = probs["cell"][1]
    area = np.bincount(inst.ravel())[1:].astype(np.float64)
    blue = np.asarray(ndi.mean(img[..., 2].astype(np.float32), inst, ids), dtype=np.float64)
    med = float(np.median(blue)) if n else 1.0
    p90 = np.asarray(ndi.labeled_comprehension(dist, inst, ids, lambda v: np.percentile(v, 90), np.float64, 0.0))
    iou = iou_matrix(R.cell_inst, inst)
    best_iou = iou.max(0) if iou.size else np.zeros(n)
    m50 = {p for _, p, _ in match_by_iou(R.cell_inst, inst, 0.5)}
    mc = {p for _, p, _ in match_by_centroid_in_mask(R.cell_inst, inst)}
    return {
        "n_gt": R.n_cells,
        "max_dist": np.asarray(ndi.maximum(dist, inst, ids), dtype=np.float64).tolist(),
        "p90_dist": p90.tolist(),
        "mean_dist": np.asarray(ndi.mean(dist, inst, ids), dtype=np.float64).tolist(),
        "mean_pfg": np.asarray(ndi.mean(p_fg, inst, ids), dtype=np.float64).tolist(),
        "mean_pint": np.asarray(ndi.mean(probs["cell"][1], inst, ids), dtype=np.float64).tolist(),
        "area": area.tolist(),
        "blue": blue.tolist(),
        "blue_rel": (blue / max(med, 1e-6)).tolist(),
        "gt_cover": np.asarray(ndi.mean((R.cell_inst > 0).astype(np.float32), inst, ids), dtype=np.float64).tolist(),
        "best_iou": np.asarray(best_iou, dtype=np.float64).tolist(),
        "matched50": [bool(i in m50) for i in ids],
        "matched_centroid": [bool(i in mc) for i in ids],
    }


def pool(tables: list[dict]) -> dict[str, np.ndarray]:
    out = {k: np.concatenate([np.asarray(t[k], dtype=np.float64) for t in tables]) if any(len(t[k]) for t in tables) else np.zeros(0)
           for k in FEATURES + ("gt_cover", "matched50", "matched_centroid")}
    out["n_gt"] = int(sum(t["n_gt"] for t in tables))
    return out


def f1_after_filter(P: dict, keep: np.ndarray, key: str = "matched50") -> dict:
    m = P[key] > 0.5
    tp = int((m & keep).sum())
    fp = int((~m & keep).sum())
    fn = P["n_gt"] - tp
    r = prf(tp, fp, fn)
    r["n_kept"] = int(keep.sum())
    r["removed_tp"] = int((m & ~keep).sum())
    r["removed_fp"] = int((~m & ~keep).sum())
    r["removed_fp_no_overlap"] = int((~m & ~keep & (P["gt_cover"] <= 0.3)).sum())
    return r


def thresholds_for(feature: str, values: np.ndarray) -> np.ndarray:
    if feature == "area":
        return np.arange(50, 401, 10, dtype=np.float64)
    if feature == "blue_rel":
        return np.round(np.arange(0.0, 1.01, 0.02), 2)
    return np.round(np.arange(0.0, 1.0001, 0.01), 2)


def select(P: dict, feature: str) -> dict:
    """Threshold on one feature that maximises nucleus F1 (IoU >= 0.5) on the pool;
    ties go to the smallest threshold (least filtering)."""
    v = P[feature]
    base = f1_after_filter(P, np.ones(len(v), bool))
    best = {"thr": 0.0, **base}
    curve = []
    for thr in thresholds_for(feature, v):
        r = f1_after_filter(P, v >= thr)
        curve.append({"thr": float(thr), "f1": r["f1"], "precision": r["precision"], "recall": r["recall"],
                      "removed_tp": r["removed_tp"], "removed_fp": r["removed_fp"]})
        if r["f1"] > best["f1"] + 1e-12:
            best = {"thr": float(thr), **r}
    return {"feature": feature, "base_f1": base["f1"], "best": best, "curve": curve}


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--data-dir", required=True)
    ap.add_argument("--runs-dir", default="runs/cv_v2")
    ap.add_argument("--out", default="runs/cv_v3/rule_study")
    ap.add_argument("--device", default="auto")
    ap.add_argument("--keys", default="")
    ap.add_argument("--rule-feature", default="max_dist", choices=FEATURES,
                    help="score the rule is built on; its per-fold thresholds go to <out>_post_config.json")
    ap.add_argument("--weights", default="", help="also choose the threshold for this all-image model from its own training fit")
    a = ap.parse_args(argv)
    samples, _ = discover(a.data_dir)
    keys = a.keys.split(",") if a.keys else [s.key for s in samples]
    os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True)
    images, rasters = {}, {}
    for s in samples:
        images[s.key] = load_image(s.image_path)
        rasters[s.key] = rasterize(s.annotation, images[s.key].shape[0], images[s.key].shape[1])
    cfg = PostConfig()
    tables: dict[str, dict[str, dict]] = {}   # fold -> image -> table
    raw_path = a.out + "_instances.json"
    if os.path.exists(raw_path):
        tables = json.load(open(raw_path))
        print(f"loaded {raw_path}")
    for k in keys:
        if k in tables and len(tables[k]) == len(samples):
            continue
        ck = os.path.join(a.runs_dir, f"fold_{k}.pt")
        pred = Predictor(ck, device=a.device, tta=1)
        assert pred.meta.get("holdout") == k and k not in pred.meta.get("train_keys", []), ck
        t0 = time.time()
        tables[k] = {}
        for s in samples:
            res, probs = pred.predict(images[s.key], cfg)
            tables[k][s.key] = instance_table(images[s.key], probs, res.cell_inst, rasters[s.key])
            tables[k][s.key]["role"] = "heldout" if s.key == k else "train"
        print(f"fold {k}: {len(samples)} images in {time.time() - t0:.0f}s", flush=True)
        with open(raw_path, "w") as f:
            json.dump(tables, f)
    # ---- selection per fold
    report = {"folds": {}, "features": FEATURES}
    for k in keys:
        train_tabs = [t for kk, t in tables[k].items() if kk != k]
        held_tab = tables[k][k]
        other_held = [tables[j][j] for j in keys if j != k]
        pools = {"train_fit": pool(train_tabs), "other_folds_heldout": pool(other_held)}
        PH = pool([held_tab])
        fold = {"pools": {}, "heldout_preview": {}}
        for name, P in pools.items():
            sel = {f: select(P, f) for f in FEATURES}
            best_f = max(FEATURES, key=lambda f: (sel[f]["best"]["f1"], -sel[f]["best"]["thr"]))
            fold["pools"][name] = {"n_pred": int(len(P["matched50"])), "n_gt": P["n_gt"],
                                   "n_fp": int((P["matched50"] < 0.5).sum()),
                                   "n_fp_no_overlap": int(((P["matched50"] < 0.5) & (P["gt_cover"] <= 0.3)).sum()),
                                   "base_f1": sel[best_f]["base_f1"], "chosen_feature": best_f,
                                   "chosen": sel[best_f]["best"], "per_feature": {f: sel[f]["best"] for f in FEATURES},
                                   "curves": {f: sel[f]["curve"] for f in FEATURES}}
            # how the chosen rule would act on the held-out image (preview only; evaluate.py gives the real numbers)
            thr = sel[best_f]["best"]["thr"]
            fold["heldout_preview"][name] = {"feature": best_f, "thr": thr,
                                             "before": f1_after_filter(PH, np.ones(len(PH[best_f]), bool)),
                                             "after": f1_after_filter(PH, PH[best_f] >= thr)}
        # distribution of each feature on matched / unmatched predictions, training pool vs held-out image
        dist = {}
        for f in FEATURES:
            d = {}
            for name, P in (("train_fit", pools["train_fit"]), ("heldout", PH)):
                m = P["matched50"] > 0.5
                nov = (~m) & (P["gt_cover"] <= 0.3)
                d[name] = {"matched_q10_q50": [float(np.percentile(P[f][m], 10)), float(np.median(P[f][m]))] if m.any() else None,
                           "unmatched_q50_q90": [float(np.median(P[f][~m])), float(np.percentile(P[f][~m], 90))] if (~m).any() else None,
                           "no_overlap_q50_q90": [float(np.median(P[f][nov])), float(np.percentile(P[f][nov], 90))] if nov.any() else None}
            dist[f] = d
        fold["distributions"] = dist
        report["folds"][k] = fold
    # ---- the rule that evaluate.py applies: one threshold per fold, chosen on that fold's training fit
    report["rule_feature"] = a.rule_feature
    post = {k: {"min_peak_dist": report["folds"][k]["pools"]["train_fit"]["per_feature"][a.rule_feature]["thr"]} for k in keys}
    with open(a.out + "_post_config.json", "w") as f:
        json.dump(post, f, indent=1)
    # ---- the all-image model: same procedure on all its training images
    if a.weights:
        pred = Predictor(a.weights, device=a.device, tta=1)
        tabs = []
        for s in samples:
            res, probs = pred.predict(images[s.key], cfg)
            tabs.append(instance_table(images[s.key], probs, res.cell_inst, rasters[s.key]))
        P = pool(tabs)
        sel = {f: select(P, f) for f in FEATURES}
        report["final"] = {"weights": a.weights, "n_pred": int(len(P["matched50"])), "n_gt": P["n_gt"],
                           "n_fp": int((P["matched50"] < 0.5).sum()), "base_f1": sel[a.rule_feature]["base_f1"],
                           "per_feature": {f: sel[f]["best"] for f in FEATURES}, "curves": {f: sel[f]["curve"] for f in FEATURES},
                           "chosen": {"feature": a.rule_feature, **sel[a.rule_feature]["best"]}}
        print(f"final model: {a.rule_feature} threshold {sel[a.rule_feature]['best']['thr']:.2f} "
              f"(training fit F1 {pct(sel[a.rule_feature]['base_f1'])} -> {pct(sel[a.rule_feature]['best']['f1'])})")
    with open(a.out + ".json", "w") as f:
        json.dump(report, f, indent=1)
    write_md(report, a.out + ".md", keys)
    print("wrote", a.out + ".md")


def pct(x):
    return f"{100 * x:.1f}%"


def write_md(rep: dict, path: str, keys: list[str]) -> None:
    L = ["# Rejection-rule study on the training images of each fold", "",
         "For every fold the model trained without image k predicts its eight training images; "
         "the per-nucleus score and threshold that maximise nucleus F1 (IoU >= 0.5) on those predictions "
         "are chosen there (`train_fit`). The nested alternative (`other_folds_heldout`) chooses on the held-out "
         "predictions of the other eight folds. Neither looks at image k. The 'held-out preview' columns apply the "
         "chosen rule to the stored prediction of image k; the real numbers come from evaluate.py.", ""]
    for name in ("train_fit", "other_folds_heldout"):
        L += [f"## Selection pool: {name}", "",
              "| Fold | pool pred / GT | pool FP (no overlap) | base F1 | chosen score | thr | pool F1 after | pool TP lost / FP removed | held-out F1 before -> after | held-out TP lost / FP removed (no overlap) |",
              "|---|---:|---:|---:|---|---:|---:|---|---|---|"]
        for k in keys:
            p = rep["folds"][k]["pools"][name]
            h = rep["folds"][k]["heldout_preview"][name]
            c = p["chosen"]
            L.append(f"| {k} | {p['n_pred']} / {p['n_gt']} | {p['n_fp']} ({p['n_fp_no_overlap']}) | {pct(p['base_f1'])} | {p['chosen_feature']} | {c['thr']:.2f} | {pct(c['f1'])} | {c['removed_tp']} / {c['removed_fp']} | "
                     f"{pct(h['before']['f1'])} -> {pct(h['after']['f1'])} | {h['after']['removed_tp']} / {h['after']['removed_fp']} ({h['after']['removed_fp_no_overlap']}) |")
        L += ["", "Best threshold per score (pool F1 after filtering; '—' = no threshold beats no filtering):", "",
              "| Fold | " + " | ".join(rep["features"]) + " |", "|---|" + "---|" * len(rep["features"])]
        for k in keys:
            p = rep["folds"][k]["pools"][name]
            cells = []
            for f in rep["features"]:
                b = p["per_feature"][f]
                cells.append("—" if b["f1"] <= p["base_f1"] + 1e-12 else f"{b['thr']:.2f} ({pct(b['f1'])})")
            L.append(f"| {k} | " + " | ".join(cells) + " |")
        L.append("")
    rf = rep.get("rule_feature", "max_dist")
    L += [f"## Rule applied by evaluate.py: `min_peak_dist` = the `{rf}` threshold of each fold's training fit", "",
          "| Fold | " + " | ".join(keys) + " |", "|---|" + "---:|" * len(keys),
          "| threshold | " + " | ".join(f"{rep['folds'][k]['pools']['train_fit']['per_feature'][rf]['thr']:.2f}" for k in keys) + " |", ""]
    if rep.get("final"):
        fi = rep["final"]
        L += [f"All-image model `{fi['weights']}`: {fi['n_pred']} predicted / {fi['n_gt']} annotated nuclei on its own training images, "
              f"{fi['n_fp']} unmatched; the same procedure picks `{rf}` >= {fi['chosen']['thr']:.2f} "
              f"(training-fit F1 {pct(fi['base_f1'])} -> {pct(fi['chosen']['f1'])}, {fi['chosen']['removed_tp']} matched and "
              f"{fi['chosen']['removed_fp']} unmatched predictions removed). This value is the shipped PostConfig default.", ""]
    L += ["## Score distributions: training pool vs held-out image (matched q10/q50; unmatched q50/q90; no-overlap unmatched q50/q90)", ""]
    for f in rep["features"]:
        L += [f"### {f}", "", "| Fold | train matched | train unmatched | train no-overlap | held-out matched | held-out unmatched | held-out no-overlap |", "|---|---|---|---|---|---|---|"]
        for k in keys:
            d = rep["folds"][k]["distributions"][f]
            def fmt(x):
                return "—" if x is None else f"{x[0]:.2f} / {x[1]:.2f}"
            L.append(f"| {k} | {fmt(d['train_fit']['matched_q10_q50'])} | {fmt(d['train_fit']['unmatched_q50_q90'])} | {fmt(d['train_fit']['no_overlap_q50_q90'])} | "
                     f"{fmt(d['heldout']['matched_q10_q50'])} | {fmt(d['heldout']['unmatched_q50_q90'])} | {fmt(d['heldout']['no_overlap_q50_q90'])} |")
        L.append("")
    with open(path, "w") as f:
        f.write("\n".join(L))


if __name__ == "__main__":
    main()