evaluate.py

#!/usr/bin/env python3
"""Leave-one-image-out cross-validation with repeated, perturbed tests per image.

    python evaluate.py cv    --data-dir DATA --runs-dir runs/cv --iters 3000
    python evaluate.py final --data-dir DATA --weights weights/cellquant.pt

``cv`` trains one model per annotated image with that image held out (unless the
fold checkpoint already exists), then tests the held-out image many times:

  * base run (no TTA), TTA x4 and TTA x8
  * 8 dihedral variants (rotate/flip the image, predict, undo the transform)
  * 6 photometric perturbations (per-channel gain, gamma, noise; fixed seeds)
  * 3 identical repeats to prove determinism
  * post-processing rule variants (how puncta get their type and how nuclei
    get the "Nostril" attribute)

``final`` reports the fit of one model on its own training images (NOT a
generalisation estimate) and writes overlays for un-annotated images.
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import subprocess
import sys
import time
from dataclasses import asdict, replace

import numpy as np
from PIL import Image

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from cellquant.annotations import discover, load_image, rasterize  # noqa: E402
from cellquant.infer import Predictor  # noqa: E402
from cellquant.labels import ALL_LABELS, CELL_LABELS, NOSTRIL_LABELS  # noqa: E402
from cellquant.metrics import count_errors, evaluate_cells, evaluate_nostrils, micro_average  # noqa: E402
from cellquant.postprocess import PostConfig, Result, analyze, render_overlay  # noqa: E402

DIHEDRAL = [(k, f) for f in (False, True) for k in range(4)]


def load_post_config(spec: str) -> dict:
    """``--post-config`` argument -> dict. Either a JSON file / literal holding
    PostConfig field overrides for every image ({"min_peak_dist": 0.63}) or a
    mapping image key -> overrides ({"C100": {"min_peak_dist": 0.61}, ...}),
    e.g. thresholds chosen per fold on that fold's training images."""
    if not spec:
        return {}
    if os.path.exists(spec):
        with open(spec) as f:
            d = json.load(f)
    else:
        d = json.loads(spec)
    if not isinstance(d, dict):
        raise SystemExit("--post-config must be a JSON object")
    return d


def post_config_for(overrides: dict, key: str) -> PostConfig:
    """PostConfig for one image from the --post-config dict (see load_post_config)."""
    if overrides and all(isinstance(v, dict) for v in overrides.values()):
        o = overrides.get(key, {})
    else:
        o = overrides
    unknown = set(o) - set(PostConfig.__dataclass_fields__)
    if unknown:
        raise SystemExit(f"--post-config: unknown PostConfig fields {sorted(unknown)}")
    return PostConfig(**o)


def fwd_np(a: np.ndarray, k: int, flip: bool) -> np.ndarray:
    if flip:
        a = a[..., ::-1]
    return np.rot90(a, k, axes=(-2, -1))


def inv_np(a: np.ndarray, k: int, flip: bool) -> np.ndarray:
    a = np.rot90(a, -k, axes=(-2, -1))
    if flip:
        a = a[..., ::-1]
    return np.ascontiguousarray(a)


def photometric(img: np.ndarray, seed: int) -> np.ndarray:
    rng = np.random.default_rng(seed)
    x = img.astype(np.float32) / 255.0
    gain = rng.uniform(0.7, 1.4, size=3).astype(np.float32)
    gamma = float(rng.uniform(0.8, 1.25))
    x = np.clip(x * gain, 0, 1) ** gamma
    x = x + rng.normal(0, float(rng.uniform(0, 0.02)), size=x.shape).astype(np.float32)
    return (np.clip(x, 0, 1) * 255).round().astype(np.uint8)


def fp_analysis(img: np.ndarray, R, res: Result) -> dict:
    """Characterise unmatched predicted nuclei: how many overlap an annotated nucleus
    at all (outline/split disagreement) versus none (extra detection), and how
    bright/large they are compared with matched ones."""
    from scipy import ndimage as ndi
    from cellquant.metrics import match_by_iou
    m = match_by_iou(R.cell_inst, res.cell_inst, 0.5)
    mp = {p for _, p, _ in m}
    mg = {g for g, _, _ in m}
    n_pred = int(res.cell_inst.max())
    if n_pred == 0:
        return {"n_fp": 0}
    ids = np.arange(1, n_pred + 1)
    blue = ndi.mean(img[..., 2].astype(np.float32), res.cell_inst, ids)
    areas = np.bincount(res.cell_inst.ravel())[1:]
    gt_cover = ndi.mean((R.cell_inst > 0).astype(np.float32), res.cell_inst, ids)
    fp = np.array([i not in mp for i in ids])
    matched = ~fp
    fn_ids = [g for g in range(1, R.n_cells + 1) if g not in mg]
    fn_cover = [float((res.cell_inst[R.cell_inst == g] > 0).mean()) for g in fn_ids]
    return {
        "n_fp": int(fp.sum()),
        "n_fp_overlapping_gt": int((fp & (gt_cover > 0.3)).sum()),
        "n_fp_no_gt_overlap": int((fp & (gt_cover <= 0.3)).sum()),
        "fp_median_blue": float(np.median(blue[fp])) if fp.any() else None,
        "matched_median_blue": float(np.median(blue[matched])) if matched.any() else None,
        "fp_median_area": float(np.median(areas[fp])) if fp.any() else None,
        "matched_median_area": float(np.median(areas[matched])) if matched.any() else None,
        "n_fn": len(fn_ids),
        "n_fn_covered_by_prediction": int(sum(c > 0.3 for c in fn_cover)),
    }


def matched_pairs(R, res: Result) -> list[dict]:
    """Per matched nucleus (IoU >= 0.5): GT label and the network's attribute probabilities,
    so that attribute thresholds can be studied across all held-out images."""
    from cellquant.metrics import match_by_iou
    by_id = {c["id"]: c for c in res.cells}
    out = []
    for g, p, iou in match_by_iou(R.cell_inst, res.cell_inst, 0.5):
        c = by_id[p]
        out.append({"gt": R.cell_labels[g], "pred": c["label"], "p_antibody": c["p_antibody"],
                    "p_nostril": c["p_nostril"], "n_puncta_inside": c["n_puncta_inside"], "iou": iou})
    return out


def score(R, res: Result, gt_counts: dict, img: np.ndarray | None = None) -> dict:
    ce = evaluate_cells(R.cell_inst, R.cell_labels, res.cell_inst, res.cell_labels)
    ne = evaluate_nostrils(R.nostril_inst, R.nostril_labels, res.nostril_inst, res.nostril_labels)
    return {
        "cells": ce,
        "nostrils": ne,
        "counts": count_errors(gt_counts, res.counts),
        "fp_analysis": fp_analysis(img, R, res) if img is not None else None,
        "pairs": matched_pairs(R, res) if img is not None else None,
        "summary": {
            "cell_det_f1": ce["detection"]["f1"],
            "cell_det_precision": ce["detection"]["precision"],
            "cell_det_recall": ce["detection"]["recall"],
            "cell_mean_iou": ce["detection"]["mean_iou_matched"],
            "cell_det_f1_iou30": ce["detection"]["iou30"]["f1"],
            "cell_det_f1_centroid": ce["detection"]["centroid"]["f1"],
            "cell_attr_antibody": ce["attribute_accuracy"]["antibody"],
            "cell_attr_nostril": ce["attribute_accuracy"]["nostril"],
            "cell_attr_both": ce["attribute_accuracy"]["both"],
            "nos_det_f1": ne["detection"]["f1"],
            "nos_det_precision": ne["detection"]["precision"],
            "nos_det_recall": ne["detection"]["recall"],
            "nos_type_acc": ne["type_accuracy"],
            **{f"f1[{c}]": ce["per_class"][c]["f1"] for c in CELL_LABELS},
            **{f"f1[{c}]": ne["per_class"][c]["f1"] for c in NOSTRIL_LABELS},
        },
    }


def inst_hash(res: Result) -> str:
    h = hashlib.sha256()
    h.update(res.cell_inst.tobytes())
    h.update(res.nostril_inst.tobytes())
    h.update(json.dumps(res.counts, sort_keys=True).encode())
    return h.hexdigest()[:16]


