Furen Xiao
e8dfe0fe35
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.
460 lines
No EOL
21 KiB
Python
460 lines
No EOL
21 KiB
Python
"""Build M6-2025 manifests: T1c volumes, labeled when GTV exists.
|
|
|
|
GTV lives in CT space; T1c is registered to CT (rigid, Mattes MI on a head box
|
|
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:
|
|
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'}
|
|
Usage: python scripts/02_build_m6_dataset.py [--workers 4] [--skip-reg]
|
|
"""
|
|
import os
|
|
import sys
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
import re
|
|
import json
|
|
import argparse
|
|
import numpy as np
|
|
import SimpleITK as sitk
|
|
from concurrent.futures import ProcessPoolExecutor, as_completed
|
|
from src.common import ROOT, DATA, d, save_jsonl, read_nii_img, path
|
|
|
|
BASE = path("m6")
|
|
T1C_RE = re.compile(r"T1.*\+C|fl3d.*\+.*c", re.I)
|
|
EXCL_RE = re.compile(r"FLAIR|DTI|vibe|dixon|t2|SWI|T2|MPR_Cor", re.I)
|
|
|
|
|
|
def pick_t1c(mrd):
|
|
cands = []
|
|
for f in sorted(os.listdir(mrd)):
|
|
if not f.endswith(".nii.gz"):
|
|
continue
|
|
b = f[: -len(".nii.gz")]
|
|
if T1C_RE.search(b) and not EXCL_RE.search(b):
|
|
pref = 0
|
|
if "MPR_Tra" in b:
|
|
pref = -2
|
|
if "_c" in b.lower() or "+C" in b:
|
|
pref -= 1
|
|
cands.append((pref, f))
|
|
cands.sort(key=lambda x: (x[0], x[1]))
|
|
return os.path.join(mrd, cands[0][1]) if cands else None
|
|
|
|
|
|
def discover():
|
|
labeled, unlabeled = [], []
|
|
for pid in sorted(os.listdir(BASE)):
|
|
fp = os.path.join(BASE, pid)
|
|
if not os.path.isdir(fp):
|
|
continue
|
|
for date in sorted(os.listdir(fp)):
|
|
dp = os.path.join(fp, date)
|
|
if not os.path.isdir(dp) or not re.match(r"\d{8}$", date):
|
|
continue
|
|
mrd = os.path.join(dp, "MR")
|
|
if not os.path.isdir(mrd):
|
|
continue
|
|
t1c = pick_t1c(mrd)
|
|
if not t1c:
|
|
continue
|
|
rt = os.path.join(dp, "RT")
|
|
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}"
|
|
row = {"key": key, "subject": pid, "date": date, "img": t1c, "source": "m6"}
|
|
if label and os.path.exists(ct):
|
|
labeled.append({**row, "label": label, "ct": ct, "gtv_registered": False})
|
|
else:
|
|
unlabeled.append({**row, "label": None})
|
|
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):
|
|
"""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
|
|
t_start = time.time()
|
|
try: # noqa
|
|
sitk.CommonProperties.SetGlobalDefaultNumberOfThreads(threads)
|
|
except AttributeError:
|
|
pass
|
|
ct = sitk.Cast(ct, sitk.sitkFloat32)
|
|
t1c = sitk.Cast(t1c, sitk.sitkFloat32)
|
|
|
|
from src.common import head_mask_from_image
|
|
ha = sitk.GetArrayFromImage(t1c)
|
|
hm = head_mask_from_image(ha)
|
|
# Registration grid = a head box located in CT space. The CTs are head-neck
|
|
# (often head-neck-chest) exams, so cropping to the MR head box would leave
|
|
# mostly neck/chest/air on the fixed image and MI locks onto the wrong
|
|
# structures. Vertex = topmost tissue slice with 10mm of continuity below
|
|
# (skips small artifacts floating above the head; bone HU is unreliable
|
|
# here, so tissue > -150 HU is used). x/y from the tissue extent of the
|
|
# 150mm band below the vertex. Box covers the whole head (x/y: extent +
|
|
# 15mm margin, z: +10 to -215 from vertex).
|
|
ca = sitk.GetArrayFromImage(ct)
|
|
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):
|
|
reg_dir = d("data/m6reg")
|
|
key = row["key"]
|
|
t1c = read_nii_img(row["img"])
|
|
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"])
|
|
tr, info, tissue_w = register_ct_to_t1c(ct, t1c)
|
|
# QA files in m6reg, everything in (oriented) MR coordinates:
|
|
# _ct: CT warped to the MR grid (linear), for overlaying _tv
|
|
# _mr: the oriented MR volume itself
|
|
# _tv: original TV warped to the MR grid (nearest) — the study label
|
|
# _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():
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--workers", type=int, default=2)
|
|
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()
|
|
labeled, unlabeled = discover()
|
|
print(f"m6 labeled candidates: {len(labeled)}, unlabeled: {len(unlabeled)}")
|
|
reg_dir = d("data/m6reg")
|
|
if not args.skip_reg:
|
|
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")
|
|
with ProcessPoolExecutor(max_workers=args.workers) as ex:
|
|
futs = {ex.submit(reg_worker, r): r for r in todo}
|
|
for fu in as_completed(futs):
|
|
k = futs[fu]["key"]
|
|
try:
|
|
key, ok, nvox = fu.result()
|
|
print(f" {key}: ok={ok} nvox={nvox}", flush=True)
|
|
except Exception as e: # noqa
|
|
print(f" {k}: ERR {e}", flush=True)
|
|
# refresh manifest state
|
|
labeled, unlabeled = discover()
|
|
for r in labeled:
|
|
p = os.path.join(reg_dir, r["key"] + "_tv.nii.gz")
|
|
accepted = True
|
|
jf = os.path.join(reg_dir, r["key"] + "_reg.json")
|
|
if os.path.exists(jf):
|
|
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")]
|
|
# 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:
|
|
r.pop("ct", None)
|
|
save_jsonl(keep, os.path.join(DATA, "manifests/m6_labeled.jsonl"))
|
|
save_jsonl(unlabeled, os.path.join(DATA, "manifests/m6_unlabeled.jsonl"))
|
|
print(f"final labeled={len(keep)} unlabeled={len(unlabeled)}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main() |