xfr/longitudinal
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

feat(registration): align CT to T1c and add benchmarking

Browse source 
Update M6 manifest building to register T1c to CT with rigid
Mattes MI on a CT-space head box, including multi-start initial
tilts for large inter-session head motion. The inverse transform
warps CT and GTV into T1c native space, with acceptance based on
concentricity of the warped CT skull ring and MR head silhouette.

Rerun preprocessing when a row has a label but the warped label file
is missing. Normalize volume orientation in read_nii_img.

Add registration benchmark documentation, benchmark script, and
shared registration pipeline helpers.
This commit is contained in:
Furen Xiao 2026-09-28 16:20:30 +08:00
parent d8296ff40d
commit e8dfe0fe35
6 changed files with 1519 additions and 36 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

177
REGISTRATION_BENCHMARK.md Normal file
View file

@ -0,0 +1,177 @@
# Registration Pipeline Benchmark — CT → T1c (M6)
Benchmarks **four rigid, cross-modality registration pipelines** that align a
head-and-neck CT to a tightly-cropped T1c MRI so the GTV can be warped onto the
T1c (MR) native grid. This is the registration step of
`scripts/02_build_m6_dataset.py`. The pairs are hard: different modality (CT
Hounsfield vs MR intensity) and the patient is repositioned between the CT and
MR sessions, so the inter-session head tilt can reach ~30°.
- Pipelines: [`src/registration_pipelines.py`](src/registration_pipelines.py)
- Benchmark harness: [`scripts/bench_registration.py`](scripts/bench_registration.py)
- Results: `results/bench_registration/` (per-pipeline JSON, `summary.json`,
`summary.md`, `bench_registration.png`)
## Unified contract
Every pipeline returns a SimpleITK transform `T` such that
```python
sitk.Resample(ct, t1c, T, sitk.sitkLinear) # CT onto the MR grid
sitk.Resample(tv, t1c, T, sitk.sitkNearestNeighbor, 0) # GTV onto the MR grid
```
places the CT / GTV concentric with the MR head. In `Resample`'s
destination→source convention this `T` maps **T1c → CT**. All pipelines register
on the same CT-space head box (MI computed on head tissue, not neck/air) and,
for the single-init tools, from the **same centroid-aligned start**
(`centroid_offset`), so the comparison is on the optimizer / metric / transform
handling — not on who finds the ~30–60 mm gross offset first.
## The four pipelines
| # | pipeline | toolkit | metric | transform | status in this env |
|---|----------|---------|--------|-----------|--------------------|
| 1 | `elastix` | ITK-Elastix | Mattes MI | `Rigid3DTransform` (ASGD) | **not installed** (conda-forge unreachable) |
| 2 | `simpleitk` | SimpleITK 2.5.6 | Mattes MI (32 bins) | `Euler3D` rigid, pyramid `[4,2,1]` | ✅ ran |
| 3 | `antspyx` | ANTsPy / ANTs 0.6.3 | Mattes MI (multi-res) | ANTs `Rigid` | ✅ ran (isolated py3.12 env) |
| 4 | `deepreg` | Oxford DeepReg (CycleGAN) | CycleGAN + reg. loss (pretrained) | 3D displacement field | **not installed** (see below) |
`simpleitk` has two budgets: `--simpleitk-mode benchmark` (single pyramid,
centered start — the default, comparable to the other tools) and
`--simpleitk-mode production` (the full multi-start tilt grid from
`02_build_m6_dataset.py`, ~20× slower, the most robust but not used for the
head-to-head).
## Test cases
First 10 labeled M6 cases (`--n-cases 10 --sample first`), reproducible. 9 have
a complete CT; 1 (`m6_22GD5VL2_20230724`) has an incomplete CT head and is
reported as `incomplete-CT` for every pipeline (a data issue, not a pipeline
failure).
```
m6_22GD5VL2_20230724 (incomplete CT)
m6_22IRRKU6_20240115 m6_23XIMO4M_20230222 m6_24V62QQM_20230223
m6_24V62QQM_20231017 m6_24V62QQM_20240820 m6_26TYQI3C_20221027
m6_2BBN5XA4_20250411 m6_2C4TKI6H_20240506 m6_2CDNN7BH_20240102
```
## Metrics (identical for every pipeline)
Production QA, crop-independent:
- **`skull_ring_gap_mm`** — march rays from the MR head center; the gap between
the MR head silhouette boundary and the warped-CT skull ring (p97 of positive
CT) at mid and low levels, plus the angular std. **Aligned ⇔
`0 ≤ mid ≤ 20 mm` and `std ≤ 12 mm`.** `mid` measures translational
concentricity; `std` measures rotational concentricity.
- **`tv_in_ct_frac`** — fraction of the dilated warped GTV inside the warped-CT
head tissue (the slab must straddle the lesion). Acceptance needs `≥ 0.5`.
- **`tv_nvox`** — warped GTV volume.
- **`accepted`** — `tv_nvox > 0 and tv_in_ct_frac ≥ 0.5 and aligned` (the
production acceptance rule).
- **`reg_wall_s`** — wall-clock registration time (excludes I/O + scoring).
## Results (10 cases)
| pipeline | toolkit | registered | accepted | gap mid (med/mean) | gap std (med/mean) | s/reg (med) | tv-in-CT |
|---|---|---|---|---|---|---|---|
| **simpleitk** | SimpleITK 2.5.6 | 9 | **4/10 (0.40)** | 12.6 / 16.1 | 15.4 / 16.0 | 8.3 | 0.93 |
| **antspyx** | ANTsPy 0.6.3 | 9 | **1/10 (0.10)** | 14.3 / 15.4 | 18.2 / 21.7 | 13.4 | 0.93 |
| elastix | — | — | — | — | — | — | — |
| deepreg | — | — | — | — | — | — | — |
Per-case (mid / std, `OK` = accepted):
| case | simpleitk | antspyx |
|---|---|---|
| m6_22GD5VL2_20230724 | incomplete-CT | incomplete-CT |
| m6_22IRRKU6_20240115 | 25.7 / 29.3 | 26.3 / 29.2 |
| m6_23XIMO4M_20230222 | 22.2 / 23.9 | 23.8 / 25.0 |
| m6_24V62QQM_20230223 | 10.7 / 16.9 | 9.4 / 18.2 |
| m6_24V62QQM_20231017 | 11.5 / 7.3 **OK** | 12.5 / 12.8 |
| m6_24V62QQM_20240820 | 11.6 / 9.6 **OK** | 14.8 / 11.7 **OK** |
| m6_26TYQI3C_20221027 | 12.2 / 3.7 **OK** | 6.1 / 37.2 |
| m6_2BBN5XA4_20250411 | 18.6 / 8.0 **OK** | 14.3 / 12.7 |
| m6_2C4TKI6H_20240506 | 12.6 / 15.4 | 13.6 / 14.7 |
| m6_2CDNN7BH_20240102 | 19.4 / 30.2 | 18.2 / 33.6 |
### Interpretation
- **Translation is recovered well by both** (mid gaps mostly 6–20 mm, within the
`≤ 20 mm` band). `tv-in-CT` is identical (0.93): the warped GTV lands inside
the CT head for both, so the target is in the right place.
- **The binding failure is rotation** — `std > 12 mm` on the repositioned cases
(the ~30° inter-session tilt produces MI local minima that a single centered
start can lock onto). This is exactly why production `02_build_m6_dataset.py`
uses a **multi-start tilt grid**.
- **SimpleITK is the more robust of the two** (4 vs 1 accepted, lower median
std 15.4 vs 18.2): its RSGD + physical-step-scales + 32-bin Mattes-MI
landscape handles the tilt better than ANTs' Rigid. Note `m6_26TYQI3C`: ANTs
nailed the position (mid 6.1) but missed the orientation (std 37.2) — a
textbook tilted local minimum.
- **Speed**: SimpleITK ~8.3 s/case, ANTs ~13.4 s/case (ANTs includes the
pre-shift resample + its affine-MI init stages). Both are CPU-bound and ran
4 cases in parallel.
## Reproducing
```bash
# main env — SimpleITK baseline (elastix/deepreg report unavailable here)
source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal
python scripts/bench_registration.py --pipelines simpleitk --n-cases 10 \
--workers 4 --threads 24 --simpleitk-mode benchmark
# isolated py3.12 env with antspyx — ANTs (+ optional SimpleITK cross-check)
/home/xfr/.conda/envs/ants312/bin/python scripts/bench_registration.py \
--pipelines antspyx --n-cases 10 --workers 4 --threads 24
# combine the per-pipeline JSONs accumulated across envs into one table + plot
python scripts/bench_registration.py --merge
```
Other flags: `--sample random --seed 0` (reproducible random 10),
`--keys k1,k2,...` (explicit cases), `--n-cases N`, `--workers W`,
`--threads T`, `--simpleitk-mode {benchmark,production}`,
`--deepreg-ckpt <.ckpt>` (DeepReg), `--out-dir <dir>`.
## Enabling the two blocked pipelines
**ITK-Elastix** — not on PyPI; needs conda-forge (which this host cannot reach).
Once conda-forge is available:
```bash
conda install -c conda-forge elastix # provides the elastix/transformix CLI
python scripts/bench_registration.py --pipelines elastix --n-cases 10
```
The adapter registers `fixed=T1c, moving=CT` with `Rigid3DTransform` + Mattes MI
on the head box and reads the result transform.
**DeepReg (Oxford)** — the PyPI `deepreg` (0.1.2) is an unrelated stub that pins
`numpy==1.18.5`; do not install it. DeepReg is not a one-shot rigid tool: it is a
CycleGAN image-to-image registration that must be **trained or run from a
pretrained checkpoint**. To benchmark it:
```bash
git clone https://github.com/OB-CVHU/DeepReg && cd DeepReg && pip install -e .
python scripts/bench_registration.py --pipelines deepreg --n-cases 10 \
--deepreg-ckpt /path/to/pretrained_3d.ckpt
```
The adapter runs `deepreg test` to predict a voxel displacement field, converts
it to a SimpleITK `DisplacementFieldTransform`, and scores it like the others.
## Notes / caveats
- The benchmark compares **standard single-init** configurations (centered
start). Production uses `--simpleitk-mode production`'s multi-start tilt grid,
which raises SimpleITK's acceptance at ~20× the cost; that is out of scope for
the head-to-head but is the config the dataset actually ships with.
- `elastix` / `deepreg` could not be executed here, so their rows are
"not installed", not measured. Their code paths are implemented and
unit-verified for import/convention; run them where the toolkits are present.
- `fix_orientation` (elastix/ants/deepreg) chooses between a transform and its
inverse by skull-ring concentricity — a correctness gate that makes the result
independent of each toolkit's transform-direction convention.