def test_image(weights: str, img: np.ndarray, R, gt_counts: dict, device: str, out_prefix: str | None = None,
               cfg: PostConfig | None = None) -> dict:
    cfg = cfg or PostConfig()
    out: dict = {"runs": {}, "post_config": asdict(cfg)}
    p1 = Predictor(weights, device=device, tta=1)
    # base + determinism
    hashes = []
    for i in range(3):
        res, probs = p1.predict(img, cfg)
        hashes.append(inst_hash(res))
    base_res, base_probs = res, probs
    out["deterministic"] = len(set(hashes)) == 1
    out["hashes"] = hashes
    out["runs"]["base"] = score(R, base_res, gt_counts, img)
    out["base_counts"] = base_res.counts
    if out_prefix:
        Image.fromarray(render_overlay(img, base_res)).save(out_prefix + "_pred.png")
    # TTA
    for tta in (4, 8):
        res, _ = Predictor(weights, device=device, tta=tta).predict(img, cfg)
        out["runs"][f"tta{tta}"] = score(R, res, gt_counts)
    # dihedral robustness
    for k, f in DIHEDRAL:
        timg = np.ascontiguousarray(fwd_np(np.moveaxis(img, -1, 0), k, f))
        timg = np.moveaxis(timg, 0, -1)
        probs = p1.probs(timg)
        probs = {kk: inv_np(v, k, f) for kk, v in probs.items()}
        res = analyze(probs, cfg)
        out["runs"][f"dihedral_k{k}_f{int(f)}"] = score(R, res, gt_counts)
    # photometric perturbations
    for seed in range(6):
        res, _ = p1.predict(photometric(img, 1000 + seed), cfg)
        out["runs"][f"photometric_{seed}"] = score(R, res, gt_counts)
    # post-processing rule variants (same probabilities as the base run)
    for nr in ("head", "geometric"):
        for cr in ("head", "contains", "either"):
            res = analyze(base_probs, replace(cfg, nostril_rule=nr, cellnos_rule=cr))
            out["runs"][f"rule_nos={nr}_cell={cr}"] = score(R, res, gt_counts)
    return out


