Blame for code/framework/report/aggregate.py

"""Aggregate per-seed metrics.json files into mean +- SD tables (over seeds).

Scans results/<exp_name>/**/seed*/metrics.json, groups by (dataset, protocol,
arch), and reports mean +- SD across seeds for every metric. Emits CSV, Markdown,
LaTeX. Overlap/clf metrics shown as percentages; HD95/ASSD as raw distances.

  python framework/report/aggregate.py --exp_name baselines [--out_root results]
"""
from __future__ import annotations

import os
import json
import glob
import argparse
from collections import defaultdict

import numpy as np

# (key, label, is_percent, higher_is_better)
METRICS = [
    ("dice", "Dice", True, True),
    ("iou", "IoU", True, True),
    ("hd95", "HD95", False, False),
    ("assd", "ASSD", False, False),
    ("sensitivity", "Sens", True, True),
    ("specificity", "Spec", True, True),
    ("precision", "Prec", True, True),
]


def load_runs(out_root, exp_name):
    runs = []
    for path in glob.glob(os.path.join(out_root, exp_name, "**", "seed*", "metrics.json"), recursive=True):
        try:
            with open(path) as f:
                runs.append(json.load(f))
        except Exception:
            pass
    return runs


# Display labels for split protocols (replaces opaque "fold01" with the real
# reporting protocol). Datasets with an official split keep "official"/"holdout";
# fold-1-of-k single splits are labelled "single-split"; PanNuke is reported as
# its official k-fold cross-validation (handled in summarize via fold-merge).
_PROTO_LABEL = {
    ("idridd_segmentation", "fold01"): "official",    # official 54/27; train re-split for val
    ("busi", "fold01"): "single-split",
    ("medsegdb_kits19", "fold01"): "single-split",
    ("pannuke_semantic", "fold01"): "single-split",   # until the other folds are added
}

# Datasets reported as official k-fold cross-validation (mean±SD OVER folds, not
# over seeds): folds are merged into one row when >1 fold is present.
_CV_DATASETS = {"pannuke_semantic"}


def _proto_label(dataset, protocol):
    return _PROTO_LABEL.get((dataset, protocol), protocol)


def _agg_over(items, key):
    vals = np.array([it.get("metrics", {}).get(f"{key}_mean", np.nan) for it in items], np.float64)
    vals = vals[~np.isnan(vals)]
    return (float(vals.mean()), float(vals.std())) if vals.size else (float("nan"), float("nan"))


def summarize(runs):
    # (dataset, arch) -> protocol -> [runs]
    by_da = defaultdict(lambda: defaultdict(list))
    for d in runs:
        by_da[(d.get("dataset"), d.get("arch"))][d.get("protocol")].append(d)
    rows = []
    for (dataset, arch), proto_map in sorted(by_da.items()):
        protos = sorted(p for p in proto_map if p is not None)
        row = {"dataset": dataset, "arch": arch}
        if dataset in _CV_DATASETS and len(protos) > 1:
            # k-fold CV: collapse seeds within each fold, then mean±SD OVER folds
            row["protocol"] = f"{len(protos)}-fold"
            row["n_seeds"] = len(protos)
            for key, _, _, _ in METRICS:
                fold_means = [m for m in (_agg_over(proto_map[p], key)[0] for p in protos)
                              if not np.isnan(m)]
                fm = np.array(fold_means, np.float64)
                row[f"{key}_mean"] = float(fm.mean()) if fm.size else float("nan")
                row[f"{key}_sd"] = float(fm.std()) if fm.size else float("nan")
        else:
            # single fixed split: mean±SD over seeds
            proto = protos[0] if protos else None
            items = proto_map.get(proto, [])
            row["protocol"] = _proto_label(dataset, proto)
            row["n_seeds"] = len(items)
            for key, _, _, _ in METRICS:
                row[f"{key}_mean"], row[f"{key}_sd"] = _agg_over(items, key)
        rows.append(row)
    return rows


def _fmt(row, key, pct):
    m, s = row[f"{key}_mean"], row[f"{key}_sd"]
    if np.isnan(m):
        return "—"
    return f"{m*100:.2f}±{s*100:.2f}" if pct else f"{m:.2f}±{s:.2f}"


def to_markdown(rows):
    cols = ["dataset", "protocol", "arch", "seeds"] + [lbl for _, lbl, _, _ in METRICS]
    out = "| " + " | ".join(cols) + " |\n" + "|" + "---|" * len(cols) + "\n"
    for r in rows:
        cells = [r["dataset"], r["protocol"], r["arch"], str(r["n_seeds"])]
        cells += [_fmt(r, k, p) for k, _, p, _ in METRICS]
        out += "| " + " | ".join(cells) + " |\n"
    return out


