| """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 |
| 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 |
|
|
| |
| 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))] |
| 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() |
|
|