def aggregate(per_key: dict[str, dict]) -> dict:
    keys = sorted(per_key)
    agg: dict = {"images": keys}
    # micro-averaged per-class scores on the base run
    cells = [per_key[k]["runs"]["base"]["cells"]["per_class"] for k in keys]
    nos = [per_key[k]["runs"]["base"]["nostrils"]["per_class"] for k in keys]
    agg["micro_base"] = {**micro_average(cells, CELL_LABELS), **micro_average(nos, NOSTRIL_LABELS)}
    for variant in ("tta4", "tta8"):
        cells = [per_key[k]["runs"][variant]["cells"]["per_class"] for k in keys]
        nos = [per_key[k]["runs"][variant]["nostrils"]["per_class"] for k in keys]
        agg[f"micro_{variant}"] = {**micro_average(cells, CELL_LABELS), **micro_average(nos, NOSTRIL_LABELS)}
    # detection totals
    for variant in ("base", "tta4", "tta8"):
        tp = sum(per_key[k]["runs"][variant]["cells"]["detection"]["tp"] for k in keys)
        fp = sum(per_key[k]["runs"][variant]["cells"]["detection"]["fp"] for k in keys)
        fn = sum(per_key[k]["runs"][variant]["cells"]["detection"]["fn"] for k in keys)
        ntp = sum(per_key[k]["runs"][variant]["nostrils"]["detection"]["tp"] for k in keys)
        nfp = sum(per_key[k]["runs"][variant]["nostrils"]["detection"]["fp"] for k in keys)
        nfn = sum(per_key[k]["runs"][variant]["nostrils"]["detection"]["fn"] for k in keys)
        from cellquant.metrics import prf
        agg[f"detection_{variant}"] = {"cells": prf(tp, fp, fn), "nostrils": prf(ntp, nfp, nfn)}
        for loose in ("iou30", "centroid"):
            ltp = sum(per_key[k]["runs"][variant]["cells"]["detection"][loose]["tp"] for k in keys)
            lfp = sum(per_key[k]["runs"][variant]["cells"]["detection"][loose]["fp"] for k in keys)
            lfn = sum(per_key[k]["runs"][variant]["cells"]["detection"][loose]["fn"] for k in keys)
            agg[f"detection_{variant}"][f"cells_{loose}"] = prf(ltp, lfp, lfn)
    # per-image summary table + robustness (mean/std over dihedral & photometric runs)
    summ = {}
    for k in keys:
        runs = per_key[k]["runs"]
        rob = [v["summary"] for n, v in runs.items() if n.startswith(("dihedral", "photometric"))]
        rob_stats = {}
        for m in ("cell_det_f1", "cell_attr_both", "nos_det_f1", "nos_type_acc"):
            vals = np.array([r[m] for r in rob])
            rob_stats[m] = {"mean": float(vals.mean()), "std": float(vals.std()), "min": float(vals.min()), "max": float(vals.max()), "n": int(len(vals))}
        summ[k] = {"base": runs["base"]["summary"], "tta4": runs["tta4"]["summary"], "tta8": runs["tta8"]["summary"],
                   "robustness": rob_stats, "deterministic": per_key[k]["deterministic"],
                   "counts": runs["base"]["counts"], "fp_analysis": runs["base"].get("fp_analysis")}
    agg["per_image"] = summ
    # rule variants: mean cell_attr_both and nostril type acc across images
    rules = {}
    for n in per_key[keys[0]]["runs"]:
        if n.startswith("rule_"):
            nm = [per_key[k]["runs"][n]["cells"]["attribute_accuracy"]["n_matched"] for k in keys]
            ok = [per_key[k]["runs"][n]["cells"]["attribute_accuracy"]["both"] * m for k, m in zip(keys, nm)]
            ntp = [per_key[k]["runs"][n]["nostrils"]["detection"]["tp"] for k in keys]
            tok = [per_key[k]["runs"][n]["nostrils"]["type_accuracy"] * t for k, t in zip(keys, ntp)]
            rules[n] = {
                "pooled_cell_attr_both": float(sum(ok) / max(1, sum(nm))),
                "pooled_nos_type_acc": float(sum(tok) / max(1, sum(ntp))),
                "micro_f1": {**micro_average([per_key[k]["runs"][n]["cells"]["per_class"] for k in keys], CELL_LABELS),
                             **micro_average([per_key[k]["runs"][n]["nostrils"]["per_class"] for k in keys], NOSTRIL_LABELS)},
            }
    agg["rule_variants"] = rules
    agg["all_deterministic"] = all(per_key[k]["deterministic"] for k in keys)
    # antibody threshold sweep on all held-out matched nuclei (base run)
    pairs = [p for k in keys for p in (per_key[k]["runs"]["base"].get("pairs") or [])]
    if pairs:
        from cellquant.labels import cell_attributes
        y = np.array([cell_attributes(p["gt"])[0] for p in pairs])      # True = antibody +
        s_ = np.array([p["p_antibody"] for p in pairs])
        sweep = []
        for thr in np.round(np.arange(0.3, 0.91, 0.05), 2):
            pred = s_ > thr
            tp_neg = int((~pred & ~y).sum()); fp_neg = int((~pred & y).sum()); fn_neg = int((pred & ~y).sum())
            from cellquant.metrics import prf
            r = prf(tp_neg, fp_neg, fn_neg)
            sweep.append({"thr": float(thr), "neg_precision": r["precision"], "neg_recall": r["recall"], "neg_f1": r["f1"],
                          "accuracy": float((pred == y).mean())})
        agg["antibody_threshold_sweep"] = {"n_matched": int(len(y)), "n_negative": int((~y).sum()), "sweep": sweep}
    return agg


def fmt_pct(x: float) -> str:
    return f"{100*x:5.1f}%"