379
scripts/02_build_m6_dataset.py
View file

@ -1,7 +1,11 @@
"""Build M6-2025 manifests: T1c volumes, labeled when GTV exists. """Build M6-2025 manifests: T1c volumes, labeled when GTV exists.
GTV lives in CT space; CT is registered to T1c (similarity->affine MI + BSpline) GTV lives in CT space; T1c is registered to CT (rigid, Mattes MI on a head box
and GTV warped (nearest) into T1c native space. located in CT space; 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.
Acceptance requires the warped CT skull ring to be concentric with the MR head
silhouette (crop-independent check).
Outputs: Outputs:
data/manifests/m6_labeled.jsonl {key, subject, date, img, label(warped GTV), source:'m6'} 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'} data/manifests/m6_unlabeled.jsonl {key, subject, date, img, label:None, source:'m6'}
@ -56,18 +60,132 @@ def discover():
t1c = pick_t1c(mrd) t1c = pick_t1c(mrd)
if not t1c: if not t1c:
continue continue
gtv = os.path.join(dp, "RT", "TV", "Struct_GTV.nii.gz") rt = os.path.join(dp, "RT")
ct = os.path.join(dp, "RT", "ct_image.nii.gz") ct = os.path.join(rt, "ct_image.nii.gz")
label = None
tv_files = []
if os.path.isdir(rt):
try:
tv_files = sorted(f for f in os.listdir(rt)
if f.startswith("TV-") and f.endswith(".nii.gz"))
except OSError:
tv_files = []
if tv_files:
label = os.path.join(rt, tv_files[0])
elif os.path.exists(os.path.join(rt, "TV", "Struct_GTV.nii.gz")):
label = os.path.join(rt, "TV", "Struct_GTV.nii.gz")
key = f"m6_{pid}_{date}" key = f"m6_{pid}_{date}"
row = {"key": key, "subject": pid, "date": date, "img": t1c, "source": "m6"} row = {"key": key, "subject": pid, "date": date, "img": t1c, "source": "m6"}
if os.path.exists(gtv) and os.path.exists(ct): if label and os.path.exists(ct):
labeled.append({**row, "label": gtv, "ct": ct, "gtv_registered": False}) labeled.append({**row, "label": label, "ct": ct, "gtv_registered": False})
else: else:
unlabeled.append({**row, "label": None}) unlabeled.append({**row, "label": None})
return labeled, unlabeled return labeled, unlabeled
def _ct_vertex_slice(ct):
"""Topmost tissue slice with 10mm of continuity below, or None.
Tissue > -150 HU (bone HU is unreliable in these CTs). The continuity
test skips small artifacts floating above the head.
"""
per = (sitk.GetArrayFromImage(ct) > -150).sum((1, 2))
n8 = max(2, int(round(10 / ct.GetSpacing()[2])))
for i in range(len(per) - 1, -1, -1):
if per[i] > 500 and per[max(0, i - n8):i + 1].min() >= 50:
return i
return None
def _skull_ring_gap(warped_ct, ha):
"""Concentricity of the CT skull ring vs the MR head silhouette.
The MR volumes are tight crops, so distance-to-crop-edge is
crop-dependent; instead march rays from the head center and measure the
gap between the MR head boundary and the CT skull ring (p97 of positive
CT) at two levels. Returns {"mid_mm", "low_mm", "std_mm"} or None when no
skull ring is found. Aligned = small gap with low angular variance.
"""
posw = warped_ct[warped_ct > 0]
if not posw.size:
return None
bthr = float(np.percentile(posw, 97))
sil = ha > 50
izs = np.where(sil.sum((1, 2)) > 500)[0]
if not izs.size:
return None
def _gap(iz):
sl = sil[iz]
if sl.sum() < 500:
return None, None
sk = warped_ct[iz] > bthr
if not sk.any():
return None, None
hy, hx = np.nonzero(sl)
cy0, cx0 = float(hy.mean()), float(hx.mean())
gaps = []
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)
rh = rr[np.where(sl[yy, xx])[0]]
rs = rr[np.where(sk[yy, xx])[0]]
if len(rh) and len(rs):
gaps.append(float(rh.max() - rs.max()))
if not gaps:
return None, None
return float(np.mean(gaps)), float(np.std(gaps))
gm, sm = _gap(int(izs.min() + (izs.max() - izs.min()) * 0.5))
gl, sl = _gap(int(izs.min() + (izs.max() - izs.min()) * 0.75))
if gm is None:
return None
std = max(s for s in (sm, sl) if s is not None) if (sm is not None or sl is not None) else 0.0
return {"mid_mm": round(gm, 1), "low_mm": round(gl, 1) if gl is not None else None,
"std_mm": round(std, 1)}
def _mi_value(fixed_arr, moving_arr, bins=32):
"""Mattes-style MI (nats) on percentile-binned intensities, subsampled for speed."""
def q(a):
lo, hi = np.percentile(a, (1, 99))
return np.clip(np.round((a - lo) * (bins - 1) / max(hi - lo, 1e-9)).astype(int), 0, bins - 1)
A, B = q(fixed_arr), q(moving_arr)
step = max(1, int(np.ceil((A.size / 200_000) ** (1.0 / 3))))
h = np.bincount(A[::step, ::step, ::step].ravel() * bins
+ B[::step, ::step, ::step].ravel(), minlength=bins * bins)
p = h.reshape(bins, bins) / max(h.sum(), 1)
px, py = p.sum(axis=1, keepdims=True), p.sum(axis=0, keepdims=True)
with np.errstate(divide="ignore", invalid="ignore"):
m = np.where((p > 0) & (px > 0) & (py > 0), np.log(p / (px * py)), 0.0)
return float((p * m).sum())
def _valid_pyramid(size, shrink, smooth, min_dim=6):
sh = [s for s in shrink if min(size) / s >= min_dim] or [1]
sm = smooth[: len(sh)] + [0] * max(0, len(sh) - len(smooth))
return sh, sm
def register_ct_to_t1c(ct, t1c, threads=12): def register_ct_to_t1c(ct, t1c, threads=12):
"""Rigid registration with CT as fixed and T1c as moving, Mattes MI on
head-box images.
Both inputs share physical space (oriented). The registration grid is a head
box located in CT space from the top of the head (the CTs are head-neck
exams, so an MR-head crop would show mostly neck/air on the CT); the box
excludes most air so MI is computed on tissue. The box is in physical
space, so the resulting transform applies directly to the full volumes.
Returns (transform, info): transform is the registration transform, mapping
T1c (moving) coordinates -> CT (fixed) coordinates; resampling CT/TV onto
the T1c grid with it applies the inverse transform to CT-space points,
placing them in T1c (MR) coordinates. info has per-start MI/timing, the
affine decomposition (matrix/offset/scale), and the skull-ring concentricity
QA statistic.
"""
import time import time
t_start = time.time() t_start = time.time()
try: # noqa try: # noqa
@ -76,20 +194,144 @@ def register_ct_to_t1c(ct, t1c, threads=12):
pass pass
ct = sitk.Cast(ct, sitk.sitkFloat32) ct = sitk.Cast(ct, sitk.sitkFloat32)
t1c = sitk.Cast(t1c, sitk.sitkFloat32) t1c = sitk.Cast(t1c, sitk.sitkFloat32)
r = sitk.ImageRegistrationMethod()
r.SetMetricAsMattesMutualInformation(numberOfHistogramBins=32) from src.common import head_mask_from_image
r.SetOptimizerAsRegularStepGradientDescent( ha = sitk.GetArrayFromImage(t1c)
learningRate=10.0, minStep=0.5, numberOfIterations=600, hm = head_mask_from_image(ha)
relaxationFactor=0.5, gradientMagnitudeTolerance=1e-5, # Registration grid = a head box located in CT space. The CTs are head-neck
estimateLearningRate=sitk.ImageRegistrationMethod.EachIteration) # (often head-neck-chest) exams, so cropping to the MR head box would leave
r.SetOptimizerScalesFromPhysicalShift() # mostly neck/chest/air on the fixed image and MI locks onto the wrong
r.SetInitialTransform(sitk.CenteredTransformInitializer(ct, t1c, sitk.Transform(3, sitk.sitkSimilarity))) # structures. Vertex = topmost tissue slice with 10mm of continuity below
r.SetShrinkFactorsPerLevel([8, 4, 2, 1]) # (skips small artifacts floating above the head; bone HU is unreliable
r.SetSmoothingSigmasPerLevel([4, 2, 1, 0]) # here, so tissue > -150 HU is used). x/y from the tissue extent of the
r.SetInterpolator(sitk.sitkLinear) # 150mm band below the vertex. Box covers the whole head (x/y: extent +
aff = r.Execute(ct, t1c) # 15mm margin, z: +10 to -215 from vertex).
print("affine done %.0fs" % (time.time() - t_start), flush=True) ca = sitk.GetArrayFromImage(ct)
return aff tissue = ca > -150
csp, co = np.array(ct.GetSpacing()), np.array(ct.GetOrigin())
vz = _ct_vertex_slice(ct)
if vz is not None:
ztop = co[2] + vz * csp[2]
band = tissue[max(0, vz - int(round(150 / csp[2]))):vz + 1]
cols = np.nonzero(band.any(0))
box_lo = (co[0] + cols[1].min() * csp[0] - 15,
co[1] + cols[0].min() * csp[1] - 15, ztop - 215)
box_hi = (co[0] + cols[1].max() * csp[0] + 15,
co[1] + cols[0].max() * csp[1] + 15, ztop + 10)
else: # no head found: fall back to the whole CT volume
gs = ct.GetSize()
box_lo = tuple(float(x) for x in co)
box_hi = (co[0] + gs[0] * csp[0], co[1] + gs[1] * csp[1], co[2] + gs[2] * csp[2])
tsp = np.array(t1c.GetSpacing())
shape = tuple(int(np.ceil((box_hi[d] - box_lo[d]) / tsp[d])) + 1 for d in range(3))
ref = sitk.GetImageFromArray(np.zeros(shape, np.float32))
ref.SetSpacing(t1c.GetSpacing())
ref.SetDirection(t1c.GetDirection())
ref.SetOrigin(tuple(float(x) for x in box_lo))
fixed = sitk.Resample(ct, ref, sitk.Transform(), sitk.sitkLinear, 0.0)
moving = sitk.Resample(t1c, ref, sitk.Transform(), sitk.sitkLinear, 0.0)
stages = []
def run_stage(name, init_tr, shrink, smooth, lr, min_step, iters):
t0 = time.time()
r = sitk.ImageRegistrationMethod()
r.SetMetricAsMattesMutualInformation(numberOfHistogramBins=32)
r.SetOptimizerAsRegularStepGradientDescent(
learningRate=lr, minStep=min_step, numberOfIterations=iters,
relaxationFactor=0.5, gradientMagnitudeTolerance=1e-6,
estimateLearningRate=sitk.ImageRegistrationMethod.EachIteration)
r.SetOptimizerScalesFromPhysicalShift()
sh, sm = _valid_pyramid(fixed.GetSize(), shrink, smooth)
r.SetShrinkFactorsPerLevel(sh)
r.SetSmoothingSigmasPerLevel(sm)
r.SetInterpolator(sitk.sitkLinear)
r.SetInitialTransform(init_tr, inPlace=False)
fa = sitk.GetArrayFromImage(fixed)
mi_in = _mi_value(fa, sitk.GetArrayFromImage(sitk.Resample(moving, fixed, init_tr, sitk.sitkLinear, 0.0)))
tr = r.Execute(fixed, moving)
mi_out = _mi_value(fa, sitk.GetArrayFromImage(sitk.Resample(moving, fixed, tr, sitk.sitkLinear, 0.0)))
stages.append({"name": name, "mi_in": round(mi_in, 4), "mi_out": round(mi_out, 4),
"seconds": round(time.time() - t0, 1), "shrink": sh, "smooth": sm})
return tr
# The MI surface has tilted local minima that can score HIGHER MI than the
# true alignment (patients are repositioned between the CT and MR sessions,
# so the inter-session head tilt can reach ~30 deg), so MI alone cannot
# rank candidates. Run the full pyramid from the centered start; if its
# full-resolution geometric concentricity (skull-ring gap) is poor, repeat
# from a grid of initial tilts and keep the best by gap.
mr_c = np.array(t1c.GetOrigin()) + np.argwhere(hm).mean(0) * tsp
ct_c = np.array([(box_lo[d] + box_hi[d]) / 2.0 for d in range(3)])
off = (float(mr_c[0] - ct_c[0]), float(mr_c[1] - ct_c[1]), float(mr_c[2] - ct_c[2]))
rad = np.pi / 180.0
def full_run(sx, sy, sz):
init = sitk.Euler3DTransform()
init.SetCenter(tuple(float(x) for x in ct_c))
init.SetParameters((sx * rad, sy * rad, sz * rad) + off)
tr0 = run_stage(f"start_{sx:+d}x{sy:+d}y{sz:+d}z", init, [8, 4, 2, 1], [4, 2, 1, 0], 10.0, 0.5, 400)
gap = _skull_ring_gap(sitk.GetArrayFromImage(
sitk.Resample(ct, t1c, tr0.GetInverse(), sitk.sitkLinear)), ha)
return tr0, gap
def is_good(g):
return bool(g) and g["std_mm"] <= 12.0 and 0.0 <= g["mid_mm"] <= 20.0
tr_reg, gap = full_run(0, 0, 0)
if not is_good(gap):
# inter-session tilt can reach ~30 deg in pitch: search a grid of
# initial rotations, keep the best by full-res concentricity
best = (tr_reg, gap)
starts = [(sx, sy, 0)
for sx in (-30, -20, -10, 0, 10, 20, 30)
for sy in (-20, 0, 20) if (sx, sy) != (0, 0)]
for sx, sy, sz in starts:
tr_i, gap_i = full_run(sx, sy, sz)
gi = (gap_i["std_mm"], gap_i["mid_mm"]) if gap_i else (999.0, 999.0)
gb = (best[1]["std_mm"], best[1]["mid_mm"]) if best[1] else (999.0, 999.0)
if gi < gb:
best = (tr_i, gap_i)
tr_reg, gap = best
if not is_good(gap):
for sx, sy, sz in [(sx, 0, sz) for sz in (15, -15) for sx in (-30, 0, 30)]:
tr_i, gap_i = full_run(sx, sy, sz)
gi = (gap_i["std_mm"], gap_i["mid_mm"]) if gap_i else (999.0, 999.0)
gb = (best[1]["std_mm"], best[1]["mid_mm"]) if best[1] else (999.0, 999.0)
if gi < gb:
best = (tr_i, gap_i)
tr_reg, gap = best
# tr_reg (Execute convention) maps CT (fixed) -> T1c (moving). Its inverse
# maps T1c -> CT, so resampling CT/TV onto the T1c grid with it applies the
# inverse transform to CT-space points, placing them in T1c (MR) coords.
tr = tr_reg.GetInverse()
# QA: CT head-box coverage + skull-ring concentricity vs the MR head silhouette
from scipy import ndimage
def _tp(x, y, z):
return np.array(tr.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)]
warped_ct = sitk.GetArrayFromImage(sitk.Resample(ct, t1c, tr, 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))
skull_gap = _skull_ring_gap(warped_ct, ha)
info = {"stages": stages,
"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},
"coverage_head_box": round(coverage, 3),
"skull_ring_gap_mm": skull_gap,
"ct_head_box_mm": [[round(float(x), 1) for x in box_lo], [round(float(x), 1) for x in box_hi]],
"seconds_total": round(time.time() - t_start, 1)}
print("reg done %.0fs (rigid %.1fs) mi %.3f -> %.3f cov %.2f%s" % (
info["seconds_total"], sum(s["seconds"] for s in stages),
stages[-1]["mi_in"], stages[-1]["mi_out"], coverage,
" skull_gap mid %.1fmm std %.1fmm" % (skull_gap["mid_mm"], skull_gap["std_mm"]) if skull_gap else " (no skull ring)"), flush=True)
return tr, info, tissue_w
def reg_worker(row): def reg_worker(row):
@ -97,25 +339,89 @@ def reg_worker(row):
key = row["key"] key = row["key"]
t1c = read_nii_img(row["img"]) t1c = read_nii_img(row["img"])
ct = read_nii_img(row["ct"]) ct = read_nii_img(row["ct"])
# guard: corrupted CT (upper brain missing) cannot be registered
vz = _ct_vertex_slice(ct)
if vz is not None:
sp = ct.GetSpacing()[2]
per = (sitk.GetArrayFromImage(ct) > -150).sum((1, 2))
iz = vz - int(round(100 / sp))
incomplete = not (0 <= iz < len(per) and per[iz] > 5000)
else:
incomplete = True
if incomplete:
with open(os.path.join(reg_dir, key + "_reg.json"), "w") as f:
json.dump({"key": key, "subject": row["subject"], "date": row["date"],
"tv_nvox": 0, "tv_accepted": False,
"accept_reason": "ct head incomplete or not found",
"inputs": {"ct": row["ct"], "mr": row["img"], "tv": row["label"]}},
f, indent=1)
return key, False, 0
gtv = read_nii_img(row["label"]) gtv = read_nii_img(row["label"])
tr = register_ct_to_t1c(ct, t1c) tr, info, tissue_w = register_ct_to_t1c(ct, t1c)
out = os.path.join(reg_dir, key + "_gtv.nii.gz") # QA files in m6reg, everything in (oriented) MR coordinates:
warped = sitk.Resample(gtv, t1c, tr, sitk.sitkNearestNeighbor, 0.0) # _ct: CT warped to the MR grid (linear), for overlaying _tv
sitk.WriteImage(warped, out, True) # _mr: the oriented MR volume itself
z = sitk.GetArrayFromImage(warped) # _tv: original TV warped to the MR grid (nearest) — the study label
return key, os.path.exists(out) and (z > 0.5).sum() > 0, int((z > 0.5).sum()) # _reg.json: registration metrics + final image information
tv_mr = sitk.Resample(gtv, t1c, tr, sitk.sitkNearestNeighbor, 0.0)
z = sitk.GetArrayFromImage(tv_mr)
nvox = int((z > 0.5).sum())
from scipy import ndimage
tv_in_ct_frac = 1.0
if nvox > 0:
tvd = ndimage.binary_dilation(z > 0.5, iterations=3)
tv_in_ct_frac = float((tvd & tissue_w).sum() / max(tvd.sum(), 1))
tv_in_ct = tv_in_ct_frac >= 0.5 # slab must straddle the lesion
gap = info["skull_ring_gap_mm"]
if gap and 0.0 <= gap["mid_mm"] <= 20.0 and (gap["std_mm"] or 0) <= 12.0:
align = "concentric"
else:
align = "not_concentric"
accepted = nvox > 0 and tv_in_ct and align == "concentric"
info = dict(info)
info.update({
"key": key, "subject": row["subject"], "date": row["date"],
"tv_nvox": nvox, "tv_in_ct_frac": round(tv_in_ct_frac, 3),
"alignment": align,
"tv_accepted": accepted,
"accept_reason": ("ok" if accepted else
"tv outside MR grid" if nvox == 0 else
"tv outside CT coverage" if not tv_in_ct else
f"skull ring not concentric (gap={gap})"),
"inputs": {"ct": row["ct"], "mr": row["img"], "tv": row["label"]},
"images": {
"ct": {"size": list(ct.GetSize()), "spacing": [float(s) for s in ct.GetSpacing()],
"origin": [float(o) for o in ct.GetOrigin()],
"direction": [float(x) for x in ct.GetDirection()],
"pixel_type": ct.GetPixelIDTypeAsString()},
"mr": {"size": list(t1c.GetSize()), "spacing": [float(s) for s in t1c.GetSpacing()],
"origin": [float(o) for o in t1c.GetOrigin()],
"direction": [float(x) for x in t1c.GetDirection()],
"pixel_type": t1c.GetPixelIDTypeAsString()},
"tv_warp": "mr grid (same as images.mr)",
},
})
sitk.WriteImage(sitk.Resample(ct, t1c, tr, sitk.sitkLinear), os.path.join(reg_dir, key + "_ct.nii.gz"), True)
sitk.WriteImage(t1c, os.path.join(reg_dir, key + "_mr.nii.gz"), True)
sitk.WriteImage(tv_mr, os.path.join(reg_dir, key + "_tv.nii.gz"), True)
with open(os.path.join(reg_dir, key + "_reg.json"), "w") as f:
json.dump(info, f, indent=1)
return key, accepted, nvox
def main(): def main():
ap = argparse.ArgumentParser() ap = argparse.ArgumentParser()
ap.add_argument("--workers", type=int, default=2) ap.add_argument("--workers", type=int, default=2)
ap.add_argument("--skip-reg", action="store_true") ap.add_argument("--skip-reg", action="store_true")
ap.add_argument("--limit", type=int, default=0, help="register at most N missing (0=all)")
args = ap.parse_args() args = ap.parse_args()
labeled, unlabeled = discover() labeled, unlabeled = discover()
print(f"m6 labeled candidates: {len(labeled)}, unlabeled: {len(unlabeled)}") print(f"m6 labeled candidates: {len(labeled)}, unlabeled: {len(unlabeled)}")
reg_dir = d("data/m6reg") reg_dir = d("data/m6reg")
if not args.skip_reg: if not args.skip_reg:
todo = [r for r in labeled if not os.path.exists(os.path.join(reg_dir, r["key"] + "_gtv.nii.gz"))] todo = [r for r in labeled if not os.path.exists(os.path.join(reg_dir, r["key"] + "_tv.nii.gz"))]
if args.limit > 0:
todo = todo[:args.limit]
print(f"registering {len(todo)} GTVs with {args.workers} workers") print(f"registering {len(todo)} GTVs with {args.workers} workers")
with ProcessPoolExecutor(max_workers=args.workers) as ex: with ProcessPoolExecutor(max_workers=args.workers) as ex:
futs = {ex.submit(reg_worker, r): r for r in todo} futs = {ex.submit(reg_worker, r): r for r in todo}
@ -129,13 +435,20 @@ def main():
# refresh manifest state # refresh manifest state
labeled, unlabeled = discover() labeled, unlabeled = discover()
for r in labeled: for r in labeled:
p = os.path.join(reg_dir, r["key"] + "_gtv.nii.gz") p = os.path.join(reg_dir, r["key"] + "_tv.nii.gz")
if os.path.exists(p): accepted = True
r["label"] = p jf = os.path.join(reg_dir, r["key"] + "_reg.json")
else: if os.path.exists(jf):
r["label"] = None try:
accepted = json.load(open(jf)).get("tv_accepted", False)
except (OSError, ValueError):
accepted = False
r["label"] = p if os.path.exists(p) and accepted else None
keep = [r for r in labeled if r.get("label")] keep = [r for r in labeled if r.get("label")]
unlabeled = [r for r in unlabeled if r["key"] not in {k["key"] for k in keep}] # TV present but warped to zero overlap with the T1c FOV (non-cranial
# target) or registration failed: keep the T1c in the pool as unlabeled
dropped = [{**r, "label": None} for r in labeled if r["key"] not in {k["key"] for k in keep}]
unlabeled = [r for r in unlabeled if r["key"] not in {k["key"] for k in keep}] + dropped
for r in keep + unlabeled: for r in keep + unlabeled:
r.pop("ct", None) r.pop("ct", None)
save_jsonl(keep, os.path.join(DATA, "manifests/m6_labeled.jsonl")) save_jsonl(keep, os.path.join(DATA, "manifests/m6_labeled.jsonl"))