def to_csv(rows):
    cols = ["dataset", "protocol", "arch", "n_seeds"]
    for k, _, _, _ in METRICS:
        cols += [f"{k}_mean", f"{k}_sd"]
    out = ",".join(cols) + "\n"
    for r in rows:
        out += ",".join(str(r[c]) for c in cols) + "\n"
    return out


def to_latex(rows):
    spec = "lll c" + "c" * len(METRICS)
    out = "\\begin{tabular}{" + spec + "}\n\\toprule\n"
    out += "Dataset & Protocol & Method & Seeds & " + " & ".join(lbl for _, lbl, _, _ in METRICS) + " \\\\\n\\midrule\n"
    for r in rows:
        cells = [str(r["dataset"]), str(r["protocol"]), str(r["arch"]), str(r["n_seeds"])]
        cells += [_fmt(r, k, p).replace("±", "$\\pm$") for k, _, p, _ in METRICS]
        out += " & ".join(cells) + " \\\\\n"
    out += "\\bottomrule\n\\end{tabular}\n"
    return out


_ARCH_ORDER = ["unet", "unetpp", "deeplabv3plus", "attention_unet", "transunet", "swinunet",
               "nnunet", "umamba"]


# (#, dataset, modality, target, classes, channels, image size, reported protocol, train/val/test)
_DATASETS_INFO = [
    ("1", "CVC-ClinicDB", "Colonoscopy (endoscopy)", "Polyp", "2", "RGB", "384×288", "official", "490 / 61 / 61"),
    ("2", "Kvasir-SEG", "GI endoscopy", "Polyp", "2", "RGB", "~622×529 (var)", "official", "800 / 100 / 100"),
    ("3", "FIVES", "Retinal fundus", "Vessel", "2", "RGB", "2048×2048", "official", "480 / 120 / 200"),
    ("4", "BUSI", "Breast ultrasound", "Tumor", "2", "grayscale¹", "variable", "single-split²", "545 / 78 / 157"),
    ("5", "REFUGE2", "Retinal fundus", "Optic disc & cup", "3", "RGB", "~2124×2056", "official", "400 / 400 / 400"),
    ("6", "ACDC", "Cardiac MRI (2D slices)", "RV / Myo / LV", "4", "grayscale", "~240×256 (var)", "official", "136 / 210 / 380"),
    ("7", "IDRiD", "Retinal fundus", "DR lesions (4) + optic disc", "6", "RGB", "4288×2848", "official", "43 / 11 / 27"),
    ("8", "PanNuke", "Histopathology (H&E)", "Nuclei (5 types)", "6", "RGB", "256×256", "official 3-fold CV", "~2.7k / 2.6k / 2.6k per fold"),
    ("9", "ISIC2018", "Dermoscopy", "Skin lesion", "2", "RGB", "256×256", "holdout", "2582 / 369 / 737"),
    ("10", "KiTS19", "Kidney CT (2D slices)", "Kidney (binary)", "2", "grayscale¹", "256×256", "single-split²", "2832 / 479 / 705"),
]

_METHODS_INFO = [
    ("UNet", "CNN encoder–decoder", "SMP, ResNet-50 encoder (ImageNet-pretrained)"),
    ("UNet++", "Nested UNet", "SMP, ResNet-50 (ImageNet)"),
    ("DeepLabV3+", "Atrous CNN", "SMP, ResNet-50 (ImageNet)"),
    ("Attention-UNet", "Attention-gated UNet", "Re-implemented, trained from scratch"),
    ("TransUNet", "CNN–Transformer hybrid", "R50-ViT-B/16 (ImageNet-pretrained)"),
    ("Swin-UNet", "Pure-Transformer UNet", "Swin-Tiny (ImageNet-pretrained)"),
    ("nnU-Net (v2)", "Self-configuring CNN", "2D config, 250 epochs"),
    ("U-Mamba", "State-space (Mamba) UNet", "U-Mamba_Bot, 100 epochs"),
]