def write_report(agg: dict, path: str, title: str, note: str = "") -> None:
    L = [f"# {title}", ""]
    if note:
        L += [note, ""]
    L += ["## Micro-averaged per-class scores (base run, no TTA)", "",
          "| Label | GT | Pred | TP | FP | FN | Precision | Recall | F1 |", "|---|---:|---:|---:|---:|---:|---:|---:|---:|"]
    for lab in ALL_LABELS:
        r = agg["micro_base"][lab]
        L.append(f"| {lab} | {r['gt_count']} | {r['pred_count']} | {r['tp']} | {r['fp']} | {r['fn']} | {fmt_pct(r['precision'])} | {fmt_pct(r['recall'])} | {fmt_pct(r['f1'])} |")
    L += ["", "Class-agnostic detection (IoU >= 0.5 for nuclei; centroid distance for puncta):", ""]
    for v in ("base", "tta4", "tta8"):
        c, n = agg[f"detection_{v}"]["cells"], agg[f"detection_{v}"]["nostrils"]
        c3, cc = agg[f"detection_{v}"]["cells_iou30"], agg[f"detection_{v}"]["cells_centroid"]
        L.append(f"- {v}: nuclei P {fmt_pct(c['precision'])} R {fmt_pct(c['recall'])} F1 {fmt_pct(c['f1'])} (IoU≥0.5); F1 {fmt_pct(c3['f1'])} (IoU≥0.3); F1 {fmt_pct(cc['f1'])} (centroid-in-nucleus) | puncta P {fmt_pct(n['precision'])} R {fmt_pct(n['recall'])} F1 {fmt_pct(n['f1'])}")
    L += ["", "## Per-image results (each image held out of training)", "",
          "| Image | Nuclei F1 (IoU≥0.5) | F1 (IoU≥0.3) | F1 (centroid) | mean IoU | Antibody acc | Nostril-attr acc | Puncta F1 | Puncta type acc | Deterministic |", "|---|---:|---:|---:|---:|---:|---:|---:|---:|---|"]
    for k, s in agg["per_image"].items():
        b = s["base"]
        L.append(f"| {k} | {fmt_pct(b['cell_det_f1'])} | {fmt_pct(b['cell_det_f1_iou30'])} | {fmt_pct(b['cell_det_f1_centroid'])} | {b['cell_mean_iou']:.3f} | {fmt_pct(b['cell_attr_antibody'])} | {fmt_pct(b['cell_attr_nostril'])} | {fmt_pct(b['nos_det_f1'])} | {fmt_pct(b['nos_type_acc'])} | {'yes' if s['deterministic'] else 'NO'} |")
    L += ["", "## Repeated tests per image (8 rotations/flips + 6 photometric perturbations = 14 runs)", "",
          "| Image | Nuclei F1 mean ± sd [min, max] | Attr acc mean ± sd | Puncta F1 mean ± sd [min, max] |", "|---|---|---|---|"]
    for k, s in agg["per_image"].items():
        r = s["robustness"]
        a, b, c = r["cell_det_f1"], r["cell_attr_both"], r["nos_det_f1"]
        L.append(f"| {k} | {100*a['mean']:.1f} ± {100*a['std']:.1f} [{100*a['min']:.1f}, {100*a['max']:.1f}] | {100*b['mean']:.1f} ± {100*b['std']:.1f} | {100*c['mean']:.1f} ± {100*c['std']:.1f} [{100*c['min']:.1f}, {100*c['max']:.1f}] |")
    L += ["", "## Why nuclei F1 is below 100 %: unmatched predictions and misses (base run)", "",
          "| Image | FP total | FP overlapping a GT nucleus | FP with no GT overlap | FP / matched median blue | FP / matched median area | FN total | FN covered by a prediction |", "|---|---:|---:|---:|---|---|---:|---:|"]
    for k, s in agg["per_image"].items():
        f = s.get("fp_analysis") or {}
        if not f or f.get("n_fp") is None:
            continue
        L.append(f"| {k} | {f['n_fp']} | {f.get('n_fp_overlapping_gt', 0)} | {f.get('n_fp_no_gt_overlap', 0)} | {f.get('fp_median_blue') or 0:.0f} / {f.get('matched_median_blue') or 0:.0f} | {f.get('fp_median_area') or 0:.0f} / {f.get('matched_median_area') or 0:.0f} | {f.get('n_fn', 0)} | {f.get('n_fn_covered_by_prediction', 0)} |")
    L += ["", "## Counts per image (ground truth / predicted, base run)", "",
          "| Image | " + " | ".join(ALL_LABELS) + " |", "|---|" + "---|" * len(ALL_LABELS)]
    for k, s in agg["per_image"].items():
        L.append(f"| {k} | " + " | ".join(f"{s['counts'][lab]['gt']} / {s['counts'][lab]['pred']}" for lab in ALL_LABELS) + " |")
    L += ["", "## Post-processing rule variants (micro F1 across held-out images; accuracies pooled over matched objects)", "",
          "| Variant | " + " | ".join(ALL_LABELS) + " | attr acc (matched nuclei) | puncta type acc (matched puncta) |", "|---|" + "---|" * (len(ALL_LABELS) + 2)]
    for n, r in agg["rule_variants"].items():
        L.append(f"| {n} | " + " | ".join(fmt_pct(r['micro_f1'][lab]['f1']) for lab in ALL_LABELS) + f" | {fmt_pct(r['pooled_cell_attr_both'])} | {fmt_pct(r['pooled_nos_type_acc'])} |")
    if agg.get("antibody_threshold_sweep"):
        sw = agg["antibody_threshold_sweep"]
        L += ["", f"## Antibody threshold sweep on the {sw['n_matched']} matched held-out nuclei ({sw['n_negative']} antibody-negative)", "",
              "| threshold | Antibody − precision | Antibody − recall | Antibody − F1 | overall accuracy |", "|---:|---:|---:|---:|---:|"]
        for r in sw["sweep"]:
            L.append(f"| {r['thr']:.2f} | {fmt_pct(r['neg_precision'])} | {fmt_pct(r['neg_recall'])} | {fmt_pct(r['neg_f1'])} | {fmt_pct(r['accuracy'])} |")
    L.append("")
    with open(path, "w") as f:
        f.write("\n".join(L))


