Spaces:
Running on Zero
Running on Zero
| #!/usr/bin/env python3 | |
| """How far does test-time augmentation help? Held-out scores of the fold models for | |
| no TTA, 4x, 8x (the full dihedral group) and 8x combined with several image scales | |
| (24x = 3 scales, 40x = 5 scales), with the per-fold peak-distance thresholds. | |
| python scripts/tta_study.py --data-dir DATA --runs-dir runs/cv_v3 \ | |
| --post-config runs/cv_v3/rule_study_post_config.json --out reports/cv_v3/tta_study.md | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import json | |
| import os | |
| import sys | |
| import time | |
| import numpy as np | |
| 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 evaluate_cells, evaluate_nostrils, prf # noqa: E402 | |
| from cellquant.postprocess import PostConfig, analyze # noqa: E402 | |
| SETTINGS = { | |
| "none (1x)": dict(tta=1, scales=(1.0,)), | |
| "4x flips": dict(tta=4, scales=(1.0,)), | |
| "8x rotations+flips": dict(tta=8, scales=(1.0,)), | |
| "8x + 3 scales 0.8/1/1.25 (24x)": dict(tta=8, scales=(0.8, 1.0, 1.25)), | |
| "8x + 5 scales 0.8..1.25 (40x)": dict(tta=8, scales=(0.8, 0.9, 1.0, 1.1, 1.25)), | |
| } | |
| def main(argv=None): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--data-dir", required=True) | |
| ap.add_argument("--runs-dir", default="runs/cv_v3") | |
| ap.add_argument("--post-config", default="") | |
| ap.add_argument("--out", default="reports/cv_v3/tta_study.md") | |
| ap.add_argument("--device", default="auto") | |
| a = ap.parse_args(argv) | |
| overrides = json.load(open(a.post_config)) if a.post_config else {} | |
| samples, _ = discover(a.data_dir, verbose=False) | |
| res: dict[str, dict] = {name: {"cells": [], "nos": [], "cen": [], "n_pred": 0, "n_gt": 0, "sec": 0.0, "per_image": {}} for name in SETTINGS} | |
| for s in samples: | |
| img = load_image(s.image_path) | |
| R = rasterize(s.annotation, img.shape[0], img.shape[1]) | |
| cfg = PostConfig(**overrides.get(s.key, {})) | |
| ck = os.path.join(a.runs_dir, f"fold_{s.key}.pt") | |
| for name, st in SETTINGS.items(): | |
| pred = Predictor(ck, device=a.device, tta=st["tta"], scales=st["scales"]) | |
| t0 = time.time() | |
| probs = pred.probs(img) | |
| dt = time.time() - t0 | |
| out = analyze(probs, cfg) | |
| ce = evaluate_cells(R.cell_inst, R.cell_labels, out.cell_inst, out.cell_labels) | |
| ne = evaluate_nostrils(R.nostril_inst, R.nostril_labels, out.nostril_inst, out.nostril_labels) | |
| r = res[name] | |
| r["cells"].append(ce["detection"]); r["cen"].append(ce["detection"]["centroid"]); r["nos"].append(ne["detection"]) | |
| r["n_pred"] += ce["n_pred"]; r["n_gt"] += ce["n_gt"]; r["sec"] += dt | |
| r["per_image"][s.key] = ce["detection"]["f1"] | |
| print(f"[{s.key}] {name}: nuclei F1 {ce['detection']['f1']:.3f} puncta F1 {ne['detection']['f1']:.3f} ({dt:.1f}s)", flush=True) | |
| L = ["# Test-time augmentation on the held-out images (v3 rule, per-fold thresholds)", "", | |
| "Each row averages the network output over the listed transforms before post-processing; the dihedral group " | |
| "(4 rotations x flip) has exactly 8 members, so more than 8x means adding image scales (the image is resized, " | |
| "predicted, and the probabilities resized back). Micro-averaged over the nine held-out images; " | |
| "seconds = mean inference time per 1024x1024 image on this machine's GPU (the Space's CPU is roughly 20-40x slower).", "", | |
| "| Setting | passes | Nucleus P | R | F1 (IoU 0.5) | F1 (centroid) | Puncta P | R | F1 | predicted nuclei / annotated | s per image |", | |
| "|---|---:|---:|---:|---:|---:|---:|---:|---:|---|---:|"] | |
| summary = {} | |
| for name, st in SETTINGS.items(): | |
| r = res[name] | |
| def tot(rows): | |
| return prf(sum(x["tp"] for x in rows), sum(x["fp"] for x in rows), sum(x["fn"] for x in rows)) | |
| c, cc, n = tot(r["cells"]), tot(r["cen"]), tot(r["nos"]) | |
| passes = st["tta"] * len(st["scales"]) | |
| summary[name] = {"passes": passes, "cells": c, "centroid": cc, "puncta": n, "n_pred": r["n_pred"], "n_gt": r["n_gt"], | |
| "sec_per_image": r["sec"] / len(samples), "per_image_f1": r["per_image"]} | |
| L.append(f"| {name} | {passes} | {100*c['precision']:.1f} | {100*c['recall']:.1f} | **{100*c['f1']:.1f}** | {100*cc['f1']:.1f} | " | |
| f"{100*n['precision']:.1f} | {100*n['recall']:.1f} | **{100*n['f1']:.1f}** | {r['n_pred']} / {r['n_gt']} | {r['sec']/len(samples):.1f} |") | |
| L += ["", "## Nucleus F1 per held-out image", "", "| Image | " + " | ".join(SETTINGS) + " |", "|---|" + "---:|" * len(SETTINGS)] | |
| for s in samples: | |
| L.append(f"| {s.key} | " + " | ".join(f"{100*res[n]['per_image'][s.key]:.1f}" for n in SETTINGS) + " |") | |
| L.append("") | |
| os.makedirs(os.path.dirname(os.path.abspath(a.out)), exist_ok=True) | |
| with open(a.out, "w") as f: | |
| f.write("\n".join(L)) | |
| with open(os.path.splitext(a.out)[0] + ".json", "w") as f: | |
| json.dump(summary, f, indent=1) | |
| print("wrote", a.out) | |
| if __name__ == "__main__": | |
| main() | |