File size: 2,886 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 | """Test-time evaluation: load best checkpoint, run on the test split, compute
Dice / IoU / HD95 (per image -> mean +- SD), optionally save overlay visualizations
and a metrics.json. Runs single-process (rank 0) for deterministic reporting.
"""
from __future__ import annotations
import os
import json
import numpy as np
import torch
from ..data.loaders import build_dataset, build_loader
from ..metrics.metrics import per_image_metrics, aggregate
from ..visualize.overlay import save_overlay
@torch.no_grad()
def evaluate(cfg, model, device, ckpt_path: str = "") -> dict:
ds = build_dataset(cfg, "test")
num_classes = ds.num_classes
loader = build_loader(cfg, "test", ds) # single-process: no DistributedSampler
ckpt_path = ckpt_path or os.path.join(cfg.out_dir(), "best.pth")
if os.path.isfile(ckpt_path):
ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
state = ckpt.get("model", ckpt)
model.load_state_dict(state)
print(f"[eval] loaded {ckpt_path}")
else:
print(f"[eval][warn] checkpoint not found: {ckpt_path} (evaluating current weights)")
model = model.to(device).eval()
use_amp = cfg.amp in ("bf16", "fp16")
amp_dtype = torch.bfloat16 if cfg.amp == "bf16" else torch.float16
records = []
vis_dir = os.path.join(cfg.out_dir(), "vis")
if cfg.visualize:
os.makedirs(vis_dir, exist_ok=True)
saved = 0
for batch in loader:
img = batch["image"].to(device, non_blocking=True)
msk = batch["mask"].numpy()
names = batch["name"]
with torch.autocast("cuda", dtype=amp_dtype, enabled=use_amp):
logits = model(img)
pred = logits.argmax(1).cpu().numpy()
for i in range(pred.shape[0]):
records.append(per_image_metrics(
pred[i], msk[i], num_classes,
include_background=cfg.include_background,
compute_hd95=cfg.compute_hd95))
if cfg.visualize and saved < cfg.vis_max:
save_overlay(img[i].cpu(), msk[i], pred[i], num_classes,
os.path.join(vis_dir, f"{names[i]}.png"))
saved += 1
agg = aggregate(records)
out = {
"dataset": cfg.dataset, "protocol": cfg.protocol, "arch": cfg.arch,
"seed": cfg.seed, "num_classes": num_classes,
"metrics": agg,
"per_image": records,
}
out_path = os.path.join(cfg.out_dir(), "metrics.json")
with open(out_path, "w") as f:
json.dump(out, f, indent=2)
print(f"[eval] dice={agg['dice_mean']:.4f}+-{agg['dice_std']:.4f} "
f"iou={agg['iou_mean']:.4f} hd95={agg['hd95_mean']:.3f} assd={agg['assd_mean']:.3f} "
f"sens={agg['sensitivity_mean']:.4f} spec={agg['specificity_mean']:.4f} "
f"prec={agg['precision_mean']:.4f} -> {out_path}")
return out
|