"""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()