def cmd_cv(a):
    samples, _ = discover(a.data_dir)
    keys = a.keys.split(",") if a.keys else [s.key for s in samples]
    os.makedirs(a.runs_dir, exist_ok=True)
    overrides = load_post_config(a.post_config)
    per_key = {}
    for s in samples:
        if s.key not in keys:
            continue
        ck = os.path.join(a.runs_dir, f"fold_{s.key}.pt")
        if not os.path.exists(ck) or a.retrain:
            cmd = [sys.executable, "-I", os.path.join(os.path.dirname(__file__), "train.py"), "--data-dir", a.data_dir,
                   "--holdout", s.key, "--iters", str(a.iters), "--out", ck, "--seed", str(a.seed), "--device", a.device]
            print("TRAIN:", " ".join(cmd), flush=True)
            with open(os.path.join(a.runs_dir, f"train_{s.key}.log"), "w") as log:
                subprocess.run(cmd, check=True, stdout=log, stderr=subprocess.STDOUT)
        from cellquant.model import load_checkpoint
        _, meta = load_checkpoint(ck)
        if meta.get("holdout") != s.key or s.key in meta.get("train_keys", []):
            raise SystemExit(f"{ck}: trained with holdout={meta.get('holdout')!r}, train_keys={meta.get('train_keys')}; "
                             f"refusing to score {s.key} with it (delete the file or use --retrain)")
        eval_json = os.path.join(a.runs_dir, f"{s.key}_eval.json")
        if a.reuse_eval and os.path.exists(eval_json) and os.path.getmtime(eval_json) > os.path.getmtime(ck):
            with open(eval_json) as f:
                per_key[s.key] = json.load(f)
            b = per_key[s.key]["runs"]["base"]["summary"]
            print(f"[{s.key}] reused {os.path.basename(eval_json)}: nuclei F1 {b['cell_det_f1']:.3f} attr {b['cell_attr_both']:.3f} puncta F1 {b['nos_det_f1']:.3f}", flush=True)
            continue
        img = load_image(s.image_path)
        R = rasterize(s.annotation, img.shape[0], img.shape[1])
        t = time.time()
        cfg = post_config_for(overrides, s.key)
        if overrides:
            changed = {k: v for k, v in asdict(cfg).items() if v != getattr(PostConfig(), k)}
            print(f"[{s.key}] post-processing overrides: {changed}", flush=True)
        res = test_image(ck, img, R, s.annotation.counts(), a.device, out_prefix=os.path.join(a.runs_dir, s.key), cfg=cfg)
        res["train_keys"] = list(meta.get("train_keys", []))
        res["checkpoint"] = os.path.basename(ck)
        per_key[s.key] = res
        b = res["runs"]["base"]["summary"]
        print(f"[{s.key}] held-out: nuclei F1 {b['cell_det_f1']:.3f} attr {b['cell_attr_both']:.3f} puncta F1 {b['nos_det_f1']:.3f} type {b['nos_type_acc']:.3f} deterministic={res['deterministic']} ({time.time()-t:.0f}s)", flush=True)
        with open(os.path.join(a.runs_dir, f"{s.key}_eval.json"), "w") as f:
            json.dump(res, f, indent=1)
    agg = aggregate(per_key)
    totals = {lab: sum(s.annotation.count(lab) for s in samples) for lab in ALL_LABELS}
    agg["dataset_totals"] = totals
    agg["post_config"] = {k: per_key[k].get("post_config", asdict(PostConfig())) for k in per_key}
    with open(os.path.join(a.runs_dir, "cv_results.json"), "w") as f:
        json.dump({"aggregate": agg, "per_image": per_key}, f, indent=1)
    n_nuc = sum(totals[c] for c in CELL_LABELS)
    n_pun = sum(totals[c] for c in NOSTRIL_LABELS)
    write_report(agg, os.path.join(a.runs_dir, "cv_report.md"), "Leave-one-image-out cross-validation",
                 f"{len(per_key)} folds over {len(samples)} annotated images ({n_nuc} nucleus and {n_pun} punctum outlines); "
                 f"each image scored by a model whose checkpoint records that image as held out. Iterations per fold: {a.iters}.")
    print("wrote", os.path.join(a.runs_dir, "cv_report.md"))


