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.
169 lines
No EOL
7.8 KiB
Python
169 lines
No EOL
7.8 KiB
Python
"""Pseudo-label the unlabeled pool with the MONAI round model (torchrun; rank = GPU).
|
|
|
|
torchrun --standalone --nproc_per_node 3 scripts_monai/03_monai_pseudo_label.py \
|
|
--ckpt runs_monai/round0/best.pt --pool data/manifests/unlabeled_pool.jsonl \
|
|
--out data/pseudo_monai/round1 [--already data/pseudo_monai/added_keys.jsonl]
|
|
|
|
Each rank (GPU) takes rows[rank::world] and computes MONAI sliding-window
|
|
(gaussian blend, 50% overlap) + 4-flip-TTA tumor probabilities, then applies
|
|
the study gates (identical to Pipelines A/B, shared from scripts_nnu/nnu_common):
|
|
pos: p_tumor >= tau_pos, median cleanup, largest-CC fraction >= min_cc_frac,
|
|
volume within the labeled-tumor [p2, p98] range (data/vols.json)
|
|
neg: >= neg_frac of interior (vol > 0.02) voxels have p_bg >= tau_neg
|
|
then the per-subject frame-invariant longitudinal consistency filter over
|
|
accepted positive timepoints (head-relative centroid shift + volume ratio).
|
|
|
|
Writes <out>/rows.jsonl, <out>/<key>_label.nii.gz (tumor mask for pos, zero
|
|
mask for neg), <out>/accepted.jsonl, <out>/summary.json. Per-rank shards are
|
|
flushed to <out>/part{rank}.jsonl so interrupted runs resume.
|
|
|
|
Usage:
|
|
python scripts_monai/03_monai_pseudo_label.py <same flags, single GPU>
|
|
torchrun --standalone --nproc_per_node 3 scripts_monai/03_monai_pseudo_label.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 numpy as np
|
|
import SimpleITK as sitk
|
|
from src.common import load_jsonl, save_jsonl, read_nii_arr, write_arr, d
|
|
from monai_common import (rank_info, init_dist, barrier, destroy_dist, load_model,
|
|
predict_probs, pos_mask, neg_frac_bg, consistency_filter,
|
|
load_voxel_stats)
|
|
|
|
|
|
def main():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--ckpt", required=True)
|
|
ap.add_argument("--pool", required=True)
|
|
ap.add_argument("--out", required=True)
|
|
ap.add_argument("--already", default=None, help="jsonl of case keys already consumed by the dataset")
|
|
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)
|
|
ap.add_argument("--tau-pos", type=float, default=0.95)
|
|
ap.add_argument("--tau-neg", type=float, default=0.98)
|
|
ap.add_argument("--neg-frac", type=float, default=0.90)
|
|
ap.add_argument("--vol-qp", type=float, nargs=2, default=[2, 98])
|
|
ap.add_argument("--min-cc-frac", type=float, default=0.2)
|
|
ap.add_argument("--max-rel-dist", type=float, default=40.0,
|
|
help="max mm of head-relative tumor centroid shift between compatible timepoints")
|
|
ap.add_argument("--vol-ratio", type=float, default=10.0,
|
|
help="max tumor volume ratio between compatible timepoints")
|
|
args = ap.parse_args()
|
|
|
|
import torch
|
|
rank, world, local_rank = rank_info()
|
|
init_dist()
|
|
torch.cuda.set_device(local_rank)
|
|
device = f"cuda:{local_rank}"
|
|
|
|
out = d(args.out)
|
|
pool_rows = load_jsonl(args.pool)
|
|
already = set()
|
|
if args.already and os.path.exists(args.already):
|
|
already = {r["key"] for r in load_jsonl(args.already)}
|
|
done = set()
|
|
for k in range(world):
|
|
pf = os.path.join(out, f"part{k}.jsonl")
|
|
if os.path.exists(pf):
|
|
done |= {r["key"] for r in load_jsonl(pf)}
|
|
rows = [r for r in pool_rows if r["key"] not in already and r["key"] not in done]
|
|
shard = rows[rank::world]
|
|
|
|
vstats = load_voxel_stats()
|
|
vol_lo = vstats.get(f"p{args.vol_qp[0]:.0f}", 1.0)
|
|
vol_hi = vstats.get(f"p{args.vol_qp[1]:.0f}", 50000.0)
|
|
if rank == 0:
|
|
print(f"[monai:pseudo:rank0] pool={len(pool_rows)} already={len(already)} "
|
|
f"done={len(done)} to_predict={len(rows)} shard={len(shard)} "
|
|
f"vol_range=[{vol_lo:.0f},{vol_hi:.0f}]mm3", flush=True)
|
|
|
|
model, _ = load_model(args.ckpt, device)
|
|
part = os.path.join(out, f"part{rank}.jsonl")
|
|
buf, n_err = [], 0
|
|
for i, r in enumerate(shard, 1):
|
|
key = r["key"]
|
|
entry = {"key": key, "subject": r["subject"], "date": r.get("date"),
|
|
"source": r.get("source"), "pimg": r["pimg"], "label": None,
|
|
"role": "rej", "vol_mm3": 0, "maxp": None}
|
|
try:
|
|
pimg_itk = sitk.ReadImage(r["pimg"])
|
|
vol = read_nii_arr(r["pimg"]).astype("float32")
|
|
pt = predict_probs(model, vol, device, args.win, args.overlap,
|
|
tta=not args.no_tta, sw_batch=args.sw_batch)
|
|
if pt.shape != vol.shape:
|
|
raise ValueError(f"prob shape {pt.shape} != image shape {vol.shape} for {key}")
|
|
entry["maxp"] = round(float(pt.max()), 4)
|
|
got = pos_mask(pt, args.tau_pos, args.min_cc_frac, vol_lo, vol_hi)
|
|
if got is not None:
|
|
mask, cc_frac, vol_mm3 = got
|
|
lp = os.path.join(out, key + "_label.nii.gz")
|
|
write_arr(mask, lp, itk_img=pimg_itk)
|
|
entry.update({"role": "pos", "label": lp, "vol_mm3": vol_mm3, "cc_frac": round(cc_frac, 3)})
|
|
else:
|
|
frac = neg_frac_bg(pt, vol, args.tau_neg, args.neg_frac)
|
|
if frac is not None:
|
|
lp = os.path.join(out, key + "_label.nii.gz")
|
|
write_arr(np.zeros(vol.shape, dtype="uint8"), lp, itk_img=pimg_itk)
|
|
entry.update({"role": "neg", "label": lp, "neg_conf": round(frac, 4)})
|
|
except Exception as e: # noqa
|
|
entry["role"] = "error"
|
|
n_err += 1
|
|
print(f"[monai:pseudo:rank{rank}] {key} ERR {e!r}", flush=True)
|
|
buf.append(entry)
|
|
if i % 20 == 0:
|
|
with open(part, "a") as f:
|
|
for b in buf:
|
|
f.write(json.dumps(b) + "\n")
|
|
buf = []
|
|
print(f"[monai:pseudo:rank{rank}] {i}/{len(shard)}", flush=True)
|
|
if buf:
|
|
with open(part, "a") as f:
|
|
for b in buf:
|
|
f.write(json.dumps(b) + "\n")
|
|
barrier()
|
|
if rank != 0:
|
|
destroy_dist()
|
|
return
|
|
|
|
merged = []
|
|
for k in range(world):
|
|
pf = os.path.join(out, f"part{k}.jsonl")
|
|
if os.path.exists(pf):
|
|
merged.extend(load_jsonl(pf))
|
|
n_pos0 = sum(1 for x in merged if x["role"] == "pos")
|
|
n_rej = consistency_filter(merged, out, args.max_rel_dist, args.vol_ratio)
|
|
accepted = [x for x in merged if x["role"] in ("pos", "neg")]
|
|
save_jsonl(merged, os.path.join(out, "rows.jsonl"))
|
|
save_jsonl(accepted, os.path.join(out, "accepted.jsonl"))
|
|
posv = [x["vol_mm3"] for x in merged if x["role"] == "pos"]
|
|
summ = {
|
|
"n_pool_predicted": len(merged),
|
|
"n_pos": sum(1 for x in merged if x["role"] == "pos"),
|
|
"n_neg": sum(1 for x in merged if x["role"] == "neg"),
|
|
"n_pos_before_consistency": n_pos0,
|
|
"n_rejected_consistency": n_rej,
|
|
"n_other_rej": sum(1 for x in merged if x["role"] == "rej"),
|
|
"n_error": n_err,
|
|
"pos_vol_mm3": {"med": float(np.median(posv)) if posv else 0,
|
|
"p5": float(np.percentile(posv, 5)) if posv else 0,
|
|
"p95": float(np.percentile(posv, 95)) if posv else 0},
|
|
"tau_pos": args.tau_pos, "tau_neg": args.tau_neg, "neg_frac": args.neg_frac,
|
|
"vol_range": [vol_lo, vol_hi], "max_rel_dist_mm": args.max_rel_dist,
|
|
"vol_ratio": args.vol_ratio, "win": args.win, "overlap": args.overlap,
|
|
"tta": not args.no_tta, "ckpt": args.ckpt,
|
|
}
|
|
with open(os.path.join(out, "summary.json"), "w") as f:
|
|
json.dump(summ, f, indent=1)
|
|
print(f"[monai:pseudo] round {os.path.basename(out)}: {json.dumps(summ)}", flush=True)
|
|
destroy_dist()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |