Spaces:
Running on Zero
Running on Zero
| #!/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() | |