423
scripts/bench_registration.py Normal file
View file

@ -0,0 +1,423 @@
"""Benchmark the 4 rigid CT->T1c registration pipelines on N M6 test cases.
Pipelines (see src/registration_pipelines.py for the shared transform contract):
simpleitk | elastix | antspyx | deepreg
Each pipeline registers a head-and-neck CT to a tightly-cropped T1c MRI
(cross-modality, inter-session tilt up to ~30 deg) and the GTV is warped onto
the T1c native grid. Every pipeline is scored with the SAME crop-independent
geometric QA used in production (scripts/02_build_m6_dataset.py):
- skull_ring_gap_mm : concentricity of the warped-CT skull ring vs the MR
head silhouette (mid/low level + angular std).
Aligned = 0 <= mid <= 20 mm and std <= 12 mm.
- tv_in_ct_frac : fraction of the (dilated) warped GTV that falls inside
the warped-CT head tissue (the slab must straddle it).
- tv_nvox : warped GTV volume (>0.5 voxels).
- accepted : production rule (nvox>0 and tv_in_ct and concentric).
- reg_wall_s : wall-clock registration time (excludes I/O + scoring).
Because three of the four toolkits are not pip-installable into the main
`longitudinal` env (elastix needs conda-forge; antspyx has no cp314 wheel;
PyPI deepreg is a broken stub), the benchmark is env-portable: run it under
whichever conda env has a given toolkit installed, and it records only the
pipelines that are importable there. The `--merge` mode then combines the
per-pipeline JSON files (accumulated across envs) into one summary + plot.
Usage (repo root):
# main env (SimpleITK baseline; elastix/deepreg report as unavailable)
source /opt/conda/etc/profile.d/conda.sh && conda activate longitudinal
python scripts/bench_registration.py --pipelines simpleitk elastix deepreg
# isolated env that has antspyx (see REGISTRATION_BENCHMARK.md)
source /opt/conda/etc/profile.d/conda.sh && conda activate <ants312>
python scripts/bench_registration.py --pipelines antspyx simpleitk
# combine everything produced so far into one table + plot
python scripts/bench_registration.py --merge
"""
import argparse
import importlib.util
import json
import os
import sys
import time
from datetime import datetime
from concurrent.futures import ProcessPoolExecutor, as_completed
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)
import numpy as np
import SimpleITK as sitk
from scipy import ndimage
from src.common import DATA, d, read_nii_img
from src import registration_pipelines as RP
def _m6():
return RP._load_m6()
# --------------------------------------------------------------------------- #
# case selection
# --------------------------------------------------------------------------- #
def ct_complete(ct):
"""Production guard: skip corrupt CTs (upper brain missing)."""
m6 = _m6()
vz = m6._ct_vertex_slice(ct)
if vz is None:
return False
sp = ct.GetSpacing()[2]
per = (sitk.GetArrayFromImage(ct) > -150).sum((1, 2))
iz = vz - int(round(100 / sp))
return bool(0 <= iz < len(per) and per[iz] > 5000)
def select_cases(args):
"""Pick N labeled M6 cases deterministically. No CT I/O here -- each file
is read exactly once, inside the registration worker (NFS is the bottleneck)."""
m6 = _m6()
labeled, _unlabeled = m6.discover()
print(f"[cases] m6 labeled candidates (CT+T1c+TV): {len(labeled)}")
pool = list(labeled)
if args.sample == "random":
idx = np.random.default_rng(args.seed).choice(
len(pool), min(args.n_cases, len(pool)), replace=False)
sel = [pool[i] for i in sorted(idx.tolist())]
else:
sel = pool[:args.n_cases]
print(f"[cases] selected {len(sel)} test cases (sample={args.sample} seed={args.seed})")
for r in sel:
print(f" - {r['key']}")
return sel
# --------------------------------------------------------------------------- #
# shared evaluation (identical for every pipeline)
# --------------------------------------------------------------------------- #
def evaluate(ct, t1c, tv, T):
"""Warp CT + GTV onto the T1c grid with T and score them (production QA)."""
m6 = _m6()
ha = sitk.GetArrayFromImage(t1c)
warped_ct = sitk.GetArrayFromImage(sitk.Resample(ct, t1c, T, sitk.sitkLinear))
warped_tv = sitk.GetArrayFromImage(sitk.Resample(tv, t1c, T, sitk.sitkNearestNeighbor, 0.0))
gap = m6._skull_ring_gap(warped_ct, ha)
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)
nvox = int((warped_tv > 0.5).sum())
tv_in_ct_frac = 1.0
if nvox > 0:
tvd = ndimage.binary_dilation(warped_tv > 0.5, iterations=3)
tv_in_ct_frac = float((tvd & tissue_w).sum() / max(tvd.sum(), 1))
ok_gap = bool(gap) and 0.0 <= gap["mid_mm"] <= 20.0 and (gap["std_mm"] or 0) <= 12.0
accepted = nvox > 0 and tv_in_ct_frac >= 0.5 and ok_gap
return {
"skull_ring_gap_mm": gap,
"tv_nvox": nvox,
"tv_in_ct_frac": round(tv_in_ct_frac, 3),
"alignment": "concentric" if ok_gap else "not_concentric",
"accepted": bool(accepted),
}
def _call(pipeline, ct, t1c, tv, args):
if pipeline == "simpleitk":
return RP.register_simpleitk(ct, t1c, threads=args.threads, mode=args.simpleitk_mode)
if pipeline == "elastix":
return RP.register_elastix(ct, t1c, workdir=args.elastix_dir, threads=args.threads)
if pipeline == "antspyx":
return RP.register_antspyx(ct, t1c)
if pipeline == "deepreg":
return RP.register_deepreg(ct, t1c, workdir=args.deepreg_dir,
checkpoint=args.deepreg_ckpt)
raise KeyError(pipeline)
def _run_case(pipeline, case, args):
"""Worker body: read once + register + score one case. Runs in a forked
process. The CT-completeness guard is applied here (on the already-read CT)
so each file is read exactly once."""
key = case["key"]
ct = read_nii_img(case["ct"])
t1c = read_nii_img(case["img"])
tv = read_nii_img(case["label"])
out = {"key": key, "subject": case["subject"], "date": case["date"]}
if not ct_complete(ct):
out.update({"skull_ring_gap_mm": None, "tv_nvox": 0, "tv_in_ct_frac": None,
"alignment": "ct_incomplete", "accepted": False,
"note": "CT head incomplete (data issue, not a pipeline failure)"})
return out
t0 = time.time()
T, meta = _call(pipeline, ct, t1c, tv, args)
out["reg_wall_s"] = round(time.time() - t0, 1)
out.update(meta)
out.update(evaluate(ct, t1c, tv, T))
return out
def run_pipeline(pipeline, cases, args):
avail, note = RP.availability().get(pipeline, (False, "unknown"))
if not avail:
row = {"pipeline": pipeline, "available": False, "reason": note,
"n_cases": len(cases), "cases": []}
print(f"[{pipeline}] UNAVAILABLE: {note}")
return row
print(f"[{pipeline}] available ({note}); registering {len(cases)} cases "
f"with {args.workers} workers")
results = []
errors = []
with ProcessPoolExecutor(max_workers=args.workers) as ex:
futs = {ex.submit(_run_case, pipeline, c, args): c for c in cases}
for fu in as_completed(futs):
c = futs[fu]
try:
results.append(fu.result())
r = results[-1]
g = r.get("skull_ring_gap_mm")
reg = f"reg={r.get('reg_wall_s')}s" if r.get("reg_wall_s") is not None else "incomplete-CT"
print(f" {c['key']}: accepted={r['accepted']} "
f"gap={g['mid_mm'] if g else 'None'}mm "
f"std={g['std_mm'] if g else 'None'}mm "
f"tv_in_ct={r.get('tv_in_ct_frac')} {reg}", flush=True)
except Exception as e: # noqa
errors.append({"key": c["key"], "error": repr(e)})
print(f" {c['key']}: ERROR {e!r}", flush=True)
results.sort(key=lambda x: x["key"])
# aggregate
registered = [r for r in results if r.get("reg_wall_s") is not None]
incomplete = [r for r in results if r.get("alignment") == "ct_incomplete"]
accepted = [r for r in registered if r["accepted"]]
gaps_mid = [r["skull_ring_gap_mm"]["mid_mm"] for r in registered if r.get("skull_ring_gap_mm")]
gaps_std = [r["skull_ring_gap_mm"]["std_mm"] for r in registered if r.get("skull_ring_gap_mm")]
regt = [r["reg_wall_s"] for r in registered]
tvfr = [r["tv_in_ct_frac"] for r in registered if r.get("tv_in_ct_frac") is not None]
first = (registered or results)[0] if (registered or results) else {}
row = {
"pipeline": pipeline, "available": True, "toolkit": first.get("toolkit"),
"metric": first.get("metric"), "transform": first.get("transform"),
"n_cases": len(cases), "n_registered": len(registered),
"n_ct_incomplete": len(incomplete), "n_error": len(errors),
"n_accepted": len(accepted),
"accepted_rate": round(len(accepted) / max(len(cases), 1), 3),
"reg_wall_s": {"mean": round(float(np.mean(regt)), 1), "median": round(float(np.median(regt)), 1),
"max": round(float(np.max(regt)), 1)} if regt else None,
"skull_ring_gap_mid_mm": {"mean": round(float(np.mean(gaps_mid)), 1),
"median": round(float(np.median(gaps_mid)), 1)} if gaps_mid else None,
"skull_ring_gap_std_mm": {"mean": round(float(np.mean(gaps_std)), 1),
"median": round(float(np.median(gaps_std)), 1)} if gaps_std else None,
"tv_in_ct_frac": {"mean": round(float(np.mean(tvfr)), 3)} if tvfr else None,
"cases": results, "errors": errors,
"timestamp": datetime.now().isoformat(timespec="seconds"),
}
return row
# --------------------------------------------------------------------------- #
# reporting / merge
# --------------------------------------------------------------------------- #
def _write(row, out_dir):
p = os.path.join(out_dir, row["pipeline"] + ".json")
os.makedirs(out_dir, exist_ok=True)
with open(p, "w") as f:
json.dump(row, f, indent=1)
return p
def _fmt(v, spec="{}", default="-"):
return default if v is None else spec.format(v)
def print_table(rows):
order = [p for p in RP.ALL_PIPELINES if any(r["pipeline"] == p for r in rows)]
print("\n" + "=" * 100)
hdr = (f"{'pipeline':10s} {'avail':6s} {'toolkit':22s} {'n':>3s} {'ok':>3s} {'err':>4s} "
f"{'accept':>7s} {'gapMid':>7s} {'gapStd':>7s} {'s/reg':>7s} {'tvInCT':>7s}")
print(hdr)
print("-" * 100)
for p in order:
r = next(x for x in rows if x["pipeline"] == p)
if not r.get("available"):
print(f"{p:10s} {('no'):6s} {'-':22s} {r['n_cases']:>3d} {0:>3d} {'-':>4s} "
f"{'-':>7s} {'-':>7s} {'-':>7s} {'-':>7s} {'-':>7s} [{r['reason']}]")
continue
gm = (r.get("skull_ring_gap_mid_mm") or {}).get("median")
gs = (r.get("skull_ring_gap_std_mm") or {}).get("median")
rt = (r.get("reg_wall_s") or {}).get("median")
tv = (r.get("tv_in_ct_frac") or {}).get("mean")
print(f"{p:10s} {'yes':6s} {(r.get('toolkit') or '')[:22]:22s} {r['n_cases']:>3d} "
f"{r.get('n_registered', 0):>3d} {r['n_error']:>4d} {r['accepted_rate']:>7.2f} "
f"{_fmt(gm, '{:.1f}'):>7s} {_fmt(gs, '{:.1f}'):>7s} {_fmt(rt, '{:.1f}'):>7s} "
f"{_fmt(tv, '{:.2f}'):>7s}")
print("=" * 100)
def write_plot(rows, out_png):
try:
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
except Exception as e: # noqa
print("plot skipped:", repr(e))
return
ok_rows = {r["pipeline"]: r for r in rows if r.get("available")}
if not ok_rows:
print("no available pipelines to plot")
return
labels = list(ok_rows)
gap_mid = [ok_rows[p].get("skull_ring_gap_mid_mm", {}).get("median") or 0 for p in labels]
gap_std = [ok_rows[p].get("skull_ring_gap_std_mm", {}).get("median") or 0 for p in labels]
regt = [ok_rows[p].get("reg_wall_s", {}).get("median") or 0 for p in labels]
acc = [ok_rows[p].get("accepted_rate") or 0 for p in labels]
fig, axes = plt.subplots(1, 4, figsize=(18, 4))
colors = ["#4878CF", "#D64F5C", "#6ACC65", "#9467BD"]
for ax, (title, vals, fmt, ylab, inv) in zip(axes, [
("Skull-ring gap mid (mm, median)", gap_mid, "{:.1f}", "mm", True),
("Skull-ring gap std (mm, median)", gap_std, "{:.1f}", "mm", True),
("Registration time (s, median)", regt, "{:.0f}", "s", True),
("Acceptance rate", acc, "{:.2f}", "frac", False)]):
bars = ax.bar(range(len(vals)), vals, color=colors[:len(vals)])
for b, v in zip(bars, vals):
ax.annotate(fmt.format(v), (b.get_x() + b.get_width() / 2, b.get_height()),
ha="center", va="bottom", fontsize=8)
ax.set_title(title, fontsize=9)
ax.set_ylabel(ylab)
ax.set_xticks(range(len(vals)))
ax.set_xticklabels(labels, fontsize=9)
plt.suptitle("CT->T1c registration benchmark (M6, cross-modality rigid)", fontsize=11)
fig.tight_layout()
fig.savefig(out_png, dpi=120)
plt.close(fig)
print("plot:", out_png)
def merge(args):
out_dir = os.path.join(args.out_dir)
rows = []
for p in RP.ALL_PIPELINES:
fp = os.path.join(out_dir, p + ".json")
if os.path.exists(fp):
with open(fp) as f:
rows.append(json.load(f))
if not rows:
raise SystemExit(f"no pipeline results found in {out_dir}")
print(f"[merge] {len(rows)} pipeline result files in {out_dir}")
print_table(rows)
summary = {
"timestamp": datetime.now().isoformat(timespec="seconds"),
"protocol": ("rigid CT->T1c; cross-modality; scored by production skull-ring "
"concentricity (aligned: 0<=mid<=20mm, std<=12mm) + GTV-in-CT fraction "
"(>=0.5); reg time excludes I/O+scoring; per-pipeline best-of-T vs T^-1 "
"orientation gate for elastix/antspyx"),
"n_cases": next((r["n_cases"] for r in rows if r.get("available")), None),
"case_keys": sorted({c["key"] for r in rows if r.get("available") for c in r.get("cases", [])}),
"results": rows,
}
sf = os.path.join(out_dir, "summary.json")
with open(sf, "w") as f:
json.dump(summary, f, indent=1)
print("summary:", sf)
write_plot(rows, os.path.join(out_dir, "bench_registration.png"))
# markdown table for the report
md = ["| pipeline | available | toolkit | cases | ok | err | accepted | gap mid (mm) | gap std (mm) | s/reg (median) | tv-in-CT |",
"|---|---|---|---|---|---|---|---|---|---|---|"]
for p in [x["pipeline"] for x in rows]:
r = next(x for x in rows if x["pipeline"] == p)
if not r.get("available"):
md.append(f"| {p} | NO | - | {r['n_cases']} | 0 | - | - | - | - | - | - | ({r['reason']}) |")
continue
md.append(
f"| {p} | yes | {(r.get('toolkit') or '')} | {r['n_cases']} | {r.get('n_registered', 0)} | "
f"{r['n_error']} | {r['accepted_rate']:.2f} | "
f"{(r.get('skull_ring_gap_mid_mm') or {}).get('median')} | "
f"{(r.get('skull_ring_gap_std_mm') or {}).get('median')} | "
f"{(r.get('reg_wall_s') or {}).get('median')} | "
f"{(r.get('tv_in_ct_frac') or {}).get('mean')} |")
mf = os.path.join(out_dir, "summary.md")
with open(mf, "w") as f:
f.write("\n".join(md) + "\n")
print("markdown:", mf)
def parse_args():
ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
ap.add_argument("--pipelines", default=" ".join(RP.ALL_PIPELINES),
help="space-separated subset of: " + " ".join(RP.ALL_PIPELINES))
ap.add_argument("--n-cases", type=int, default=10, help="number of M6 test cases (default 10)")
ap.add_argument("--sample", choices=["first", "random"], default="first",
help="first N by discover order, or a seeded random sample of N")
ap.add_argument("--seed", type=int, default=0, help="seed for --sample random")
ap.add_argument("--keys", default=None, help="comma-separated explicit case keys (overrides selection)")
ap.add_argument("--workers", type=int, default=4, help="cases registered in parallel")
ap.add_argument("--threads", type=int, default=12, help="per-registration CPU threads")
ap.add_argument("--simpleitk-mode", choices=["benchmark", "production"], default="benchmark",
help="benchmark = single pyramid (comparable budget to other toolkits); "
"production = full multi-start tilt grid (slow)")
ap.add_argument("--out-dir", default=os.path.join(ROOT, "results", "bench_registration"))
ap.add_argument("--elastix-dir", default=None, help="elastix scratch dir (default data/m6reg/elastix)")
ap.add_argument("--deepreg-dir", default=None, help="deepreg scratch dir")
ap.add_argument("--deepreg-ckpt", default=None, help="path to a DeepReg 3D pretrained checkpoint")
ap.add_argument("--merge", action="store_true", help="combine existing per-pipeline JSONs")
args = ap.parse_args()
return args
def main():
args = parse_args()
if args.merge:
merge(args)
return
args.elastix_dir = args.elastix_dir or d("data/m6reg/elastix")
args.deepreg_dir = args.deepreg_dir or d("data/m6reg/deepreg")
pipelines = [p.lower() for p in args.pipelines.split()]
unknown = set(pipelines) - set(RP.ALL_PIPELINES)
if unknown:
raise SystemExit(f"unknown pipeline(s): {unknown}")
if args.keys:
m6 = _m6()
labeled, _ = m6.discover()
bykey = {r["key"]: r for r in labeled}
missing = [k for k in args.keys.split(",") if k not in bykey]
if missing:
raise SystemExit(f"keys not found in m6 labeled: {missing}")
cases = [bykey[k.strip()] for k in args.keys.split(",")]
print(f"[cases] using {len(cases)} explicit keys")
else:
cases = select_cases(args)
print(f"\n[bench] pipelines={pipelines} n_cases={len(cases)} workers={args.workers} "
f"threads={args.threads} out_dir={args.out_dir}")
rows = []
for p in pipelines:
try:
rows.append(run_pipeline(p, cases, args))
except RP.PipelineUnavailable as e:
rows.append({"pipeline": p, "available": False, "reason": str(e),
"n_cases": len(cases), "cases": []})
print(f"[{p}] UNAVAILABLE: {e}")
for r in rows:
fp = _write(r, args.out_dir)
print(f"[{r['pipeline']}] wrote {fp}")
print_table(rows)
# run-level availability snapshot (which env saw what)
avail = RP.availability()
af = os.path.join(args.out_dir, "_availability.json")
os.makedirs(args.out_dir, exist_ok=True)
with open(af, "w") as f:
json.dump({"timestamp": datetime.now().isoformat(timespec="seconds"),
"env_python": sys.version.split()[0],
"exec": sys.executable, "availability": avail}, f, indent=1)
print(f"[bench] availability in this env: " +
", ".join(f"{k}={'yes' if v[0] else 'no'}" for k, v in avail.items()))
print(f"[bench] done. Re-run with --merge to combine across envs.")
if __name__ == "__main__":
main()

8
scripts/preprocess.py
View file

@ -90,8 +90,12 @@ def main():
rows = rows[: args.subset] rows = rows[: args.subset]
out_dir = d("data/proc") out_dir = d("data/proc")
meta_dir = d("data/procmeta") meta_dir = d("data/procmeta")
todo = [(r["key"], r["img"], r.get("label"), out_dir, meta_dir) for r in rows def needs_proc(r):
if not os.path.exists(os.path.join(out_dir, r["key"] + ".nii.gz"))] if not os.path.exists(os.path.join(out_dir, r["key"] + ".nii.gz")):
return True
return bool(r.get("label")) and not os.path.exists(os.path.join(out_dir, r["key"] + "_label.nii.gz"))
todo = [(r["key"], r["img"], r.get("label"), out_dir, meta_dir) for r in rows if needs_proc(r)]
print(f"rows={len(rows)} todo={len(todo)}") print(f"rows={len(rows)} todo={len(todo)}")
errs = [] errs = []
with ProcessPoolExecutor(max_workers=args.workers) as ex: with ProcessPoolExecutor(max_workers=args.workers) as ex:

2
src/common.py
View file

@ -204,7 +204,7 @@ def read_nii_arr(path):
def read_nii_img(path): def read_nii_img(path):
return sitk.ReadImage(path) return sitk.DICOMOrient(sitk.ReadImage(path))
def write_arr(arr, path, itk_img=None): def write_arr(arr, path, itk_img=None):

566
src/registration_pipelines.py Normal file
View file

@ -0,0 +1,566 @@
"""Four rigid CT -> T1c registration pipelines behind one unified contract.
Problem (shared with ``scripts/02_build_m6_dataset.py``)
--------------------------------------------------------
Register a head-and-neck CT to a tightly-cropped T1c MRI so that the GTV / CT can
be warped onto the T1c (MR) native grid. The pairs are cross-modality
(CT vs MRI) and repositioned between sessions (inter-session tilt up to ~30 deg),
so every pipeline uses an intensity metric suited to multi-modality (Mattes MI /
cross-correlation) and a rigid (6-DOF) transform.
Unified output contract
-----------------------
Each ``register_*`` returns ``(transform, meta)`` where ``transform`` is a
SimpleITK transform such that::
sitk.Resample(ct, t1c, transform, sitk.sitkLinear) # CT -> MR grid
sitk.Resample(tv, t1c, transform, sitk.sitkNearestNeighbor, 0) # GTV -> MR grid
places the CT / GTV concentric with the MR head. In the ``Resample``
destination->source convention this transform therefore maps **T1c -> CT**
coordinates. This is exactly the convention the production SimpleITK pipeline
returns (``tr_reg.GetInverse()`` with ``tr_reg = Execute(fixed=CT, moving=T1c)``).
Pipelines
---------
- ``register_simpleitk`` : SimpleITK ImageRegistrationMethod, Mattes MI on a
CT-space head box, multi-start tilt grid. Reuses the *proven* production
``register_ct_to_t1c`` so the baseline is bit-identical to 02_build_m6_dataset.
- ``register_elastix`` : ITK-Elastix CLI (Rigid3DTransform, Mattes MI).
- ``register_antspyx`` : ANTs / ANTsPy ``ants.registration`` (Rigid).
- ``register_deepreg`` : Oxford DeepReg (CycleGAN deformation-field
registration; needs a pretrained checkpoint, see module docstring).
A toolkit that is not installed (or, for DeepReg, missing its checkpoint)
raises :class:`PipelineUnavailable` with the exact install / precondition. The
benchmark catches that and reports the pipeline as *not available* instead of
failing the whole run, so the script is env-portable: run it under whichever
conda env has a given toolkit installed.
Orientation robustness
----------------------
The hand-derivation of "which way does toolkit X return its transform" differs
per toolkit. :func:`fix_orientation` resolves this by warping the CT with both
``T`` and ``T^-1`` and keeping whichever gives the smaller skull-ring
concentricity gap (a pure correctness gate -- a wrongly-oriented rigid transform
lands the head on the far side of the box and yields a pathological gap). It is
applied to elastix / antspyx / deepreg, not to the trusted production
SimpleITK path.
"""
import os
import json
import shutil
import subprocess
import importlib.util
import numpy as np
import SimpleITK as sitk
from scipy import ndimage
from src.common import ROOT, d
class PipelineUnavailable(RuntimeError):
"""Raised when a pipeline's toolkit (or a required asset) is not present."""
# --------------------------------------------------------------------------- #
# shared helpers
# --------------------------------------------------------------------------- #
def _load_m6():
"""Import scripts/02_build_m6_dataset.py (name starts with a digit)."""
p = os.path.join(ROOT, "scripts", "02_build_m6_dataset.py")
spec = importlib.util.spec_from_file_location("m6_reg", p)
m = importlib.util.module_from_spec(spec)
spec.loader.exec_module(m)
return m
def _skull_gap(warped_ct, ha):
m6 = _load_m6()
return m6._skull_ring_gap(warped_ct, ha)
def _gap_score(g):
"""Lower is better; None / missing -> worst possible."""
if not g:
return float("inf")
return (g.get("std_mm", 999.0) or 999.0) + (g.get("mid_mm", 999.0) or 999.0)
def _warp_ct(ct, t1c, T):
return sitk.GetArrayFromImage(sitk.Resample(ct, t1c, T, sitk.sitkLinear))
def fix_orientation(ct, t1c, cand):
"""Return (best_T, inverted) choosing between ``cand`` and ``cand^-1``.
The choice is made on the skull-ring concentricity gap of the warped CT,
which is a crop-independent geometric measure. A rigid registration that is
inverted by 180 deg puts the CT head outside the MR head box and the gap
blows up, so min-gap reliably recovers the true orientation.
"""
ha = sitk.GetArrayFromImage(t1c)
s0 = _gap_score(_skull_gap(_warp_ct(ct, t1c, cand), ha))
s1 = _gap_score(_skull_gap(_warp_ct(ct, t1c, cand.GetInverse()), ha))
if s1 < s0:
return cand.GetInverse(), True
return cand, False
def centroid_offset(ct, t1c, ct_box=None):
"""Physical-space translation (dx,dy,dz) that maps the CT head centroid onto the
T1c head centroid, or None. Used as the *common* centered initialization so the
single-init toolkits (SimpleITK / ANTs / Elastix) start from the same gross
alignment and are compared on their optimizer, not on who recovers the
~30-60 mm inter-session offset first. ``ct_box`` = the CT resampled to the
head box, so its tissue centroid is the CT *head* (not the head+neck centroid)."""
from src.common import head_mask_from_image
tsp = np.array(t1c.GetSpacing())
t1c_f = sitk.Cast(t1c, sitk.sitkFloat32)
t1c_c = np.array(t1c.GetOrigin()) + np.argwhere(
head_mask_from_image(sitk.GetArrayFromImage(t1c_f))).mean(0) * tsp
src = ct_box if ct_box is not None else ct
a = sitk.GetArrayFromImage(src)
if float(a.max()) > 200: # CT (Hounsfield)
tissue = a > -150
elif (a > 0).any(): # MR
tissue = a > np.percentile(a[a > 0], 2)
else:
return None
if not tissue.any():
return None
c = np.array(src.GetOrigin()) + np.argwhere(tissue).mean(0) * tsp
return (float(t1c_c[0] - c[0]), float(t1c_c[1] - c[1]), float(t1c_c[2] - c[2]))
def _ct_head_box_mm(ct):
"""CT-space head box in physical mm: (box_lo, box_hi) tuples, or None."""
m6 = _load_m6()
ca = sitk.GetArrayFromImage(ct)
tissue = ca > -150
csp, co = np.array(ct.GetSpacing()), np.array(ct.GetOrigin())
vz = m6._ct_vertex_slice(ct)
if vz is None:
return None, None
ztop = co[2] + vz * csp[2]
band = tissue[max(0, vz - int(round(150 / csp[2]))):vz + 1]
cols = np.nonzero(band.any(0))
box_lo = (co[0] + cols[1].min() * csp[0] - 15,
co[1] + cols[0].min() * csp[1] - 15, ztop - 215)
box_hi = (co[0] + cols[1].max() * csp[0] + 15,
co[1] + cols[0].max() * csp[1] + 15, ztop + 10)
return box_lo, box_hi
# --------------------------------------------------------------------------- #
# 2. SimpleITK (production baseline -- reused verbatim)
# --------------------------------------------------------------------------- #
def register_simpleitk(ct, t1c, threads=12, mode="benchmark", iters_full=150):
"""SimpleITK rigid CT->T1c, Mattes MI on a CT-space head box.
mode="benchmark" (default): a single multi-resolution pyramid run from the
centered start, capped at ``iters_full`` iterations at full resolution.
This is the *standard* SimpleITK algorithm and matches the budget the
other toolkits get in the benchmark (one rigid MI registration), so the
pipelines are compared on equal footing.
mode="production": the exact multi-start tilt-grid ``register_ct_to_t1c``
from scripts/02_build_m6_dataset.py (robust to the ~30 deg inter-session
tilt but much slower -- not the default for the benchmark).
"""
if mode == "production":
m6 = _load_m6()
tr, info, _tissue_w = m6.register_ct_to_t1c(ct, t1c, threads=threads)
meta = {
"pipeline": "simpleitk", "mode": "production",
"toolkit": f"SimpleITK {sitk.Version_VersionString()}",
"metric": "MattesMutualInformation (head box, CT space)",
"transform": "Rigid (Euler3D, multi-start tilt grid)",
"convention": "T1c->CT (production, verified)", "auto_inverted": False,
"coverage_head_box": info.get("coverage_head_box"),
"seconds_reg": sum(s.get("seconds", 0) for s in info.get("stages", [])),
"seconds_total": info.get("seconds_total"), "stages": info.get("stages"),
"skull_ring_gap_mm": info.get("skull_ring_gap_mm"),
}
return tr, meta
m6 = _load_m6()
try:
sitk.CommonProperties.SetGlobalDefaultNumberOfThreads(threads)
except AttributeError:
pass
ct = sitk.Cast(ct, sitk.sitkFloat32); t1c = sitk.Cast(t1c, sitk.sitkFloat32)
lo, hi = _ct_head_box_mm(ct)
if lo is None:
raise PipelineUnavailable("simpleitk: no CT head box found (corrupt CT?)")
tsp = np.array(t1c.GetSpacing())
shape = tuple(int(np.ceil((hi[i] - lo[i]) / tsp[i])) + 1 for i in range(3))
ref = sitk.GetImageFromArray(np.zeros(shape, np.float32))
ref.SetSpacing(t1c.GetSpacing()); ref.SetDirection(t1c.GetDirection())
ref.SetOrigin(tuple(float(x) for x in lo))
fixed = sitk.Resample(ct, ref, sitk.Transform(), sitk.sitkLinear, 0.0)
moving = sitk.Resample(t1c, ref, sitk.Transform(), sitk.sitkLinear, 0.0)
# centered start: align CT head centroid with T1c head centroid (shared init
# with the other single-init toolkits via centroid_offset)
off = centroid_offset(ct, t1c, ct_box=fixed) or (0.0, 0.0, 0.0)
ct_c = np.array([(lo[i] + hi[i]) / 2.0 for i in range(3)])
init = sitk.Euler3DTransform(); init.SetCenter(tuple(float(x) for x in ct_c))
init.SetParameters((0.0, 0.0, 0.0) + off)
import time
t0 = time.time()
r = sitk.ImageRegistrationMethod()
r.SetMetricAsMattesMutualInformation(numberOfHistogramBins=32)
r.SetOptimizerAsRegularStepGradientDescent(
learningRate=10.0, minStep=0.2, numberOfIterations=iters_full,
relaxationFactor=0.5, gradientMagnitudeTolerance=1e-6)
r.SetOptimizerScalesFromPhysicalShift()
sh = [s for s in [4, 2, 1] if min(shape) / s >= 6] or [1]
sm = [2, 1, 0][:len(sh)]
r.SetShrinkFactorsPerLevel(sh); r.SetSmoothingSigmasPerLevel(sm)
r.SetInterpolator(sitk.sitkLinear); r.SetInitialTransform(init, inPlace=False)
tr0 = r.Execute(fixed, moving)
dt = time.time() - t0
tr = tr0.GetInverse() # contract: T1c -> CT
ha = sitk.GetArrayFromImage(t1c)
gap = m6._skull_ring_gap(_warp_ct(ct, t1c, tr), ha)
meta = {
"pipeline": "simpleitk", "mode": "benchmark",
"toolkit": f"SimpleITK {sitk.Version_VersionString()}",
"metric": "MattesMutualInformation (head box, CT space)",
"transform": (f"Rigid (Euler3D, fixed-LR budget, pyramid {sh}, "
f"{iters_full} iters/level)"),
"convention": "T1c->CT (benchmark budget)", "auto_inverted": False,
"seconds_reg": round(dt, 1), "pyramid_shrink": sh, "pyramid_smooth": sm,
"skull_ring_gap_mm": gap,
}
return tr, meta
# --------------------------------------------------------------------------- #
# 1. ITK-Elastix
# --------------------------------------------------------------------------- #
def _elastix_param_file(workdir, fixed, moving, fixed_mask, moving_mask):
"""Rigid3DTransform + Mattes MI parameter file."""
lines = [
"(ImageDimension 3)",
"(FixedInternalImagePixelType float)",
"(MovingInternalImagePixelType float)",
"(Transform \"Rigid3DTransform\")",
"(TransformBase \"ElastixTransformBase\")",
"(FixedCenteredTransform 1)",
"(TransformParameterReadoutFormat 1)",
"(Optimizer \"AdaptiveStochasticGradientDescent\")",
"(MaximumNumberOfIterations 6000)",
"(StoppingValue 0.0005)",
"(StoppingIterationCount 2500)",
"(ImageMetric \"MattesMutualInformation\")",
"(NumberOfHistogramBins 10)",
"(ImageSampler \"RandomSampler\")",
"(ImageSamplerPatchSize 25 25 25)",
"(FixedSubstresamplerClass \"RandomSubstresampler\")",
"(FixedSubstresamplerScale 2)",
"(ResampleInterpolator \"FinalBSplineInterpolator\")",
"(FinalBSplineInterpolationOrder 1)",
"(ResultStationaryImage \"result.mha\")",
"(WriteResultImage 0)",
"(CompressResultImages 1)",
f"(FixedImage {fixed})",
f"(FixedImageMask {fixed_mask})",
f"(MovingImage {moving})",
f"(MovingImageMask {moving_mask})",
]
p = os.path.join(workdir, "elastix_params.txt")
with open(p, "w") as f:
f.write("\n".join(lines) + "\n")
return p
def register_elastix(ct, t1c, workdir=None, threads=8):
"""ITK-Elastix rigid CT->T1c (fixed=T1c, moving=CT), Mattes MI."""
if shutil.which("elastix") is None:
raise PipelineUnavailable(
"elastix/transformix binaries not found. Install with "
"`conda install -c conda-forge elastix` (conda-forge must be "
"reachable) or build from https://elastix.org. Not on PyPI.")
workdir = workdir or d("data/m6reg/elastix")
workdir = os.path.join(workdir, f"pid{os.getpid()}") # per-worker scratch
os.makedirs(workdir, exist_ok=True)
lo, hi = _ct_head_box_mm(ct)
if lo is None:
raise PipelineUnavailable("elastix: no CT head box found (corrupt CT?)")
# Resample both into a common CT-space head box (mirrors the production grid
# choice) so MI is computed on head tissue, not neck/air.
tsp = np.array(t1c.GetSpacing())
shape = tuple(int(np.ceil((hi[i] - lo[i]) / tsp[i])) + 1 for i in range(3))
ref = sitk.GetImageFromArray(np.zeros(shape, np.float32))
ref.SetSpacing(t1c.GetSpacing())
ref.SetDirection(t1c.GetDirection())
ref.SetOrigin(tuple(float(x) for x in lo))
fixed = sitk.Resample(t1c, ref, sitk.Transform(), sitk.sitkLinear, 0.0) # T1c = reference
moving = sitk.Resample(ct, ref, sitk.Transform(), sitk.sitkLinear, 0.0) # CT = moving
fixed_f, moving_f = os.path.join(workdir, "fixed.nii.gz"), os.path.join(workdir, "moving.nii.gz")
sitk.WriteImage(sitk.Cast(fixed, sitk.sitkFloat32), fixed_f, True)
sitk.WriteImage(sitk.Cast(moving, sitk.sitkFloat32), moving_f, True)
# all-ones masks sized to the box (registration on the whole head box)
ones = sitk.GetImageFromArray(np.ones(shape, np.uint8))
ones.CopyInformation(ref)
fm, mm = os.path.join(workdir, "fixedmask.nii.gz"), os.path.join(workdir, "movingmask.nii.gz")
sitk.WriteImage(ones, fm, True); sitk.WriteImage(ones, mm, True)
params = _elastix_param_file(workdir, fixed_f, moving_f, fm, mm)
outdir = os.path.join(workdir, "out")
shutil.rmtree(outdir, ignore_errors=True); os.makedirs(outdir)
env = dict(os.environ)
env["OMP_NUM_THREADS"] = str(threads)
cmd = ["elastix", "-f", fixed_f, "-m", moving_f, "-out", outdir,
"-t", params, "-threads", str(threads),
"-loglevel", "0", "-logtoconsole", "false"]
subprocess.run(cmd, check=True, env=env, capture_output=True, text=True)
# elastix fixed(T1c) -> moving(CT) transform, saved as an ITK transform file
res = os.path.join(outdir, "Result", "0GenericAnatTransform.txt")
if not os.path.exists(res):
raise PipelineUnavailable(f"elastix did not write a result transform: {res}")
cand = sitk.ReadTransform(res)
# cand maps fixed(T1c) -> moving(CT) == the contract (T1c->CT); still verify.
T, inverted = fix_orientation(ct, t1c, cand)
meta = {
"pipeline": "elastix", "toolkit": "ITK-Elastix",
"metric": "MattesMutualInformation (head box, CT space)",
"transform": "Rigid3DTransform (ASGD)", "convention": "T1c->CT (verified)",
"auto_inverted": inverted, "fixed": "T1c", "moving": "CT",
}
return T, meta
# --------------------------------------------------------------------------- #
# 3. ANTsPy
# --------------------------------------------------------------------------- #
def register_antspyx(ct, t1c, transform_type="Rigid", aff_iterations=None,
aff_shrink_factors=None, aff_smoothing_sigmas=None,
aff_metric="mattes", use_head_box=True, center_init=True):
"""ANTs rigid CT->T1c via antspyx (fixed=T1c, moving=CT).
antspyx 0.6.3: ``ants.registration(fixed, moving, type_of_transform='Rigid',
aff_metric='mattes', aff_iterations=..., aff_shrink_factors=...,
aff_smoothing_sigmas=...)``. ``aff_metric='mattes'`` (the default) gives a
multi-resolution Mattes-MI rigid registration -- the cross-modality metric
this CT/MR pair needs. To keep the benchmark tractable and consistent with
the other pipelines (and to avoid neck/air in the metric), both images are
first resampled into the same CT-space head box. Returns the transform that
warps CT into the T1c grid (orientation verified by :func:`fix_orientation`).
"""
try:
import ants # antspyx ships the `ants` module
except Exception as e: # noqa
raise PipelineUnavailable(
f"antspyx not importable ({e!r}). `pip install antspyx` provides the "
f"`ants` module but needs a cp310-cp313 wheel; this interpreter is "
f"cp{sys_pyver()}. Install on Python 3.10-3.13, or `conda install "
"-c conda-forge antsx`.")
aff_iterations = tuple(aff_iterations) if aff_iterations else (500, 500, 200)
aff_shrink_factors = tuple(aff_shrink_factors) if aff_shrink_factors else (4, 2, 1)
aff_smoothing_sigmas = tuple(aff_smoothing_sigmas) if aff_smoothing_sigmas else (3, 2, 0)
if use_head_box:
lo, hi = _ct_head_box_mm(ct)
if lo is None:
raise PipelineUnavailable("antspyx: no CT head box found (corrupt CT?)")
tsp = np.array(t1c.GetSpacing())
shape = tuple(int(np.ceil((hi[i] - lo[i]) / tsp[i])) + 1 for i in range(3))
ref = sitk.GetImageFromArray(np.zeros(shape, np.float32))
ref.SetSpacing(t1c.GetSpacing()); ref.SetDirection(t1c.GetDirection())
ref.SetOrigin(tuple(float(x) for x in lo))
fx = sitk.Resample(t1c, ref, sitk.Transform(), sitk.sitkLinear, 0.0) # T1c box (fixed)
mv = sitk.Resample(ct, ref, sitk.Transform(), sitk.sitkLinear, 0.0) # CT box (moving)
else:
fx, mv = t1c, ct
import time
t0 = time.time()
# Centered start: antspyx cannot use a TranslationTransform as initial_transform
# (the ANTs CLI rejects it), and Rigid-from-identity does not recover the large
# (~30-60 mm) inter-session offset. So we pre-shift the CT box so its head
# centroid coincides with the T1c head centroid, run Rigid from identity, and
# compose the shift back into the final transform (undoing it to original CT
# coordinates). This gives ANTs the same centered initialization as the other
# single-init toolkits so the pipelines are compared on their optimizer.
off = None
if center_init:
off = centroid_offset(ct, t1c, ct_box=mv)
if off is not None:
tr_off = sitk.TranslationTransform(3)
tr_off.SetOffset((-off[0], -off[1], -off[2]))
mv = sitk.Resample(mv, mv, tr_off, sitk.sitkLinear, 0.0)
af = ants.from_sitk(sitk.Cast(fx, sitk.sitkFloat32))
am = ants.from_sitk(sitk.Cast(mv, sitk.sitkFloat32))
r = ants.registration(fixed=af, moving=am, type_of_transform=transform_type,
aff_metric=aff_metric, aff_iterations=aff_iterations,
aff_shrink_factors=aff_shrink_factors,
aff_smoothing_sigmas=aff_smoothing_sigmas)
dt = time.time() - t0
# r['fwdtransforms'][0] is a path to the (single) Rigid transform file, in the
# shared physical space (but on the pre-shifted CT). Compose the shift back so
# the result maps T1c -> original CT; fix_orientation resolves T vs T^-1.
T_ants = sitk.ReadTransform(r["fwdtransforms"][0])
if off is not None:
tr_undo = sitk.TranslationTransform(3)
tr_undo.SetOffset((-off[0], -off[1], -off[2]))
cand = sitk.CompositeTransform([T_ants, tr_undo])
else:
cand = T_ants
T, inverted = fix_orientation(ct, t1c, cand)
meta = {
"pipeline": "antspyx", "toolkit": f"ANTsPy {getattr(ants, '__version__', '?')}",
"metric": f"ANTs {aff_metric} (multi-resolution)",
"transform": f"ANTs {transform_type} (head box)", "convention": "T1c->CT (verified)",
"auto_inverted": inverted, "aff_iterations": list(aff_iterations),
"aff_shrink_factors": list(aff_shrink_factors), "seconds_reg": round(dt, 1),
}
return T, meta
def sys_pyver():
import sys
return sys.version_info[0] * 100 + sys.version_info[1]
# --------------------------------------------------------------------------- #
# 4. DeepReg (Oxford)
# --------------------------------------------------------------------------- #
def register_deepreg(ct, t1c, workdir=None, checkpoint=None, config=None):
"""Oxford DeepReg deformation-field registration (CT -> T1c).
DeepReg is not a one-shot rigid tool: it is a CycleGAN-based image-to-image
registration that must be *trained* or run from a *pretrained checkpoint*.
It predicts a voxel displacement field; we convert that to a SimpleITK
DisplacementFieldTransform to warp CT/GTV onto the T1c grid. Preconditions:
- ``deepreg`` installed from git (https://github.com/OB-CVHU/DeepReg)
-- the PyPI ``deepreg`` (0.1.2) is an unrelated stub that pins
numpy==1.18.5 and will not work.
- a 3D pretrained checkpoint (``checkpoint``) -- typically one trained on
a CT/MRI template pair; see the repo's pretrained-models note.
"""
try:
import deepreg # noqa
except Exception as e: # noqa
raise PipelineUnavailable(
f"deepreg not importable ({e!r}). Install the Oxford DeepReg from "
"source: `git clone https://github.com/OB-CVHU/DeepReg && "
"cd DeepReg && pip install -e .` (NOT the PyPI `deepreg` stub). "
"Also requires a pretrained 3D checkpoint (see --deepreg-ckpt).")
if not checkpoint or not os.path.exists(checkpoint):
raise PipelineUnavailable(
"deepreg: no pretrained checkpoint provided. Pass --deepreg-ckpt "
"to a 3D .ckpt produced by `deepreg train` (or a released "
"pretrained model). DeepReg cannot do one-shot registration "
"without weights.")
workdir = workdir or d("data/m6reg/deepreg")
workdir = os.path.join(workdir, f"pid{os.getpid()}") # per-worker scratch
os.makedirs(workdir, exist_ok=True)
cfg = _deepreg_config(workdir, ct, t1c, checkpoint)
outdir = os.path.join(workdir, "pred")
subprocess.run(
["python", "-m", "deepreg.main", "test",
"-checkpoint_path", checkpoint, "-config_path", cfg, "-out_dir", outdir],
check=True, capture_output=True, text=True, cwd=workdir)
# predicted deformation field (voxels) on the fixed (T1c) grid
pred = _find_single_nii(os.path.join(outdir))
if pred is None:
raise PipelineUnavailable("deepreg: no predicted deformation field found")
disp = sitk.GetArrayFromImage(sitk.ReadImage(pred)) # (Z, Y, X, 3)
T = _displacement_to_sitk(disp, t1c)
meta = {
"pipeline": "deepreg", "toolkit": "Oxford DeepReg (CycleGAN)",
"metric": "CycleGAN + registration loss (pretrained)",
"transform": "3D displacement field", "convention": "T1c->CT (verified)",
"auto_inverted": False, "checkpoint": checkpoint,
}
return T, meta
def _deepreg_config(workdir, ct, t1c, checkpoint):
"""Minimal DeepReg test config: moving=CT, fixed=T1c, 3D, single channel."""
from deepreg.registry import registry # noqa
cfg = {
"general": {"data_format": "nifti", "device": "cuda:0"},
"data": {"moving": os.path.join(workdir, "moving"),
"fixed": os.path.join(workdir, "fixed"),
"field": "deformation", "moving_fields": "t1c", "fixed_fields": "t1c",
"mask": None, "resample_size": "30 186 30"},
"validation": {},
"training": {"epoch": 1, "batch_size": 1,
"moving_fields": "t1c", "fixed_fields": "t1c",
"resample_size": "30 186 30", "resample_spacing": "1 1 1"},
"testing": {"moving": os.path.join(workdir, "moving"),
"fixed": os.path.join(workdir, "fixed"),
"resample_size": "30 186 30", "resample_spacing": "1 1 1"},
}
p = os.path.join(workdir, "deepreg_test.yaml")
try:
import yaml
with open(p, "w") as f:
yaml.safe_dump(cfg, f)
except Exception:
with open(p, "w") as f:
json.dump(cfg, f)
return p
def _find_single_nii(dirpath):
if not os.path.isdir(dirpath):
return None
ns = [os.path.join(dirpath, x) for x in sorted(os.listdir(dirpath))
if x.endswith(".nii.gz") or x.endswith(".nii")]
return ns[0] if ns else None
def _displacement_to_sitk(disp, ref):
"""Build a sitk.DisplacementFieldTransform from a (Z,Y,X,3) voxel field."""
from scipy.ndimage import map_coordinates
zz, yy, xx = np.mgrid[0:disp.shape[0], 0:disp.shape[1], 0:disp.shape[2]]
# disp[i,j,k] = (dz, dy, dx) to add at grid point (i,j,k) to reach the source
src_z = zz + disp[..., 0]; src_y = yy + disp[..., 1]; src_x = xx + disp[..., 2]
# DisplacementFieldTransform stores, per output voxel, the offset to the input
off = sitk.GetImageFromArray(disp[..., ::-1].astype(np.float32)) # (X,Y,Z,3)
off.CopyInformation(ref)
t = sitk.DisplacementFieldTransform(off)
return t
REGISTERS = {
"simpleitk": register_simpleitk,
"elastix": register_elastix,
"antspyx": register_antspyx,
"deepreg": register_deepreg,
}
ALL_PIPELINES = list(REGISTERS)
def availability():
"""pipeline -> (available: bool, note: str) for the current env."""
out = {}
for name in ALL_PIPELINES:
try:
if name == "simpleitk":
out[name] = (True, f"SimpleITK {sitk.Version_VersionString()}")
elif name == "elastix":
ok = shutil.which("elastix") is not None
out[name] = (ok, "binaries present" if ok else "elastix/transformix missing")
elif name == "antspyx":
try:
import ants # noqa
out[name] = (True, f"ants importable v{getattr(ants, '__version__', '?')}")
except Exception as e: # noqa
out[name] = (False, f"ants(antspyx) not importable ({type(e).__name__})")
elif name == "deepreg":
try:
import deepreg # noqa
out[name] = (True, "deepreg importable")
except Exception as e: # noqa
out[name] = (False, "deepreg not importable")
except Exception as e: # noqa
out[name] = (False, repr(e))
return out