# (metric, formula/definition, direction, unit, chinese role)
_METRICS_INFO = [
    ("Dice (DSC)", "2·|P∩G| / (|P|+|G|) = 2TP / (2TP+FP+FN)", "↑ higher", "%",
     "区域重叠度(分割主指标);对类别不平衡较鲁棒,综合反映分割整体准确度。"),
    ("IoU (Jaccard)", "|P∩G| / |P∪G| = TP / (TP+FP+FN)", "↑ higher", "%",
     "交并比,同样衡量区域重叠但更严格(对错误惩罚更重),常与 Dice 并列报告。"),
    ("Sensitivity / Recall", "TP / (TP+FN)", "↑ higher", "%",
     "召回率/敏感度:真值前景中被正确分出的比例,反映“漏检/漏分割”程度(越高漏得越少)。"),
    ("Specificity", "TN / (TN+FP)  (pixel-level, one-vs-rest)", "↑ higher", "%",
     "特异度:真值背景中被正确判为背景的比例,反映对背景区域的误报控制能力。"),
    ("Precision", "TP / (TP+FP)", "↑ higher", "%",
     "精确率:预测为前景的像素中真正属于前景的比例,反映“过分割/误报”程度(越高误分越少)。"),
    ("HD95", "95th-percentile Hausdorff distance between predicted & GT boundaries", "↓ lower", "pixels",
     "边界距离:预测与真值轮廓间偏差的 95% 分位 Hausdorff 距离,衡量最大边界误差(取95%分位以抑制离群点),越小边界越贴合。"),
    ("ASSD", "Average symmetric surface distance between boundaries", "↓ lower", "pixels",
     "平均对称表面距离:两条轮廓间双向平均距离,衡量整体边界吻合度,越小边界越准。"),
]


def _intro_html():
    h = ["<div class='ov'>"]
    h.append("<h2>Project overview</h2>")
    h.append("<p>A unified benchmark of <b>8 2D medical-image segmentation methods</b> across "
             "<b>10 public datasets</b> spanning <b>7 imaging modalities</b> (endoscopy, retinal fundus, "
             "ultrasound, cardiac MRI, dermoscopy, histopathology, abdominal CT). Every method is trained "
             "and evaluated through one pipeline (bf16 AMP on A100, identical metrics). Reported as "
             "<b>mean±SD</b> over <b>3 seeds</b> for fixed-split datasets and over <b>folds</b> for "
             "cross-validation datasets (PanNuke: official 3-fold). Each (dataset,method) cell aggregates "
             "~tens–thousands of test images; the suite totals ≈20k images. Per-method efficiency "
             "(params / FLOPs / throughput) is in <code>efficiency.md</code>.</p>")

    h.append("<h2>Datasets</h2>")
    h.append("<table><tr><th>#</th><th>Dataset</th><th>Modality</th><th>Target</th><th>Classes</th>"
             "<th>Channels</th><th>Image size</th><th>Reported protocol</th><th>Train / Val / Test</th></tr>")
    for r in _DATASETS_INFO:
        h.append("<tr><td>%s</td><td class='l'>%s</td><td class='l'>%s</td><td class='l'>%s</td>"
                 "<td>%s</td><td>%s</td><td>%s</td><td>%s</td><td>%s</td></tr>" % r)
    h.append("</table>")
    h.append("<div class='cap'>¹ BUSI/KiTS19 are grayscale in content but stored as 3-channel PNG — read as "
             "grayscale. ² BUSI/KiTS19 have no canonical split, so one fixed fold (of a 5-fold partition) is "
             "reported with 3 seeds; datasets with an official split use it. Mask labels are 0…C-1 "
             "(0=background); multi-class metrics are macro-averaged over foreground classes.</div>")

    h.append("<h2>Segmentation methods</h2>")
    h.append("<table><tr><th>Method</th><th>Family</th><th>Backbone / setup</th></tr>")
    for m in _METHODS_INFO:
        h.append("<tr><td class='l'>%s</td><td class='l'>%s</td><td class='l'>%s</td></tr>" % m)
    h.append("</table>")

    h.append("<h2>Evaluation metrics</h2>")
    h.append("<table><tr><th>Metric</th><th>Definition</th><th>Direction</th><th>Unit</th>"
             "<th>作用 / 含义（中文）</th></tr>")
    for m in _METRICS_INFO:
        h.append("<tr><td class='l'>%s</td><td class='l'>%s</td><td>%s</td><td>%s</td>"
                 "<td class='l'>%s</td></tr>" % m)
    h.append("</table>")
    h.append("<p style='font-size:12.5px'>P = predicted foreground, G = ground-truth foreground; "
             "TP/FP/FN/TN are pixel counts. Overlap metrics (Dice/IoU/Sens/Spec/Prec) are exact pixel "
             "computations; boundary metrics (HD95/ASSD) are computed with MONAI on the mask edges.</p>")

    h.append("<h2>Evaluation protocol</h2>")
    h.append("<p>For each test image, every metric is computed <b>per foreground class</b> (background "
             "excluded) and macro-averaged over the classes present; classes absent in <i>both</i> "
             "prediction and ground truth are skipped so they do not dilute the score. Per-image values are "
             "averaged to one score per (dataset, method, run); we then report <b>mean±SD across runs</b> — "
             "3 random seeds for fixed-split datasets, or across folds for CV datasets. For surface metrics, "
             "an empty–empty (no foreground in either) pair scores 0, while empty-vs-non-empty is undefined "
             "(NaN) and excluded from the average. All methods share the same held-out test images per "
             "dataset, so comparisons are paired.</p>")
    h.append("<hr style='margin:18px 0;border:none;border-top:1px solid #ddd'>")
    h.append("</div>")
    return "\n".join(h)


