feat(registration): implement best-of-ANTs registration for CT to T1c alignment and enhance benchmarking scripts
This commit is contained in:
parent
12143e5bb8
commit
41c88b2a67
3 changed files with 270 additions and 156 deletions
208
BENCHMARKING.md
208
BENCHMARKING.md
|
|
@ -1,152 +1,96 @@
|
|||
# Benchmarking Pipelines A / B / C
|
||||
# Pipeline benchmark (A / B / C)
|
||||
|
||||
`scripts/benchmark_pipelines.py` compares the three segmentation pipelines on
|
||||
**one GPU** with a unified protocol:
|
||||
Benchmark of the three segmentation pipelines — A in-house Unet3D, B nnU-Net v2
|
||||
(NTUHLPLTrainer), C MONAI UNet — driven by `scripts/benchmark_pipelines.py`
|
||||
(DICE, training time, inference time, params, VRAM). Protocol details live in
|
||||
the script docstring; this file records what has been run and what the machine
|
||||
currently supports.
|
||||
|
||||
| pipeline | backbone | data protocol |
|
||||
|---|---|---|
|
||||
| A | in-house `Unet3D(base=16, depth=4)` | processed 1 mm head-cropped percentile-normalized volumes (`pimg`/`plabel`) |
|
||||
| B | nnU-Net v2 `3d_fullres` | native source volumes (`img`/`label`), per Pipeline B's design |
|
||||
| C | MONAI `UNet(16→128)` | processed 1 mm volumes (`pimg`/`plabel`) |
|
||||
## Data chain status (2026-10-04)
|
||||
|
||||
Each pipeline is benchmarked with **its own standard settings** (batch, patch,
|
||||
augmentation, TTA, optimizer), so the results reflect the pipelines as used in
|
||||
the study.
|
||||
The M6 labeled set was fully rebuilt on 2026-10-04 (see
|
||||
`scripts/02_build_m6_dataset.py`): T1c fixed, best-of-ANTs
|
||||
(`ants_rigid` / `ants_affine` / `ants_syn`, winner by Mattes MI), TV warped
|
||||
onto the T1c grid with SMD+linear. After the repair pass:
|
||||
|
||||
## Metrics
|
||||
- `data/manifests/m6_labeled.jsonl`: **1312** volumes (was 1094 pre-repair),
|
||||
`m6_unlabeled.jsonl`: **229**
|
||||
- all 1541 m6 processed volumes (`data/proc/<key>{,_label}.nii.gz`, 1mm
|
||||
cropped/normalized) regenerated, 0 errors
|
||||
- patient-level splits rebuilt (`scripts/05_build_splits.py`, seed 0):
|
||||
**4200** labeled rows (ntuh + m6) → train 3570 / val 211 / test 419
|
||||
volumes, 3893 subjects; unlabeled pool 13280
|
||||
- tumor volume stats: median 3116 mm3, p95 24.8k, max 3.2M, zero_frac 0.12%
|
||||
|
||||
- **DICE** — tumor probability thresholded at 0.5 vs ground truth on the
|
||||
held-out patient-level test split (`split_test.jsonl`), the primary-metric
|
||||
convention of all three pipelines. B additionally reports the **hard-seg**
|
||||
(argmax) Dice it exports natively.
|
||||
- **Training** — wall-clock total and s/epoch for `--train-epochs` on the
|
||||
first `--max-rows` train rows (single GPU). For B, the one-time raw-dataset
|
||||
build and `plan_and_preprocess` wall time are reported separately
|
||||
(`raw_build_s`, `plan_preprocess_s`).
|
||||
- **Inference** — mean **s/volume** on `--max-eval-rows` held-out test volumes
|
||||
(first `--eval-warmup` cases untimed), each pipeline's standard
|
||||
sliding-window + TTA protocol (`--no-tta` disables TTA everywhere).
|
||||
- **Model size** — parameter count, checkpoint weight file size (MB), and peak
|
||||
**VRAM** (MB) for the training run and for the inference run.
|
||||
(A/C training VRAM via `torch.cuda.max_memory_allocated`; B via nvidia-smi
|
||||
polling around the `nnUNetv2_train` subprocess.)
|
||||
## Smoke run (2026-10-04)
|
||||
|
||||
## Prerequisites
|
||||
|
||||
- The `longitudinal-313` conda env and the GPU(s) it sees:
|
||||
```bash
|
||||
source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal-313
|
||||
```
|
||||
- Data for Pipelines A and C (and split membership for all three):
|
||||
`data/proc/*` + `data/manifests/split_{train,val,test}.jsonl`, i.e.
|
||||
`scripts/preprocess.py` and `scripts/05_build_splits.py` have been run.
|
||||
- Data for Pipeline B (train/eval mode): the labeled source manifests
|
||||
`data/manifests/{ntuh,m6_labeled}.jsonl` with reachable native NIfTIs
|
||||
(the benchmark joins native `img`/`label` paths by case key).
|
||||
- `--skip-train` mode additionally needs the study checkpoints:
|
||||
`runs/round0/best.pt` (A), `runs_monai/round0/best.pt` (C), and the study
|
||||
nnU-Net results tree `nnu/results/Dataset210_NTUH_T1C_PL/...` (B).
|
||||
|
||||
## Quick start
|
||||
|
||||
```bash
|
||||
cd /mnt/b4/xfr/git26/longitudinal
|
||||
conda activate longitudinal-313
|
||||
|
||||
# small/fast check (what the validation run used): ~minutes per pipeline
|
||||
python scripts/benchmark_pipelines.py \
|
||||
--train-epochs 1 --max-rows 24 --nnu-cases 32 --max-eval-rows 4
|
||||
|
||||
# default full benchmark: train 100 epochs on 80% of train cases (~1986 rows),
|
||||
# eval on 10% of test volumes (~30)
|
||||
python scripts/benchmark_pipelines.py
|
||||
```
|
||||
|
||||
Outputs (repo root):
|
||||
Small/fast end-to-end check that all three pipelines run on the rebuilt data
|
||||
(1 epoch, 24 train rows, 3 eval rows, TTA on, GPU 0):
|
||||
|
||||
```
|
||||
results/benchmark_pipelines.json # settings + all metrics, per pipeline
|
||||
results/benchmark_pipelines_per_row.jsonl
|
||||
results/benchmark_pipelines_bench.png # 6-panel bar chart
|
||||
python scripts/benchmark_pipelines.py --gpu 0 --max-rows 24 --nnu-cases 24 \
|
||||
--max-eval-rows 3 --train-epochs 1 --out results/benchmark_pipelines_smoke.json
|
||||
```
|
||||
|
||||
plus the artifacts it created (safe to delete): `runs_bench/{A,C}/best.pt`,
|
||||
`nnu/{raw,preprocessed,results}/Dataset221_T1C_BENCH`,
|
||||
`logs/nnu/bench_{plan,train}.log`.
|
||||
| pipeline | wall | s/epoch | tr VRAM | s/vol (TTA) | in VRAM | params | weights | dice (1 epoch) |
|
||||
|---|---|---|---|---|---|---|---|---|
|
||||
| A Unet3D(16, depth 4) | 17 s | 2.5 (8 steps) | 2.0 GB | 4.12 | 0.5 GB | 1.46 M | 5.9 MB | 0.0000 |
|
||||
| B nnU-Net 3d_fullres | 285 s (incl. 83 s plan&preprocess, 4 s raw build) | 73.3 | 9.8 GB | 34.73 | 1.8 GB | 30.8 M | 246 MB | 0.0000 |
|
||||
| C MONAI UNet(16→128) | 5 s | 2.5 (8 steps) | 1.7 GB | 0.33 | 0.5 GB | 1.19 M | 4.8 MB | 0.0005 |
|
||||
|
||||
## All options
|
||||
DICE ≈ 0 is expected (models trained for a single epoch on 24 volumes); the
|
||||
point of the run was validating the full chain — new m6 native/processed
|
||||
manifests, split membership, B's `ntuh + m6_labeled` native join, nnU-Net
|
||||
isolated dataset 221 build, inference + dice scoring. All three passed.
|
||||
Eval rows for the smoke run were the first three `split_test` volumes (all
|
||||
ntuh). Artifacts: `results/benchmark_pipelines_smoke.json`,
|
||||
`_per_row.jsonl`, `_bench.png`.
|
||||
|
||||
| flag | default | meaning |
|
||||
|---|---|---|
|
||||
| `--pipelines` | `A B C` | space-separated subset, e.g. `"A C"` |
|
||||
| `--gpu` | `0` | physical GPU index; all pipelines run serially on this one GPU |
|
||||
| `--train-epochs` | `100` | training budget for A and C (and B via `NNU_PL_EPOCHS`) |
|
||||
| `--max-rows` | `80%` of `split_train` | train rows for A/C (first N of `split_train.jsonl` that have labels; default 80% of the split, ~1986 rows) |
|
||||
| `--nnu-cases` | `80%` of `split_train` | train cases for B (first N labeled native cases of the same split; default same as `--max-rows`) |
|
||||
| `--max-eval-rows` | `10%` of `split_test` | held-out test volumes for inference + DICE (default 10% of the split, ~30 rows) |
|
||||
| `--eval-warmup` | `1` | untimed warmup cases before inference timing |
|
||||
| `--workers` | `4` | data-loader workers (A/C) |
|
||||
| `--batch` | `3` | per-GPU batch (A/C) |
|
||||
| `--patch` | `96` | patch/sliding-window size (A/C) |
|
||||
| `--no-tta` | off | disable each pipeline's standard test-time augmentation |
|
||||
| `--skip-train` | off | evaluate existing study checkpoints instead of training |
|
||||
| `--nnu-dsid` | `221` train / `210` skip-train | nnU-Net dataset id: in training mode an isolated `Dataset{id}_T1C_BENCH` (the study's 210 is never touched); in `--skip-train` mode the study's 210 by default, or a previously benchmarked id to re-evaluate it |
|
||||
| `--nnu-skip-prep` | off | B (train mode): skip `plan_and_preprocess` and reuse the existing raw + preprocessed tree for `--nnu-dsid` (e.g. retime one more training epoch on an already prepared dataset) |
|
||||
| `--out` | `results/benchmark_pipelines.json` | output JSON path (the `_per_row.jsonl` / `_bench.png` are derived from it) |
|
||||
### Environment fixes made during the check
|
||||
|
||||
## Examples
|
||||
- `monai` was missing from `longitudinal-313` (pipeline C failed with
|
||||
ModuleNotFoundError) — installed `pip install monai==1.6.1 --no-deps`
|
||||
(torch 2.14+cu126 untouched); `import monai` verified.
|
||||
|
||||
```bash
|
||||
# longer training budget, bigger eval set (quality-focused run)
|
||||
python scripts/benchmark_pipelines.py --train-epochs 5 --max-rows 240 --max-eval-rows 20
|
||||
## Resuming training (`--resume`)
|
||||
|
||||
# timing only, no TTA
|
||||
python scripts/benchmark_pipelines.py --no-tta --max-eval-rows 20
|
||||
All three pipelines can continue from their latest benchmark checkpoint:
|
||||
|
||||
# re-evaluate the study's trained checkpoints (no training)
|
||||
python scripts/benchmark_pipelines.py --skip-train --max-eval-rows 24
|
||||
|
||||
# benchmark B only; keep the trained Dataset221 model for a second inference-only pass
|
||||
python scripts/benchmark_pipelines.py --pipelines B
|
||||
python scripts/benchmark_pipelines.py --pipelines B --skip-train --nnu-dsid 221
|
||||
|
||||
# use GPU 1 and write to a custom location
|
||||
python scripts/benchmark_pipelines.py --gpu 1 --out results/bench_gpu1.json
|
||||
```
|
||||
python scripts/benchmark_pipelines.py --pipelines "A B C" --resume \
|
||||
--train-epochs 200 --nnu-skip-prep
|
||||
```
|
||||
|
||||
## Interpreting the results
|
||||
- `--train-epochs` is the **total budget** (not epochs added). Each pipeline
|
||||
trains from where its checkpoint left off up to that total.
|
||||
- **A / C**: load `runs_bench/{A,C}/best.pt` (model + optimizer + scheduler
|
||||
state — checkpoints written before 2026-10-04 carry model + epoch only and
|
||||
fall back to fresh optimizer/scheduler) and run the remaining epochs; the
|
||||
run is a no-op train (eval only) when the checkpoint already reached the
|
||||
total.
|
||||
- **B**: appends nnU-Net's `--c` flag (continue `fold_0` from
|
||||
`checkpoint_latest.pth`, native optimizer + LR-schedule restore) and
|
||||
`NNUNet_pl_EPOCHS = --train-epochs`. Pair with `--nnu-skip-prep` so the
|
||||
preprocessed tree isn't `--clean`ed; the isolated Dataset221 is rebuilt
|
||||
otherwise and a finished run's `checkpoint_latest.pth` is gone, so a
|
||||
completed run + a bigger budget retrains B from scratch (interrupted runs
|
||||
keep `checkpoint_latest.pth` and resume correctly).
|
||||
- Verified 2026-10-04: A/C resumed epoch 1 → 3 (model restored, fresh
|
||||
opt/sched for the old-format ckpt), re-run at 3/3 trained 0 epochs, B
|
||||
continued from epoch 1 with LR restored mid-decay (0.00536 @ start of
|
||||
epoch 2).
|
||||
|
||||
- **"1 epoch" is pipeline-native.** A/C: one pass over the N train volumes
|
||||
(batch 3, 96³ random patches, AdamW 3e-4 + warmup/cosine, bf16).
|
||||
B: one pass over the N train cases at nnU-Net's planned batch and patch size
|
||||
(typically 2 and ~256³, AdamW 1e-2 PolyLR, fp16). Epoch times are therefore
|
||||
comparable as *wall-clock cost per pass over the training set in each
|
||||
pipeline's standard configuration*, not per identical number of patches.
|
||||
- **DICE reflects the training budget.** Short benchmarks measure the
|
||||
machinery, not the final quality of the pipelines. Compare DICE across
|
||||
pipelines only at the same `--train-epochs`/`--max-rows` budget, or use
|
||||
`--skip-train` on fully trained study checkpoints.
|
||||
- **Inference times are single-GPU, standard-protocol numbers**
|
||||
(A: 96³ step 48 + 4-flip TTA; B: nnU-Net gaussian sliding window +
|
||||
mirroring, native volumes; C: MONAI gaussian 50% + 4-flip TTA).
|
||||
B is the slowest per volume on this data set because it runs on uncropped
|
||||
native volumes with full mirroring TTA.
|
||||
- **B's one-time cost.** `raw_build_s` + `plan_preprocess_s` are paid once per
|
||||
(dataset, round) — in the iterative study every round repays the
|
||||
plan/preprocess on the growing dataset; the benchmark reports it for
|
||||
context only.
|
||||
- **VRAM notes.** Training VRAM includes optimizer states (AdamW) and the
|
||||
largest batch; B's number is the nvidia-smi peak of the whole
|
||||
`nnUNetv2_train` process tree on that GPU. Inference VRAM is per
|
||||
pipeline as measured in-process (B: torch peak over the timed cases).
|
||||
## Readiness notes for a full run
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
| symptom | fix |
|
||||
|---|---|
|
||||
| `split_train.jsonl has no labeled processed rows` | run `scripts/preprocess.py` over the manifests and `scripts/05_build_splits.py` first |
|
||||
| `B: key ... not in ntuh/m6_labeled manifests` | build the source manifests (`scripts/01_build_ntuh_manifest.py`, `scripts/02_build_m6_dataset.py`) and confirm the native NIfTIs are reachable |
|
||||
| `--skip-train: .../best.pt not found` | the study checkpoint for that pipeline doesn't exist yet — train a round first, or drop `--skip-train` |
|
||||
| B `--skip-train` fails on the model folder | the study's `nnu/results/Dataset210_NTUH_T1C_PL/NTUHLPLTrainer__nnUNetPlans__3d_fullres/fold_0/` must exist (produced by `scripts_nnu/03_nnu_train.py`) |
|
||||
| OOM on B | B trains at the planned 256³-class patches and needs ~10–20 GB; free other GPU users or run with a smaller `--nnu-cases` set (patch size derives from the median case size) |
|
||||
| inflated/inconsistent timings | do not run two benchmarks on the same `--gpu` concurrently; note that data-loader (NFS) latency is included in training and inference wall times |
|
||||
- Budget knobs: `--train-epochs` (default 100), `--max-rows` (default 80% of
|
||||
split_train ≈ 2856 for A/C and for B's `--nnu-cases`), `--max-eval-rows`
|
||||
(default 10% of split_test ≈ 42), `--no-tta`.
|
||||
- `--skip-train` is only usable for **A** (`runs/round0/best.pt`, trained
|
||||
2026-09-25 on the *previous* split — a few patients moved between
|
||||
train/test in the rebuild, so its test dice carries a small leakage
|
||||
caveat). No study checkpoints exist for B (Dataset210) or C
|
||||
(`runs_monai/round0/best.pt`); they must be trained.
|
||||
- GPUs: 3× 80 GB available; the benchmark uses one (`--gpu`).
|
||||
- B's one-time cost scales with `--nnu-cases` (plan_and_preprocess ≈ 3.5 s/case
|
||||
at 24 cases with 8 preprocess procs, plus training time at the planned
|
||||
batch/patch).
|
||||
|
|
@ -1,12 +1,20 @@
|
|||
"""Build M6-2025 manifests: T1c volumes, labeled when GTV exists.
|
||||
|
||||
GTV lives in CT space; T1c is registered to CT (rigid, Mattes MI on the full
|
||||
volumes -- no mask, no crop / pad; multi-start initial tilts because
|
||||
inter-session head tilt reaches ~30 deg and tilted MI local minima can
|
||||
outscore the true alignment).
|
||||
The inverse transform then warps CT and GTV (nearest) into T1c native space.
|
||||
GTV lives in CT space; the T1c is the fixed image and the CT is registered
|
||||
to it with the best of three ANTsPy pipelines (src/registration_pipelines.py,
|
||||
full volumes, no mask / crop / pad, each toolkit's native COM init):
|
||||
ants_rigid : Rigid[0.25] (Mattes, 500x500x200, 4x2x1)
|
||||
ants_affine : Affine[0.25] (Mattes, 2100x1200x1200x10, 6x4x2x1)
|
||||
ants_syn : SyN (Affine[0.25] + SyN[0.2,3,0])
|
||||
The winner is the candidate with the highest Mattes-style MI between the
|
||||
warped CT and the T1c (full volumes). All three return the
|
||||
transform in the unified contract direction (T1c -> CT), so resampling the
|
||||
CT / GTV onto the T1c grid with it places them in T1c native space -- the
|
||||
TV is warped with SMD + linear (signed distance to the TV boundary, linearly
|
||||
resampled, thresholded at the mean CT spacing halved; the validated recipe
|
||||
from scripts/qa_label_interpolation.py).
|
||||
Acceptance requires the warped CT skull ring to be concentric with the MR head
|
||||
silhouette (crop-independent check).
|
||||
silhouette.
|
||||
Outputs:
|
||||
data/manifests/m6_labeled.jsonl {key, subject, date, img, label(warped GTV), source:'m6'}
|
||||
data/manifests/m6_unlabeled.jsonl {key, subject, date, img, label:None, source:'m6'}
|
||||
|
|
@ -133,8 +141,8 @@ def _skull_ring_gap(warped_ct, ha):
|
|||
for a in np.linspace(0, 2 * np.pi, 72, endpoint=False):
|
||||
dy, dx = np.sin(a), np.cos(a)
|
||||
rr = np.arange(0.0, 160.0, 0.5)
|
||||
yy = np.clip((cy0 + rr * dy).astype(int), 0, ha.shape[0] - 1)
|
||||
xx = np.clip((cx0 + rr * dx).astype(int), 0, ha.shape[1] - 1)
|
||||
yy = np.clip((cy0 + rr * dy).astype(int), 0, sl.shape[0] - 1)
|
||||
xx = np.clip((cx0 + rr * dx).astype(int), 0, sl.shape[1] - 1)
|
||||
rh = rr[np.where(sl[yy, xx])[0]]
|
||||
rs = rr[np.where(sk[yy, xx])[0]]
|
||||
if len(rh) and len(rs):
|
||||
|
|
@ -321,6 +329,112 @@ def register_ct_to_t1c(ct, t1c, threads=12):
|
|||
return tr, info, tissue_w
|
||||
|
||||
|
||||
_QA = None
|
||||
|
||||
|
||||
def _load_qa():
|
||||
"""Import scripts/qa_label_interpolation.py (the validated SMD + linear
|
||||
label-warp recipe: smd_threshold / make_tv_bin / warp_tv). Cached per
|
||||
process (each worker process keeps its own copy)."""
|
||||
global _QA
|
||||
if _QA is None:
|
||||
import importlib.util
|
||||
p = os.path.join(os.path.dirname(os.path.abspath(__file__)),
|
||||
"qa_label_interpolation.py")
|
||||
spec = importlib.util.spec_from_file_location("qa_label_interp", p)
|
||||
m = importlib.util.module_from_spec(spec)
|
||||
spec.loader.exec_module(m)
|
||||
_QA = m
|
||||
return _QA
|
||||
|
||||
|
||||
def register_ants_best(ct, t1c):
|
||||
"""Best-of-ANTs CT -> T1c registration with the T1c as the fixed image.
|
||||
|
||||
Runs the three ANTsPy pipelines from src/registration_pipelines.py
|
||||
(fixed = T1c, moving = CT, full volumes, no mask / crop / pad, each
|
||||
toolkit's native COM init):
|
||||
ants_rigid : Rigid[0.25] (Mattes, 500x500x200, 4x2x1)
|
||||
ants_affine : Affine[0.25] (Mattes, 2100x1200x1200x10, 6x4x2x1)
|
||||
ants_syn : SyN (Affine[0.25] + SyN[0.2,3,0])
|
||||
Each returns a transform in the unified contract direction (T1c -> CT),
|
||||
so sitk.Resample(ct, t1c, T, ...) / sitk.Resample(tv, t1c, T, ...) place
|
||||
the CT / TV concentric with the MR head -- no inversion (the T1c was the
|
||||
fixed image). The winner is the candidate with the highest Mattes-style
|
||||
MI between the warped CT and the T1c (_mi_value, full volumes). The
|
||||
per-candidate skull-ring gap is still recorded and gates acceptance.
|
||||
|
||||
Returns (transform, info, tissue_w) like register_ct_to_t1c.
|
||||
"""
|
||||
import time
|
||||
from scipy import ndimage
|
||||
from src import registration_pipelines as RP
|
||||
from src.common import head_mask_from_image
|
||||
t_start = time.time()
|
||||
ct = sitk.Cast(ct, sitk.sitkFloat32)
|
||||
t1c = sitk.Cast(t1c, sitk.sitkFloat32)
|
||||
calls = {
|
||||
"ants_rigid": lambda: RP.register_antspyx(ct, t1c),
|
||||
"ants_affine": lambda: RP.register_ants_affine_variant(ct, t1c, "antspyx"),
|
||||
"ants_syn": lambda: RP.register_ants_variant_deformable(ct, t1c, "ants_syn"),
|
||||
}
|
||||
ha = sitk.GetArrayFromImage(t1c)
|
||||
cands = {}
|
||||
for name, fn in calls.items():
|
||||
try:
|
||||
T, meta = fn()
|
||||
except Exception as e: # noqa
|
||||
cands[name] = {"error": repr(e)}
|
||||
print(f" {name}: FAILED {e!r}", flush=True)
|
||||
continue
|
||||
warped_ct = RP._warp_ct(ct, t1c, T)
|
||||
mi = _mi_value(ha, warped_ct)
|
||||
gap = RP._skull_gap(warped_ct, ha)
|
||||
cands[name] = {"T": T, "meta": meta, "mi": mi, "gap": gap}
|
||||
print(f" {name}: mi={mi:.4f} gap={gap} ({meta.get('seconds_reg')}s)",
|
||||
flush=True)
|
||||
ok = {n: c for n, c in cands.items() if "T" in c}
|
||||
if not ok:
|
||||
raise RuntimeError("all ANTs candidates failed: "
|
||||
+ repr({n: c.get("error") for n, c in cands.items()}))
|
||||
best_name = max(ok, key=lambda n: ok[n]["mi"])
|
||||
T, meta, gap = ok[best_name]["T"], ok[best_name]["meta"], ok[best_name]["gap"]
|
||||
|
||||
def _tp(x, y, z):
|
||||
return np.array(T.TransformPoint((float(x), float(y), float(z))))
|
||||
|
||||
t0 = _tp(0, 0, 0)
|
||||
mat = np.column_stack([_tp(1, 0, 0) - t0, _tp(0, 1, 0) - t0, _tp(0, 0, 1) - t0])
|
||||
scale = [round(float(np.linalg.norm(mat[:, i])), 4) for i in range(3)]
|
||||
hm = head_mask_from_image(ha)
|
||||
warped_ct = sitk.GetArrayFromImage(sitk.Resample(ct, t1c, T, sitk.sitkLinear))
|
||||
posw = warped_ct[warped_ct > 0]
|
||||
tissue_w = ((warped_ct > max(2.0, float(np.percentile(posw, 2))))
|
||||
if posw.size else np.zeros_like(warped_ct, bool))
|
||||
tissue_w = ndimage.binary_opening(tissue_w, iterations=2)
|
||||
coverage = float(tissue_w[hm].sum() / max(hm.sum(), 1))
|
||||
info = {
|
||||
"pipeline": "ants_best",
|
||||
"fixed": "T1c", "moving": "CT",
|
||||
"best": best_name,
|
||||
"candidates": {
|
||||
n: ({**c["meta"], "mi": round(c["mi"], 4),
|
||||
"skull_ring_gap_mm": c["gap"], "selected": n == best_name}
|
||||
if "T" in c else c)
|
||||
for n, c in cands.items()},
|
||||
"affine_stage": {"matrix": [[round(float(v), 4) for v in row] for row in mat.tolist()],
|
||||
"offset": [round(float(x), 3) for x in t0], "scale": scale,
|
||||
"transform_type": type(T).__name__},
|
||||
"coverage_head_box": round(coverage, 3),
|
||||
"skull_ring_gap_mm": gap,
|
||||
"seconds_total": round(time.time() - t_start, 1),
|
||||
}
|
||||
print("reg done %.0fs best=%s mi=%.4f gap=%s cov %.2f" % (
|
||||
info["seconds_total"], best_name, ok[best_name]["mi"], gap, coverage),
|
||||
flush=True)
|
||||
return T, info, tissue_w
|
||||
|
||||
|
||||
def reg_worker(row):
|
||||
reg_dir = d("data/m6reg")
|
||||
key = row["key"]
|
||||
|
|
@ -344,14 +458,17 @@ def reg_worker(row):
|
|||
f, indent=1)
|
||||
return key, False, 0
|
||||
gtv = read_nii_img(row["label"])
|
||||
tr, info, tissue_w = register_ct_to_t1c(ct, t1c)
|
||||
tr, info, tissue_w = register_ants_best(ct, t1c)
|
||||
# QA files in m6reg, everything in (oriented) MR coordinates:
|
||||
# _ct: CT warped to the MR grid (linear), for overlaying _tv
|
||||
# _mr: the oriented MR volume itself
|
||||
# _tv: original TV warped to the MR grid (nearest) — the study label
|
||||
# _tv: original TV warped to the MR grid (SMD + linear) — the study label
|
||||
# _reg.json: registration metrics + final image information
|
||||
tv_mr = sitk.Resample(gtv, t1c, tr, sitk.sitkNearestNeighbor, 0.0)
|
||||
z = sitk.GetArrayFromImage(tv_mr)
|
||||
qa = _load_qa()
|
||||
thr = qa.smd_threshold([float(s) for s in gtv.GetSpacing()])[0]
|
||||
z = qa.warp_tv(qa.make_tv_bin(gtv), t1c, tr, "smd_linear", thr)
|
||||
tv_mr = sitk.GetImageFromArray(z)
|
||||
tv_mr.CopyInformation(t1c)
|
||||
nvox = int((z > 0.5).sum())
|
||||
from scipy import ndimage
|
||||
tv_in_ct_frac = 1.0
|
||||
|
|
@ -369,6 +486,7 @@ def reg_worker(row):
|
|||
info.update({
|
||||
"key": key, "subject": row["subject"], "date": row["date"],
|
||||
"tv_nvox": nvox, "tv_in_ct_frac": round(tv_in_ct_frac, 3),
|
||||
"tv_smd_thr_mm": round(thr, 3),
|
||||
"alignment": align,
|
||||
"tv_accepted": accepted,
|
||||
"accept_reason": ("ok" if accepted else
|
||||
|
|
|
|||
|
|
@ -28,8 +28,11 @@ Protocol
|
|||
|
||||
Usage (repo root, inside the `longitudinal-313` conda env):
|
||||
python scripts/benchmark_pipelines.py # train + eval A B C
|
||||
python scripts/benchmark_pipelines.py --pipelines A C
|
||||
python scripts/benchmark_pipelines.py --train-epochs 3 --max-rows 240
|
||||
python scripts/benchmark_pipelines.py --pipelines A C
|
||||
python scripts/benchmark_pipelines.py --train-epochs 3 --max-rows 240
|
||||
python scripts/benchmark_pipelines.py --pipelines "A B C" --resume --train-epochs 200
|
||||
# ^ continue from runs_bench/{A,C}/best.pt and the nnU-Net fold_0
|
||||
# checkpoint_latest.pth; --train-epochs is the TOTAL budget
|
||||
python scripts/benchmark_pipelines.py --no-tta # no test-time augmentation
|
||||
python scripts/benchmark_pipelines.py --skip-train # eval existing study checkpoints
|
||||
Outputs: results/benchmark_pipelines.json (+ _per_row.jsonl, _bench.png)
|
||||
|
|
@ -81,6 +84,12 @@ def parse_args():
|
|||
ap.add_argument("--skip-train", action="store_true",
|
||||
help="eval existing study checkpoints (runs/round0, runs_monai/round0, nnu/results 210) "
|
||||
"instead of training")
|
||||
ap.add_argument("--resume", action="store_true",
|
||||
help="continue training from the latest benchmark checkpoint. A/C: load "
|
||||
"runs_bench/{A,C}/best.pt (model + optimizer + scheduler when present) "
|
||||
"and train the remaining epochs. B: append nnUNet's --c flag (continue "
|
||||
"from fold_0/checkpoint_latest.pth; pair with --nnu-skip-prep). "
|
||||
"--train-epochs is the TOTAL budget, not epochs added")
|
||||
ap.add_argument("--nnu-dsid", type=int, default=None,
|
||||
help="nnU-Net dataset id: training mode default 221 (isolated; the study's 210 is never "
|
||||
"touched), --skip-train mode default 210 (study model); pass explicitly to evaluate "
|
||||
|
|
@ -208,7 +217,7 @@ def bench_A(args, train_rows, test_rows):
|
|||
model = A.build_model(base=16, device="cuda")
|
||||
patch = (args.patch, args.patch, args.patch)
|
||||
_, dl = make_dataloader(train_rows, patch, args.batch, True, num_workers=args.workers,
|
||||
seed=0, persistent=False)
|
||||
seed=0, persistent=False)
|
||||
steps = max(len(dl), 1)
|
||||
total_steps = steps * args.train_epochs
|
||||
opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-4)
|
||||
|
|
@ -218,10 +227,26 @@ def bench_A(args, train_rows, test_rows):
|
|||
torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=max(total_steps - warmup, 1),
|
||||
eta_min=3e-4 * 0.05)],
|
||||
milestones=[warmup])
|
||||
ckpt = os.path.join(ckpt_dir, "best.pt")
|
||||
start_epoch = 0
|
||||
if args.resume and os.path.exists(ckpt):
|
||||
sd = torch.load(ckpt, map_location=dev)
|
||||
model.load_state_dict(sd["model"])
|
||||
if "optimizer" in sd:
|
||||
opt.load_state_dict(sd["optimizer"])
|
||||
if "scheduler" in sd:
|
||||
sched.load_state_dict(sd["scheduler"])
|
||||
start_epoch = int(sd.get("epoch", 0))
|
||||
res["resumed_from_epoch"] = start_epoch
|
||||
print(f"[A] resume: {ckpt} epoch {start_epoch}/{args.train_epochs} "
|
||||
f"(optimizer={'restored' if 'optimizer' in sd else 'fresh'}, "
|
||||
f"scheduler={'restored' if 'scheduler' in sd else 'fresh'})", flush=True)
|
||||
elif args.resume:
|
||||
print(f"[A] --resume: {ckpt} not found, training from scratch", flush=True)
|
||||
model.train()
|
||||
reset_peak(dev)
|
||||
t0 = time.time()
|
||||
for ep in range(args.train_epochs):
|
||||
for ep in range(start_epoch, args.train_epochs):
|
||||
if hasattr(dl.sampler, "set_epoch"):
|
||||
dl.sampler.set_epoch(ep)
|
||||
loss_sum = n = 0.0
|
||||
|
|
@ -245,11 +270,11 @@ def bench_A(args, train_rows, test_rows):
|
|||
res.setdefault("per_epoch_loss", []).append(round(loss_sum / max(n, 1), 4))
|
||||
dt = time.time() - t0
|
||||
res["train_s"] = round(dt, 1)
|
||||
res["train_s_per_epoch"] = round(dt / max(args.train_epochs, 1), 1)
|
||||
res["train_s_per_epoch"] = round(dt / max(args.train_epochs - start_epoch, 1), 1)
|
||||
res["train_vram_mb"] = peak_mb(dev)
|
||||
res["steps_per_epoch"] = steps
|
||||
ckpt = os.path.join(ckpt_dir, "best.pt")
|
||||
torch.save({"model": model.state_dict(), "epoch": args.train_epochs}, ckpt)
|
||||
torch.save({"model": model.state_dict(), "epoch": args.train_epochs,
|
||||
"optimizer": opt.state_dict(), "scheduler": sched.state_dict()}, ckpt)
|
||||
del dl, opt, sched
|
||||
res["ckpt"] = res.get("ckpt", os.path.join(ckpt_dir, "best.pt"))
|
||||
res["params"] = int(sum(p.numel() for p in model.parameters()))
|
||||
|
|
@ -326,10 +351,26 @@ def bench_C(args, train_rows, test_rows):
|
|||
eta_min=3e-4 * 0.05)],
|
||||
milestones=[warmup])
|
||||
loss_f = WeightedDiceCELoss()
|
||||
ckpt = os.path.join(ckpt_dir, "best.pt")
|
||||
start_epoch = 0
|
||||
if args.resume and os.path.exists(ckpt):
|
||||
sd = torch.load(ckpt, map_location=dev)
|
||||
model.load_state_dict(sd["model"])
|
||||
if "optimizer" in sd:
|
||||
opt.load_state_dict(sd["optimizer"])
|
||||
if "scheduler" in sd:
|
||||
sched.load_state_dict(sd["scheduler"])
|
||||
start_epoch = int(sd.get("epoch", 0))
|
||||
res["resumed_from_epoch"] = start_epoch
|
||||
print(f"[C] resume: {ckpt} epoch {start_epoch}/{args.train_epochs} "
|
||||
f"(optimizer={'restored' if 'optimizer' in sd else 'fresh'}, "
|
||||
f"scheduler={'restored' if 'scheduler' in sd else 'fresh'})", flush=True)
|
||||
elif args.resume:
|
||||
print(f"[C] --resume: {ckpt} not found, training from scratch", flush=True)
|
||||
model.train()
|
||||
reset_peak(dev)
|
||||
t0 = time.time()
|
||||
for ep in range(args.train_epochs):
|
||||
for ep in range(start_epoch, args.train_epochs):
|
||||
if hasattr(dl.sampler, "set_epoch"):
|
||||
dl.sampler.set_epoch(ep)
|
||||
loss_sum = n = 0.0
|
||||
|
|
@ -350,11 +391,11 @@ def bench_C(args, train_rows, test_rows):
|
|||
res.setdefault("per_epoch_loss", []).append(round(loss_sum / max(n, 1), 4))
|
||||
dt = time.time() - t0
|
||||
res["train_s"] = round(dt, 1)
|
||||
res["train_s_per_epoch"] = round(dt / max(args.train_epochs, 1), 1)
|
||||
res["train_s_per_epoch"] = round(dt / max(args.train_epochs - start_epoch, 1), 1)
|
||||
res["train_vram_mb"] = peak_mb(dev)
|
||||
res["steps_per_epoch"] = steps
|
||||
ckpt = os.path.join(ckpt_dir, "best.pt")
|
||||
torch.save({"model": model.state_dict(), "epoch": args.train_epochs}, ckpt)
|
||||
torch.save({"model": model.state_dict(), "epoch": args.train_epochs,
|
||||
"optimizer": opt.state_dict(), "scheduler": sched.state_dict()}, ckpt)
|
||||
del dl, opt, sched, loss_f
|
||||
res["ckpt"] = res.get("ckpt", os.path.join(ckpt_dir, "best.pt"))
|
||||
res["params"] = int(sum(p.numel() for p in model.parameters()))
|
||||
|
|
@ -432,10 +473,21 @@ def bench_B(args, train_rows, test_rows):
|
|||
|
||||
env = N.nnu_env(epoch=args.train_epochs, lr=1e-2)
|
||||
env["CUDA_VISIBLE_DEVICES"] = str(args.gpu)
|
||||
bcmd = N.train_cmd(1)
|
||||
if args.resume:
|
||||
# nnUNet --c: continue fold_0 from checkpoint_latest.pth (optimizer +
|
||||
# lr-schedule state are restored natively); trains up to PL_EPOCHS
|
||||
# (= --train-epochs, the total budget)
|
||||
bcmd.append("--c")
|
||||
res["resumed_from_checkpoint_latest"] = True
|
||||
if not args.nnu_skip_prep:
|
||||
print("[B] --resume without --nnu-skip-prep: plan_and_preprocess "
|
||||
"still runs with --clean (the trained checkpoint is kept)",
|
||||
flush=True)
|
||||
poller = SmiPoller(args.gpu)
|
||||
poller.start()
|
||||
t0 = time.time()
|
||||
N.run(N.train_cmd(1), os.path.join(d("logs/nnu"), "bench_train.log"), env=env, retries=0)
|
||||
N.run(bcmd, os.path.join(d("logs/nnu"), "bench_train.log"), env=env, retries=0)
|
||||
res["train_s"] = round(time.time() - t0, 1)
|
||||
res["train_s_per_epoch"] = round(res["train_s"] / max(args.train_epochs, 1), 1)
|
||||
res["train_vram_mb"] = round(poller.max_mib * 1.048576, 1) if poller.max_mib else None
|
||||
|
|
|
|||
Loading…
Reference in a new issue