def cmd_final(a):
    samples, unann = discover(a.data_dir)
    os.makedirs(a.out_dir, exist_ok=True)
    overrides = load_post_config(a.post_config)
    pred = Predictor(a.weights, device=a.device, tta=1)
    rows = {"post_config": asdict(post_config_for(overrides, ""))}
    for s in samples:
        img = load_image(s.image_path)
        R = rasterize(s.annotation, img.shape[0], img.shape[1])
        res, _ = pred.predict(img, post_config_for(overrides, s.key))
        rows[s.key] = score(R, res, s.annotation.counts())
        Image.fromarray(render_overlay(img, res)).save(os.path.join(a.out_dir, f"{s.key}_pred.png"))
        b = rows[s.key]["summary"]
        print(f"[{s.key}] TRAINING-SET fit: nuclei F1 {b['cell_det_f1']:.3f} attr {b['cell_attr_both']:.3f} puncta F1 {b['nos_det_f1']:.3f}")
    for p in unann:
        img = load_image(p)
        res, _ = pred.predict(img, post_config_for(overrides, ""))
        key = os.path.basename(p).split(" ")[0]
        Image.fromarray(render_overlay(img, res)).save(os.path.join(a.out_dir, f"{key}_unlabeled_pred.png"))
        rows[f"{key} (no annotation)"] = {"counts": res.counts}
        print(f"[{key}] un-annotated image counts: {res.counts}")
    with open(os.path.join(a.out_dir, "final_fit.json"), "w") as f:
        json.dump(rows, f, indent=1)


def main(argv=None):
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)
    c = sub.add_parser("cv")
    c.add_argument("--data-dir", required=True)
    c.add_argument("--runs-dir", default="runs/cv")
    c.add_argument("--iters", type=int, default=3000)
    c.add_argument("--keys", default="")
    c.add_argument("--seed", type=int, default=0)
    c.add_argument("--device", default="auto")
    c.add_argument("--retrain", action="store_true")
    c.add_argument("--reuse-eval", action="store_true", help="reuse <key>_eval.json files newer than their checkpoint")
    c.add_argument("--post-config", default="", help="JSON (file or literal) of PostConfig overrides, global or per image key")
    f = sub.add_parser("final")
    f.add_argument("--data-dir", required=True)
    f.add_argument("--weights", required=True)
    f.add_argument("--out-dir", default="reports/final")
    f.add_argument("--device", default="auto")
    f.add_argument("--post-config", default="", help="JSON (file or literal) of PostConfig overrides")
    a = ap.parse_args(argv)
    {"cv": cmd_cv, "final": cmd_final}[a.cmd](a)


if __name__ == "__main__":
    main()