def to_html(rows, title="SegGen baselines"):
    cell, dslist = {}, []
    for r in rows:
        ds = f"{r['dataset']}/{r['protocol']}"
        if ds not in dslist:
            dslist.append(ds)
        cell[(ds, r["arch"])] = r
    archs = [a for a in _ARCH_ORDER if any(r["arch"] == a for r in rows)] or \
            sorted({r["arch"] for r in rows})
    h = ["<!doctype html><html><head><meta charset='utf-8'><title>%s</title><style>" % title,
         "body{font-family:Arial,Helvetica,sans-serif;margin:24px;color:#222}",
         "h1{font-size:20px}h2{margin-top:28px;font-size:16px;color:#0a5}",
         "table{border-collapse:collapse;margin:6px 0 18px}",
         "th,td{border:1px solid #ccc;padding:5px 9px;text-align:center;font-size:12.5px}",
         "th{background:#f3f3f3}td.ds{text-align:left;font-weight:bold;background:#fafafa}",
         "td.l{text-align:left}.ov{font-size:12.5px;line-height:1.5;max-width:1100px}",
         "b{color:#0a6}.cap{color:#777;font-size:12px;margin:2px 0 14px}</style></head><body>"]
    nseeds = {f"{r['dataset']}/{r['protocol']}": r["n_seeds"] for r in rows}
    h.append(f"<h1>{title} — mean±SD</h1>")
    h.append("<div class='cap'>Higher is better for Dice/IoU/Sens/Spec/Prec (%); "
             "lower for HD95/ASSD. <b>Bold</b>=best (highest/lowest mean) per dataset. "
             "Variance is over seeds for single-split datasets and over folds for k-fold (CV) datasets. "
             "'seeds' column = #seeds (single-split) or #folds (CV).</div>")
    h.append(_intro_html())
    h.append("<h1 style='font-size:18px'>Results</h1>")
    for key, label, pct, hib in METRICS:
        h.append(f"<h2>{label}{' (%)' if pct else ' (lower=better)'}</h2>")
        h.append("<table><tr><th>Dataset</th>" + "".join(f"<th>{a}</th>" for a in archs)
                 + "<th>seeds</th></tr>")
        for ds in dslist:
            vals = {a: cell[(ds, a)][f"{key}_mean"] for a in archs
                    if (ds, a) in cell and cell[(ds, a)][f"{key}_mean"] == cell[(ds, a)][f"{key}_mean"]}
            best = (max if hib else min)(vals, key=vals.get) if vals else None
            bold = {best} if best is not None else set()
            tds = [f"<td class='ds'>{ds}</td>"]
            for a in archs:
                if (ds, a) in cell:
                    txt = _fmt(cell[(ds, a)], key, pct)
                    tds.append(f"<td>{'<b>'+txt+'</b>' if a in bold else txt}</td>")
                else:
                    tds.append("<td>—</td>")
            tds.append(f"<td>{nseeds.get(ds,'')}</td>")
            h.append("<tr>" + "".join(tds) + "</tr>")
        h.append("</table>")
    h.append("</body></html>")
    return "\n".join(h)


def main():
    p = argparse.ArgumentParser()
    p.add_argument("--exp_name", required=True)
    p.add_argument("--out_root", default="results")
    args = p.parse_args()

    runs = load_runs(args.out_root, args.exp_name)
    if not runs:
        print(f"no metrics.json under {args.out_root}/{args.exp_name}")
        return
    rows = summarize(runs)
    base = os.path.join(args.out_root, args.exp_name)
    open(os.path.join(base, "summary.csv"), "w").write(to_csv(rows))
    open(os.path.join(base, "summary.md"), "w").write(to_markdown(rows))
    open(os.path.join(base, "summary.tex"), "w").write(to_latex(rows))
    open(os.path.join(base, "summary.html"), "w").write(to_html(rows, title=f"SegGen baselines ({args.exp_name})"))
    print(to_markdown(rows))
    print(f"{len(runs)} runs -> {len(rows)} (dataset,arch) cells; written {base}/summary.{{csv,md,tex,html}}")


if __name__ == "__main__":
    main()