File size: 5,751 Bytes
b8fae22 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | """Evaluate a trained framework checkpoint on a chosen dataset/split, with the full
metric set INCLUDING boundary metrics. Two uses, one script (no retraining):
* C7 (boundary fidelity, in-domain): point --eval_dataset at the same dataset to
recompute boundary-Dice / NSD from a saved best.pth.
* C4 (cross-center): set --eval_dataset/--eval_protocol to a DIFFERENT but
label-compatible dataset (e.g. train busi? no — train cvc_clinicdb, eval
kvasir_seg; both binary polyp). num_classes/in_channels must match.
Run from project root (…/NPJ), env seggen:
CUDA_VISIBLE_DEVICES=5 python -m framework.eval_ckpt \
--ckpt results/baselines/cvc_clinicdb_official/unet/seed0/best.pth \
--arch unet --encoder resnet50 \
--data_root $DR --dataset cvc_clinicdb --protocol official \
--eval_dataset kvasir_seg --eval_protocol official \
--out_json results/crosscenter/cvc2kvasir_unet_seed0.json
"""
from __future__ import annotations
import argparse
import json
import os
import sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
import numpy as np
import torch
from torch.utils.data import DataLoader
from framework.config import Config # noqa: F401 (kept for parity / future YAML use)
from framework.models.registry import build_model, required_img_size
from framework.data.unified_dataset import UnifiedSegDataset
from framework.data.transforms import build_transform
from framework.metrics.metrics import per_image_metrics
from framework.metrics.boundary import boundary_metrics
def get_args():
p = argparse.ArgumentParser("Evaluate a checkpoint (in-domain or cross-center) + boundary metrics")
p.add_argument("--ckpt", required=True)
p.add_argument("--arch", default="unet")
p.add_argument("--encoder", default="resnet50")
p.add_argument("--data_root", required=True)
p.add_argument("--dataset", required=True, help="dataset the ckpt was TRAINED on (sets in_ch/num_classes)")
p.add_argument("--protocol", required=True)
p.add_argument("--eval_dataset", default="", help="dataset to evaluate ON (default = --dataset)")
p.add_argument("--eval_protocol", default="", help="default = --protocol")
p.add_argument("--split", default="test")
p.add_argument("--img_size", type=int, default=256)
p.add_argument("--normalize", default="auto")
p.add_argument("--tol", type=float, default=2.0)
p.add_argument("--batch_size", type=int, default=16)
p.add_argument("--num_workers", type=int, default=6)
p.add_argument("--out_json", default="")
return p.parse_args()
def main():
a = get_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
img_size = required_img_size(a.arch) or a.img_size
eval_ds_name = a.eval_dataset or a.dataset
eval_proto = a.eval_protocol or a.protocol
# in_ch / num_classes are fixed by the TRAIN dataset (how the model was built)
train_probe = UnifiedSegDataset(a.data_root, a.dataset, a.protocol, "test", transform=None)
in_ch, n_cls = train_probe.in_channels, train_probe.num_classes
model = build_model(a.arch, in_channels=in_ch, num_classes=n_cls, img_size=img_size,
encoder=a.encoder, encoder_weights="none", pretrained_ckpt="")
ckpt = torch.load(a.ckpt, map_location="cpu", weights_only=False)
model.load_state_dict(ckpt["model"])
model.to(device).eval()
tf = build_transform(img_size, in_ch, train=False, aug="none", normalize=a.normalize)
ds = UnifiedSegDataset(a.data_root, eval_ds_name, eval_proto, a.split,
transform=tf, in_channels=in_ch, num_classes=n_cls)
if ds.num_classes != n_cls:
raise ValueError(f"label mismatch: train num_classes={n_cls} vs eval={ds.num_classes} "
f"({a.dataset}->{eval_ds_name}) — not label-compatible for cross-center.")
cross = (eval_ds_name != a.dataset) or (eval_proto != a.protocol)
print(f"[eval] ckpt={os.path.basename(a.ckpt)} train={a.dataset}/{a.protocol} "
f"eval={eval_ds_name}/{eval_proto}/{a.split} cross_center={cross} "
f"in_ch={in_ch} num_classes={n_cls} n={len(ds)}", flush=True)
loader = DataLoader(ds, batch_size=a.batch_size, shuffle=False, num_workers=a.num_workers)
recs = []
for batch in loader:
img = batch["image"].to(device, non_blocking=True)
msk = batch["mask"].numpy()
with torch.no_grad():
pred = model(img).argmax(1).cpu().numpy()
for i in range(pred.shape[0]):
m = per_image_metrics(pred[i], msk[i], n_cls,
include_background=False, compute_hd95=True)
b = boundary_metrics(pred[i], msk[i], n_cls, tol=a.tol)
m.update(b)
recs.append(m)
keys = [k for k, val in recs[0].items() if isinstance(val, (int, float))] # skip per_class dict
summary = {}
for k in keys:
v = np.array([r[k] for r in recs], dtype=np.float64)
v = v[~np.isnan(v)]
summary[k] = {"mean": round(float(v.mean()), 4) if v.size else None,
"std": round(float(v.std()), 4) if v.size else None, "n": int(v.size)}
out = {"ckpt": a.ckpt, "train": f"{a.dataset}/{a.protocol}",
"eval": f"{eval_ds_name}/{eval_proto}/{a.split}", "cross_center": cross,
"num_images": len(ds), "metrics": summary}
print(json.dumps(out, indent=2), flush=True)
if a.out_json:
os.makedirs(os.path.dirname(os.path.abspath(a.out_json)) or ".", exist_ok=True)
with open(a.out_json, "w") as f:
json.dump(out, f, indent=2)
print(f"[eval] wrote {a.out_json}", flush=True)
if __name__ == "__main__":
main()
|