Furen Xiao
77adc2b3af
Introduce the third segmentation pipeline using MONAI (1.6.x) to allow direct comparison with Pipelines A and B. This includes the implementation of the iterative pseudo-labeling workflow, training scripts, and inference protocols. - Add `scripts_monai/` directory containing the MONAI pipeline scripts. - Update documentation in `README.md` and `AGENTS.md` to include MONAI package requirements and pipeline details. - Configure `.gitignore` to exclude MONAI-specific run directories. - Update data directory descriptions to include MONAI pseudo-labels.
103 lines
No EOL
4.1 KiB
Python
103 lines
No EOL
4.1 KiB
Python
"""Held-out test evaluation of a MONAI round model (torchrun; rank = GPU).
|
|
|
|
torchrun --standalone --nproc_per_node 3 scripts_monai/04_monai_eval_test.py \
|
|
--rows data/manifests/split_test.jsonl --ckpt runs_monai/round0/best.pt \
|
|
--out results/round0_test_monai.json [--no-tta]
|
|
|
|
Per case: MONAI sliding-window (gaussian blend, 50% overlap) + 4-flip-TTA tumor
|
|
probability map; Dice at threshold 0.5 vs the held-out label — the same
|
|
convention as Pipeline A's 08_eval and Pipeline B's primary metric.
|
|
|
|
Per-rank results are appended to <out>_part{rank}.jsonl for resume; rank 0
|
|
merges them into <out> (json) + <out>_per_row.jsonl.
|
|
|
|
Usage:
|
|
python scripts_monai/04_monai_eval_test.py <same flags, single GPU>
|
|
torchrun --standalone --nproc_per_node 3 scripts_monai/04_monai_eval_test.py ...
|
|
"""
|
|
import argparse
|
|
import os
|
|
import sys
|
|
import json
|
|
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
import torch
|
|
from src.common import load_jsonl, save_jsonl, read_nii_arr
|
|
from monai_common import (rank_info, init_dist, barrier, destroy_dist, load_model,
|
|
predict_probs, prob_dice)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--rows", required=True)
|
|
ap.add_argument("--ckpt", required=True)
|
|
ap.add_argument("--out", required=True)
|
|
ap.add_argument("--no-tta", action="store_true")
|
|
ap.add_argument("--win", type=int, default=96)
|
|
ap.add_argument("--overlap", type=float, default=0.5)
|
|
ap.add_argument("--sw-batch", type=int, default=8)
|
|
args = ap.parse_args()
|
|
|
|
rank, world, local_rank = rank_info()
|
|
init_dist()
|
|
torch.cuda.set_device(local_rank)
|
|
device = f"cuda:{local_rank}"
|
|
|
|
rows = load_jsonl(args.rows)
|
|
outp = os.path.join(os.path.dirname(os.path.abspath(args.out)),
|
|
os.path.basename(args.out) + f"_part{rank}.jsonl")
|
|
done = set()
|
|
if os.path.exists(outp):
|
|
done = {r["key"] for r in load_jsonl(outp)}
|
|
shard = [r for r in rows if r["key"] not in done][rank::world]
|
|
if rank == 0:
|
|
print(f"[monai:eval:rank0] n_test={len(rows)} done={len(done)} "
|
|
f"shard(world={world})={len(shard)}", flush=True)
|
|
|
|
model, _ = load_model(args.ckpt, device)
|
|
with open(outp, "a") as f:
|
|
for i, r in enumerate(shard, 1):
|
|
key = r["key"]
|
|
entry = {"key": key, "dice": None}
|
|
try:
|
|
vol = read_nii_arr(r["pimg"]).astype("float32")
|
|
lab = read_nii_arr(r.get("plabel") or r.get("label")).astype("uint8")
|
|
p = predict_probs(model, vol, device, args.win, args.overlap,
|
|
tta=not args.no_tta, sw_batch=args.sw_batch)
|
|
if p.shape == lab.shape:
|
|
entry["dice"] = round(prob_dice(p, lab), 4)
|
|
else:
|
|
raise ValueError(f"prob {p.shape} vs label {lab.shape}")
|
|
except Exception as e: # noqa
|
|
print(f"[monai:eval:rank{rank}] {key} ERR {e!r}", flush=True)
|
|
f.write(json.dumps(entry) + "\n")
|
|
f.flush()
|
|
if i % 20 == 0:
|
|
print(f"[monai:eval:rank{rank}] {i}/{len(shard)}", flush=True)
|
|
barrier()
|
|
if rank != 0:
|
|
destroy_dist()
|
|
return
|
|
|
|
per = []
|
|
for k in range(world):
|
|
pf = os.path.join(os.path.dirname(os.path.abspath(args.out)),
|
|
os.path.basename(args.out) + f"_part{k}.jsonl")
|
|
if os.path.exists(pf):
|
|
per.extend(load_jsonl(pf))
|
|
dice_ok = [e["dice"] for e in per if e["dice"] is not None]
|
|
res = {"ckpt": args.ckpt, "n": len(per), "n_ok": len(dice_ok),
|
|
"dice": float(sum(dice_ok) / len(dice_ok)) if dice_ok else 0.0,
|
|
"tta": not args.no_tta}
|
|
os.makedirs(os.path.dirname(os.path.abspath(args.out)), exist_ok=True)
|
|
with open(args.out, "w") as f:
|
|
json.dump(res, f, indent=1)
|
|
save_jsonl(per, args.out.replace(".json", "_per_row.jsonl"))
|
|
print(json.dumps({k: res[k] for k in ("n", "n_ok", "dice")}, indent=1), flush=True)
|
|
print("saved", args.out, flush=True)
|
|
destroy_dist()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |