601 lines
No EOL
26 KiB
Python
601 lines
No EOL
26 KiB
Python
"""Investigate the low performance of the ``synthmorph`` pipeline on the M6
|
|
CT->T1c rigid benchmark (see REGISTRATION_BENCHMARK_RIGID.md).
|
|
|
|
The benchmark median MI is 0.277 for synthmorph vs ~0.65 for the intensity
|
|
tools, with the ~30 deg repositioned cases collapsing to 0.06-0.11. This
|
|
script decomposes *why*, scoring everything with the exact benchmark scoring
|
|
(``scripts/bench_registration.py::evaluate``):
|
|
|
|
A. Transform decomposition -- for each case, run synthmorph ``rigid`` (the
|
|
benchmark config; deterministic, so it reproduces the bench JSON) and a
|
|
``simpleitk`` benchmark reference, and report per case:
|
|
- achieved MI / ImpactLoss / skull-ring gap for both,
|
|
- rotation angle + axis of each transform (world RAS frame),
|
|
- the *relative* rotation error angle(R_ref^T R_sm),
|
|
- the head-centre placement error |T_sm(c) - T_ref(c)| (c = T1c head
|
|
centroid),
|
|
- the identity-transform score (the "do nothing" baseline),
|
|
- a numeric rigidity check of the synthmorph matrix (R^T R ~ I, det = 1).
|
|
|
|
B. Affine model -- run synthmorph ``affine`` (12-DOF matrix model) on the
|
|
same cases and score it: does the bigger matrix family recover the tilt
|
|
that rigid does not? The linear part is decomposed into rotation /
|
|
scale / shear (polar).
|
|
|
|
C. Joint (affine + deformable) upper bound -- for the worst cases, run
|
|
``-m joint`` directly in Docker (outside the rigid contract: the output is
|
|
a displacement field, applied with the tool's own ``-o`` resampling) and
|
|
score the warped CT image. Answers: is the low score a matrix-model
|
|
limitation or a model-family (pretrained-on-MRI) limitation?
|
|
|
|
D. Network-space mechanics -- static (no registration) check of what the
|
|
network actually sees, emulating ``network_space`` + ``normalize`` from
|
|
the image code (isotropic 1-mm LIA grid, extent 256, centred on each
|
|
volume's own centre, fill 0, min-max into [0,1]):
|
|
- head truncation: fraction of head-mask voxels (the T1c head mapped into
|
|
CT space via the reference transform) outside the 256 mm box, for the CT
|
|
and T1c network spaces,
|
|
- intensity band: the min-max-normalized values over the head region for
|
|
the T1c, the [0,80]-clipped CT (the benchmark default) and the raw CT.
|
|
|
|
Usage (repo root, conda env longitudinal-313):
|
|
python scripts/investigate_synthmorph.py --workers 3
|
|
python scripts/investigate_synthmorph.py --parts D # static only
|
|
python scripts/investigate_synthmorph.py --keys m6_2BBN5XA4_20250411
|
|
"""
|
|
import argparse
|
|
import importlib.util
|
|
import json
|
|
import math
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
import time
|
|
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.ndimage import center_of_mass
|
|
|
|
from src.common import d, read_nii_img, head_mask_from_image
|
|
from src import registration_pipelines as RP
|
|
|
|
|
|
def _load_module(name, rel):
|
|
p = os.path.join(ROOT, rel)
|
|
spec = importlib.util.spec_from_file_location(name, p)
|
|
m = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(m)
|
|
return m
|
|
|
|
|
|
BR = _load_module("bench_registration", "scripts/bench_registration.py")
|
|
M6 = RP._load_m6()
|
|
|
|
# the worst 3 bench cases (rigid MI): ~30 deg tilt x2 + a mid-tier case
|
|
JOINT_KEYS = ("m6_2BBN5XA4_20250411", "m6_24V62QQM_20240820",
|
|
"m6_24V62QQM_20230223")
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# transform decomposition
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _rot_angle(R):
|
|
c = (np.trace(R) - 1.0) / 2.0
|
|
return math.degrees(math.acos(float(np.clip(c, -1.0, 1.0))))
|
|
|
|
|
|
def _rot_axis(R, theta_deg):
|
|
if theta_deg <= 0.5:
|
|
return np.zeros(3)
|
|
axis = np.array([R[2, 1] - R[1, 2], R[0, 2] - R[2, 0], R[1, 0] - R[0, 1]])
|
|
n = np.linalg.norm(axis)
|
|
return axis / n if n > 1e-9 else np.zeros(3)
|
|
|
|
|
|
def _to_affine(T):
|
|
"""Recover (R, t) of any SimpleITK transform by point mapping (world RAS)."""
|
|
o = np.array(T.TransformPoint((0.0, 0.0, 0.0)), dtype=float)
|
|
R = np.empty((3, 3), dtype=float)
|
|
for i in range(3):
|
|
p = [0.0, 0.0, 0.0]
|
|
p[i] = 1.0
|
|
R[:, i] = np.array(T.TransformPoint(tuple(p)), dtype=float) - o
|
|
return R, o
|
|
|
|
|
|
def _decompose_affine(T):
|
|
"""(R, t, theta_deg, axis, rigidity) of a SimpleITK transform in world RAS."""
|
|
R, t = _to_affine(T)
|
|
theta = _rot_angle(R)
|
|
rig = {"det": round(float(np.linalg.det(R)), 6),
|
|
"max_devi_from_I": round(float(np.linalg.norm(R.T @ R - np.eye(3))), 6)}
|
|
return R, t, theta, _rot_axis(R, theta), rig
|
|
|
|
|
|
def _angle_between_rotations(A, B):
|
|
"""angle of A^T B in degrees (rotation B needs on top of A)."""
|
|
return _rot_angle(A.T @ B)
|
|
|
|
|
|
def _polar_rotation(M):
|
|
"""rotation part of a general 3x3 linear map (SVD polar)."""
|
|
U, _s, Vt = np.linalg.svd(M)
|
|
R_p = U @ Vt
|
|
if np.linalg.det(R_p) < 0:
|
|
U[:, -1] *= -1
|
|
R_p = U @ Vt
|
|
return R_p
|
|
|
|
|
|
def head_centroid_world(t1c):
|
|
ha = sitk.GetArrayFromImage(t1c).astype("float32")
|
|
idx = center_of_mass(head_mask_from_image(ha))
|
|
sp = np.array(t1c.GetSpacing())
|
|
o = np.array(t1c.GetOrigin())
|
|
D = np.array(t1c.GetDirection()).reshape(3, 3)
|
|
return o + D @ (np.array(idx[::-1]) * sp)
|
|
|
|
|
|
def score_identity(ct, t1c):
|
|
"""The "do nothing" baseline: CT resampled onto the T1c grid with the
|
|
identity transform (i.e. no registration at all)."""
|
|
ha = sitk.GetArrayFromImage(t1c).astype("float32")
|
|
wc = sitk.GetArrayFromImage(sitk.Resample(ct, t1c, sitk.Transform(),
|
|
sitk.sitkLinear))
|
|
gap = M6._skull_ring_gap(wc, ha)
|
|
mi = BR._mi_masked(wc, ha, head_mask_from_image(ha))
|
|
return {"mi": mi, "gap": gap}
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# part A/B/D worker
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _case_task(case, args):
|
|
out = {"key": case["key"], "subject": case["subject"], "date": case["date"]}
|
|
ct = read_nii_img(case["ct"])
|
|
t1c = read_nii_img(case["img"])
|
|
tv = read_nii_img(case["label"])
|
|
if not BR.ct_complete(ct):
|
|
out["note"] = "ct_incomplete (skipped, same as the benchmark)"
|
|
return out
|
|
|
|
gpu = (args.gpu_base + abs(hash(out["key"])) % max(1, args.n_gpus)) % max(1, args.n_gpus)
|
|
parts = set(args.parts)
|
|
|
|
# geometry facts
|
|
for nm, img in (("ct", ct), ("t1c", t1c)):
|
|
sz = img.GetSize()
|
|
sp = img.GetSpacing()
|
|
c = np.array(img.TransformIndexToPhysicalPoint(
|
|
((sz[0] - 1) // 2, (sz[1] - 1) // 2, (sz[2] - 1) // 2)))
|
|
out[nm] = {"shape_xyz": list(sz),
|
|
"spacing_xyz": [round(float(s), 3) for s in sp],
|
|
"fov_mm_xyz": [round(float(s0 * s1), 1) for s0, s1 in zip(sz, sp)],
|
|
"center_world": [round(float(x), 1) for x in c]}
|
|
c_t1c = head_centroid_world(t1c)
|
|
out["t1c_head_centroid_world"] = [round(float(x), 1) for x in c_t1c]
|
|
|
|
R_ref = t_ref = None
|
|
if parts & {"A", "B", "D"}:
|
|
# reference: simpleitk benchmark mode (deterministic at fixed threads)
|
|
T_ref, meta_ref = RP.register_simpleitk(ct, t1c, threads=16,
|
|
mode="benchmark")
|
|
sref = BR.evaluate(ct, t1c, tv, T_ref)
|
|
R_ref, t_ref, th_ref, ax_ref, _ = _decompose_affine(T_ref)
|
|
out["reference"] = {
|
|
"tool": meta_ref["toolkit"],
|
|
"R": R_ref.tolist(), "t": t_ref.tolist(),
|
|
"mi": sref["mi_ct_mr"], "impact_loss": sref["impact_loss"],
|
|
"gap": sref["skull_ring_gap_mm"],
|
|
"rotation_deg": round(th_ref, 2),
|
|
"axis": [round(float(a), 3) for a in ax_ref],
|
|
"translation": [round(float(x), 1) for x in t_ref],
|
|
"head_centre_ct": [round(float(x), 1) for x in (R_ref @ c_t1c + t_ref)],
|
|
}
|
|
out["identity"] = score_identity(ct, t1c)
|
|
|
|
if "A" in parts:
|
|
t0 = time.time()
|
|
T_sm, meta_sm = RP.register_synthmorph(ct, t1c, model="rigid",
|
|
gpu=int(gpu), extent=256)
|
|
ssm = BR.evaluate(ct, t1c, tv, T_sm)
|
|
R_sm, t_sm, th_sm, ax_sm, rig = _decompose_affine(T_sm)
|
|
p_ref = R_ref @ c_t1c + t_ref
|
|
p_sm = R_sm @ c_t1c + t_sm
|
|
out["rigid"] = {
|
|
"seconds": round(time.time() - t0, 1),
|
|
"auto_inverted": meta_sm["auto_inverted"],
|
|
"clip_ct": meta_sm["clip_ct"],
|
|
"mi": ssm["mi_ct_mr"], "impact_loss": ssm["impact_loss"],
|
|
"gap": ssm["skull_ring_gap_mm"],
|
|
"rotation_deg": round(th_sm, 2),
|
|
"axis": [round(float(a), 3) for a in ax_sm],
|
|
"translation": [round(float(x), 1) for x in t_sm],
|
|
"matrix_rigidity": rig,
|
|
"rot_err_vs_ref_deg": round(_angle_between_rotations(R_ref, R_sm), 2),
|
|
"head_centre_err_mm": round(float(np.linalg.norm(p_sm - p_ref)), 1),
|
|
}
|
|
|
|
if "B" in parts:
|
|
t0 = time.time()
|
|
T_af, meta_af = RP.register_synthmorph(ct, t1c, model="affine",
|
|
gpu=int(gpu), extent=256)
|
|
saf = BR.evaluate(ct, t1c, tv, T_af)
|
|
M_af, t_af = _to_affine(T_af)
|
|
svals = np.linalg.svd(M_af, compute_uv=False)
|
|
lin = {"scale_mean": round(float(np.mean(svals)), 4),
|
|
"scale_iso_dev": round(float(np.std(svals)) / max(float(np.mean(svals)), 1e-12), 4),
|
|
"rotation_deg": round(_rot_angle(_polar_rotation(M_af)), 2),
|
|
"det": round(float(np.linalg.det(M_af)), 4)}
|
|
row = {"seconds": round(time.time() - t0, 1),
|
|
"auto_inverted": meta_af["auto_inverted"],
|
|
"mi": saf["mi_ct_mr"], "impact_loss": saf["impact_loss"],
|
|
"gap": saf["skull_ring_gap_mm"], "linear": lin}
|
|
if R_ref is not None:
|
|
row["rot_err_vs_ref_deg"] = round(
|
|
_angle_between_rotations(R_ref, _polar_rotation(M_af)), 2)
|
|
row["head_centre_err_mm"] = round(
|
|
float(np.linalg.norm(M_af @ c_t1c + t_af - (R_ref @ c_t1c + t_ref))), 1)
|
|
out["affine"] = row
|
|
|
|
if "D" in parts:
|
|
out["network_space"] = _network_space_analysis(ct, t1c, out,
|
|
R_ref=R_ref, t_ref=t_ref)
|
|
|
|
return out
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# part D: what the network sees
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _network_grid(img, extent=256):
|
|
"""The matrix-model network space of ``img``: isotropic 1-mm LIA grid,
|
|
``extent``^3, centred on the volume's own centre (mirrors
|
|
``network_space`` in the container's registration.py)."""
|
|
sz = img.GetSize()
|
|
c = np.array(img.TransformIndexToPhysicalPoint(
|
|
((sz[0] - 1) // 2, (sz[1] - 1) // 2, (sz[2] - 1) // 2)))
|
|
# LIA axes in the RAS world: x_L = -x_R, y_I = -y_A, z_A = +z_R
|
|
origin = (c[0] + extent / 2.0 - 0.5, c[1] + extent / 2.0 - 0.5,
|
|
c[2] - extent / 2.0 + 0.5)
|
|
net = sitk.Image((extent, extent, extent), sitk.sitkFloat32)
|
|
net.SetSpacing((1.0, 1.0, 1.0))
|
|
net.SetOrigin(tuple(float(x) for x in origin))
|
|
net.SetDirection((-1.0, 0.0, 0.0, 0.0, -1.0, 0.0, 0.0, 0.0, 1.0))
|
|
return net, c
|
|
|
|
|
|
def _to_network(img, net):
|
|
a = sitk.GetArrayFromImage(
|
|
sitk.Resample(img, net, sitk.Transform(), sitk.sitkLinear, 0.0))
|
|
lo, hi = float(a.min()), float(a.max())
|
|
if hi - lo < 1e-9:
|
|
return np.zeros_like(a), lo, hi
|
|
return (a - lo) / (hi - lo), lo, hi
|
|
|
|
|
|
def _network_space_analysis(ct, t1c, out, extent=256, R_ref=None, t_ref=None):
|
|
"""Extent truncation + min-max intensity band (what the model sees)."""
|
|
res = {"extent": extent}
|
|
ha = sitk.GetArrayFromImage(t1c).astype("float32")
|
|
head_t1c = sitk.GetImageFromArray(head_mask_from_image(ha).astype("float32"))
|
|
head_t1c.CopyInformation(t1c)
|
|
|
|
if R_ref is not None:
|
|
T_ref = sitk.AffineTransform(3)
|
|
T_ref.SetMatrix(list(np.array(R_ref).ravel()))
|
|
T_ref.SetTranslation(list(np.array(t_ref)))
|
|
# T_ref maps T1c -> CT, so rasterising the T1c head onto the CT grid
|
|
# needs the inverse as the destination -> source map
|
|
head_ct = sitk.Resample(head_t1c, ct, T_ref.GetInverse(),
|
|
sitk.sitkNearestNeighbor, 0.0)
|
|
else:
|
|
head_ct = head_t1c
|
|
|
|
# truncation: head voxels outside each volume's OWN extent box
|
|
for nm, img, head in (("ct", ct, head_ct), ("t1c", t1c, head_t1c)):
|
|
_net, c = _network_grid(img, extent)
|
|
idx = np.argwhere(sitk.GetArrayFromImage(head) > 0.5)
|
|
if idx.size == 0:
|
|
res[f"{nm}_head_trunc_frac"] = None
|
|
continue
|
|
sp = np.array(img.GetSpacing())
|
|
o = np.array(img.GetOrigin())
|
|
D = np.array(img.GetDirection()).reshape(3, 3)
|
|
world = o[None, :] + (idx[:, ::-1] * sp) @ D.T # idx is (z,y,x)
|
|
half = extent / 2.0
|
|
inside = (np.abs(world - c[None, :]) <= half).all(axis=1)
|
|
head_c = world.mean(axis=0)
|
|
res[f"{nm}_head_voxels"] = int(idx.shape[0])
|
|
res[f"{nm}_head_trunc_frac"] = round(float(1.0 - inside.mean()), 4)
|
|
res[f"{nm}_head_off_volume_centre_mm"] = \
|
|
[round(float(x), 1) for x in (head_c - c)]
|
|
|
|
# min-max normalized band over the head region (emulating normalize=True)
|
|
bands = {}
|
|
for nm, img, head in (("t1c", t1c, head_t1c),
|
|
("ct_clip", RP._clip_intensity(ct, 0.0, 80.0), head_ct),
|
|
("ct_raw", ct, head_ct)):
|
|
net, _c = _network_grid(img, extent)
|
|
a, lo, hi = _to_network(img, net)
|
|
hnet = sitk.GetArrayFromImage(
|
|
sitk.Resample(head, net, sitk.Transform(),
|
|
sitk.sitkNearestNeighbor, 0.0)) > 0.5
|
|
v = a[hnet]
|
|
bands[nm] = {"src_min": round(lo, 1), "src_max": round(hi, 1),
|
|
"head_nvox_net": int(v.size),
|
|
"p5": round(float(np.percentile(v, 5)), 4) if v.size else None,
|
|
"p50": round(float(np.percentile(v, 50)), 4) if v.size else None,
|
|
"p95": round(float(np.percentile(v, 95)), 4) if v.size else None,
|
|
"band_width_p95_p5": (round(float(np.percentile(v, 95)
|
|
- np.percentile(v, 5)), 4)
|
|
if v.size else None)}
|
|
res["normalized_band"] = bands
|
|
return res
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# part C: joint (affine + deformable) upper bound
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _joint_task(case, args):
|
|
out = {"key": case["key"], "subject": case["subject"], "date": case["date"],
|
|
"model": "joint (affine + deformable)"}
|
|
ct = read_nii_img(case["ct"])
|
|
t1c = read_nii_img(case["img"])
|
|
if not BR.ct_complete(ct):
|
|
out["note"] = "ct_incomplete (skipped)"
|
|
return out
|
|
gpu = (args.gpu_base + abs(hash(out["key"])) % max(1, args.n_gpus)) % max(1, args.n_gpus)
|
|
tmp = tempfile.mkdtemp(prefix="sm_joint_")
|
|
try:
|
|
ct_clip = RP._clip_intensity(ct, 0.0, 80.0)
|
|
sitk.WriteImage(sitk.Cast(ct_clip, sitk.sitkFloat32),
|
|
os.path.join(tmp, "moving.nii.gz"), True)
|
|
sitk.WriteImage(sitk.Cast(t1c, sitk.sitkFloat32),
|
|
os.path.join(tmp, "fixed.nii.gz"), True)
|
|
# CLI order is (moving, fixed); -o saves moving registered to fixed
|
|
cmd = ["docker", "run", "--rm",
|
|
"--user", f"{os.getuid()}:{os.getgid()}",
|
|
"-e", "TF_FORCE_GPU_ALLOW_GROWTH=true",
|
|
"--gpus", "all", "-e", f"CUDA_VISIBLE_DEVICES={gpu}",
|
|
"-v", f"{tmp}:/data", RP.SYNTHMORPH_IMAGE,
|
|
"register", "-m", "joint", "-e", "256", "-g",
|
|
"-o", "/data/warped.nii.gz",
|
|
"/data/moving.nii.gz", "/data/fixed.nii.gz"]
|
|
t0 = time.time()
|
|
proc = subprocess.run(cmd, capture_output=True, text=True, timeout=7200)
|
|
dt = time.time() - t0
|
|
if proc.returncode != 0:
|
|
out["error"] = (f"docker failed rc={proc.returncode}:\n"
|
|
+ "\n".join((proc.stdout + proc.stderr).splitlines()[-15:]))
|
|
return out
|
|
warped = sitk.GetArrayFromImage(
|
|
sitk.ReadImage(os.path.join(tmp, "warped.nii.gz"))).astype("float32")
|
|
ha = sitk.GetArrayFromImage(t1c).astype("float32")
|
|
gap = M6._skull_ring_gap(warped, ha)
|
|
mi = BR._mi_masked(warped, ha, head_mask_from_image(ha))
|
|
il = BR._impact_loss_value(warped, ha)
|
|
out.update({"seconds": round(dt, 1), "mi": mi, "impact_loss": il,
|
|
"gap": gap, "gpu": int(gpu)})
|
|
finally:
|
|
if not args.keep_scratch:
|
|
shutil.rmtree(tmp, ignore_errors=True)
|
|
return out
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# part F: re-centred CT probe (volume-centre anchoring hypothesis)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _recenter_task(case, args, ref_map):
|
|
"""H5 probe: the matrix model's network space is centred on each volume's
|
|
OWN centre. The T1c is a tight head crop (centre ~ head centre), but the
|
|
head+neck CT's volume centre sits 15-80 mm below the head. If the model
|
|
under-corrects that offset, cropping the CT to the head box (re-centring
|
|
the volume on the head) should remove the translation miss. The crop is
|
|
world-preserving, so the returned transform scores against the FULL CT
|
|
exactly like the benchmark."""
|
|
out = {"key": case["key"], "subject": case["subject"], "date": case["date"]}
|
|
ct = read_nii_img(case["ct"])
|
|
t1c = read_nii_img(case["img"])
|
|
tv = read_nii_img(case["label"])
|
|
if not BR.ct_complete(ct):
|
|
out["note"] = "ct_incomplete (skipped)"
|
|
return out
|
|
ref = ref_map.get(case["key"])
|
|
if not ref:
|
|
out["error"] = "no reference transform in partABD.json"
|
|
return out
|
|
R = np.array(ref["R"])
|
|
t = np.array(ref["t"])
|
|
T_ref = sitk.AffineTransform(3)
|
|
T_ref.SetMatrix(list(R.ravel()))
|
|
T_ref.SetTranslation(list(t))
|
|
|
|
gpu = (args.gpu_base + abs(hash(out["key"])) % max(1, args.n_gpus)) % max(1, args.n_gpus)
|
|
|
|
# head in CT space (T_ref maps T1c -> CT, so rasterise with its inverse)
|
|
ha = sitk.GetArrayFromImage(t1c).astype("float32")
|
|
head_t1c = sitk.GetImageFromArray(head_mask_from_image(ha).astype("float32"))
|
|
head_t1c.CopyInformation(t1c)
|
|
head_ct = sitk.Resample(head_t1c, ct, T_ref.GetInverse(),
|
|
sitk.sitkNearestNeighbor, 0.0)
|
|
m = sitk.GetArrayFromImage(head_ct) > 0.5
|
|
idx = np.argwhere(m) # (z, y, x)
|
|
lo, hi = idx.min(0), idx.max(0)
|
|
cen = (lo + hi + 1) / 2.0
|
|
half = (hi - lo) / 2.0 + 15.0 / np.array(ct.GetSpacing())[::-1] # (z,y,x)
|
|
sz = np.array(ct.GetSize())[::-1] # GetSize is (x,y,z); work in (z,y,x)
|
|
b0 = np.clip(np.floor(cen - half).astype(int), 0, sz - 1)
|
|
b1 = np.clip(np.ceil(cen + half).astype(int), 0, sz)
|
|
|
|
def _volcentre(img):
|
|
s = img.GetSize()
|
|
return np.array(img.TransformIndexToPhysicalPoint(
|
|
((s[0] - 1) // 2, (s[1] - 1) // 2, (s[2] - 1) // 2)))
|
|
|
|
hc = T_ref.TransformPoint(tuple(float(x) for x in head_centroid_world(t1c)))
|
|
out["before"] = {"head_off_volume_centre_mm": [
|
|
round(float(x), 1) for x in (np.array(hc) - _volcentre(ct))]}
|
|
|
|
# SimpleITK image slicing is in (x, y, z) order; b0/b1 are (z, y, x)
|
|
ct_crop = ct[int(b0[2]):int(b1[2]), int(b0[1]):int(b1[1]),
|
|
int(b0[0]):int(b1[0])]
|
|
out["crop"] = {"shape": list(ct_crop.GetSize()),
|
|
"head_off_volume_centre_mm": [
|
|
round(float(x), 1)
|
|
for x in (np.array(hc) - _volcentre(ct_crop))]}
|
|
|
|
t0 = time.time()
|
|
T_sm, meta_sm = RP.register_synthmorph(ct_crop, t1c, model="rigid",
|
|
gpu=int(gpu), extent=256)
|
|
ssm = BR.evaluate(ct, t1c, tv, T_sm) # full CT, as in the benchmark
|
|
R_sm, t_sm = _to_affine(T_sm)
|
|
c_t1c = head_centroid_world(t1c)
|
|
out["rigid_recentered"] = {
|
|
"seconds": round(time.time() - t0, 1),
|
|
"auto_inverted": meta_sm["auto_inverted"],
|
|
"mi": ssm["mi_ct_mr"], "impact_loss": ssm["impact_loss"],
|
|
"gap": ssm["skull_ring_gap_mm"],
|
|
"rot_err_vs_ref_deg": round(_angle_between_rotations(R, R_sm), 2),
|
|
"head_centre_err_mm": round(
|
|
float(np.linalg.norm((R_sm @ c_t1c + t_sm) - (R @ c_t1c + t))), 1),
|
|
}
|
|
return out
|
|
|
|
|
|
# --------------------------------------------------------------------------- #
|
|
# main
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def parse_args():
|
|
ap = argparse.ArgumentParser(description=__doc__,
|
|
formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
ap.add_argument("--parts", default="ABCD",
|
|
help="which parts to run (letters A/B/C/D/F; F = re-centred "
|
|
"CT probe, needs the partABD.json reference)")
|
|
ap.add_argument("--n-cases", type=int, default=10)
|
|
ap.add_argument("--keys", default=None, help="comma-separated case keys")
|
|
ap.add_argument("--workers", type=int, default=3)
|
|
ap.add_argument("--gpu-base", type=int, default=0)
|
|
ap.add_argument("--n-gpus", type=int, default=3)
|
|
ap.add_argument("--out-dir", default=os.path.join(ROOT, "results",
|
|
"synthmorph_investigation"))
|
|
ap.add_argument("--keep-scratch", action="store_true")
|
|
return ap.parse_args()
|
|
|
|
|
|
def main():
|
|
args = parse_args()
|
|
os.makedirs(args.out_dir, exist_ok=True)
|
|
labeled, _ = M6.discover()
|
|
pool = list(labeled)
|
|
if args.keys:
|
|
bykey = {r["key"]: r for r in pool}
|
|
missing = [k for k in args.keys.split(",") if k not in bykey]
|
|
if missing:
|
|
raise SystemExit(f"keys not found: {missing}")
|
|
cases = [bykey[k.strip()] for k in args.keys.split(",")]
|
|
else:
|
|
cases = pool[:args.n_cases]
|
|
print(f"[investigate] {len(cases)} cases, parts={args.parts}")
|
|
|
|
rows, joint_rows, recenter_rows = [], [], []
|
|
do_cases = bool(set(args.parts) & {"A", "B", "D"})
|
|
do_joint = "C" in args.parts
|
|
do_recenter = "F" in args.parts
|
|
|
|
ref_map = {}
|
|
if do_recenter:
|
|
fp0 = os.path.join(args.out_dir, "partABD.json")
|
|
if os.path.exists(fp0):
|
|
with open(fp0) as f:
|
|
ref_map = {r["key"]: r["reference"]
|
|
for r in json.load(f) if "reference" in r}
|
|
|
|
with ProcessPoolExecutor(max_workers=args.workers) as ex:
|
|
futs = {}
|
|
if do_cases:
|
|
for c in cases:
|
|
futs[ex.submit(_case_task, c, args)] = ("case", c)
|
|
if do_joint:
|
|
for c in cases:
|
|
if c["key"] in JOINT_KEYS:
|
|
futs[ex.submit(_joint_task, c, args)] = ("joint", c)
|
|
if do_recenter:
|
|
for c in cases:
|
|
futs[ex.submit(_recenter_task, c, args, ref_map)] = ("recenter", c)
|
|
for fu in as_completed(futs):
|
|
kind, c = futs[fu]
|
|
try:
|
|
r = fu.result()
|
|
if kind == "joint":
|
|
joint_rows.append(r)
|
|
print(f" [joint] {c['key']}: mi={r.get('mi')} "
|
|
f"il={r.get('impact_loss')} {r.get('seconds')}s "
|
|
f"err={(r.get('error') or '')[:100]}", flush=True)
|
|
elif kind == "recenter":
|
|
recenter_rows.append(r)
|
|
rr = r.get("rigid_recentered", {})
|
|
print(f" [recenter] {c['key']}: mi={rr.get('mi')} "
|
|
f"ctr_err={rr.get('head_centre_err_mm')} "
|
|
f"rot_err={rr.get('rot_err_vs_ref_deg')} "
|
|
f"off_before={r.get('before', {}).get('head_off_volume_centre_mm')} "
|
|
f"off_after={r.get('crop', {}).get('head_off_volume_centre_mm')} "
|
|
f"note={(r.get('note') or r.get('error') or '')[:60]}", flush=True)
|
|
else:
|
|
rows.append(r)
|
|
print(f" [case] {c['key']}: ref.mi={r.get('reference', {}).get('mi')} "
|
|
f"rigid.mi={r.get('rigid', {}).get('mi')} "
|
|
f"affine.mi={r.get('affine', {}).get('mi')} "
|
|
f"note={r.get('note', r.get('error', ''))[:60]}", flush=True)
|
|
except Exception as e: # noqa
|
|
print(f" {c['key']}: ERROR {e!r}", flush=True)
|
|
(joint_rows if kind == "joint" else
|
|
recenter_rows if kind == "recenter" else rows).append(
|
|
{"key": c["key"], "error": repr(e)})
|
|
rows.sort(key=lambda x: x["key"])
|
|
joint_rows.sort(key=lambda x: x["key"])
|
|
recenter_rows.sort(key=lambda x: x["key"])
|
|
|
|
if do_cases:
|
|
fp = os.path.join(args.out_dir, "partABD.json")
|
|
merged = {r["key"]: r for r in rows}
|
|
if os.path.exists(fp):
|
|
with open(fp) as f:
|
|
for r in json.load(f):
|
|
# merge missing sections (a partial run must not wipe
|
|
# data from a fuller earlier run)
|
|
if r["key"] in merged:
|
|
for k, v in r.items():
|
|
merged[r["key"]].setdefault(k, v)
|
|
else:
|
|
merged[r["key"]] = r
|
|
with open(fp, "w") as f:
|
|
json.dump(sorted(merged.values(), key=lambda x: x["key"]),
|
|
f, indent=1)
|
|
print("wrote", fp)
|
|
if do_joint:
|
|
fp = os.path.join(args.out_dir, "partC_joint.json")
|
|
with open(fp, "w") as f:
|
|
json.dump(joint_rows, f, indent=1)
|
|
print("wrote", fp)
|
|
if do_recenter:
|
|
fp = os.path.join(args.out_dir, "partF_recenter.json")
|
|
with open(fp, "w") as f:
|
|
json.dump(recenter_rows, f, indent=1)
|
|
print("wrote", fp)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |