"""Build M6-2025 manifests: T1c volumes, labeled when GTV exists. GTV lives in CT space; the T1c is the fixed image and the CT is registered to it with the best of three ANTsPy pipelines (src/registration_pipelines.py, full volumes, no mask / crop / pad, each toolkit's native COM init): ants_rigid : Rigid[0.25] (Mattes, 500x500x200, 4x2x1) ants_affine : Affine[0.25] (Mattes, 2100x1200x1200x10, 6x4x2x1) ants_syn : SyN (Affine[0.25] + SyN[0.2,3,0]) The winner is the candidate with the highest Mattes-style MI between the warped CT and the T1c (full volumes). All three return the transform in the unified contract direction (T1c -> CT), so resampling the CT / GTV onto the T1c grid with it places them in T1c native space -- the TV is warped with SMD + linear (signed distance to the TV boundary, linearly resampled, thresholded at the mean CT spacing halved; the validated recipe from scripts/qa_label_interpolation.py). Acceptance requires the warped CT skull ring to be concentric with the MR head silhouette. 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. The ring test is ``>=`` (not ``>``) so it also works on a CT clipped to a brain window (e.g. [0, 80] HU), where the p97 threshold sits exactly on the clipped-skull atom and a strict comparison would give an empty ring. """ 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, sl.shape[0] - 1) xx = np.clip((cx0 + rr * dx).astype(int), 0, sl.shape[1] - 1) rh = rr[np.where(sl[yy, xx])[0]] rs = rr[np.where(sk[yy, xx])[0]] if len(rh) and len(rs): 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 the full volumes (no mask, no crop / pad). Both inputs share physical space (oriented). The transform is computed on the full CT / T1c volumes, so it applies directly to any resampling. 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) # full volumes, no mask / crop / pad: fixed = CT, moving = T1c fixed, moving = ct, t1c tsp = np.array(t1c.GetSpacing()) 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. # Euler rotation center = the CT head centroid (the 150 mm band below the # vertex, same head definition as before); off maps it onto the MR head # centroid. Both are init-only quantities; the registration itself runs on # the full volumes. ca = sitk.GetArrayFromImage(ct) csp, co = np.array(ct.GetSpacing()), np.array(ct.GetOrigin()) tissue = ca > -150 vz = _ct_vertex_slice(ct) if vz is not None: z0 = max(0, vz - int(round(150 / csp[2]))) idx = np.array([z0, 0, 0]) + np.argwhere(tissue[z0:vz + 1]).mean(0) else: idx = np.argwhere(tissue).mean(0) ct_c = np.array(co) + idx * csp mr_c = np.array(t1c.GetOrigin()) + np.argwhere(hm).mean(0) * tsp 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): # |mid| (not 0 <= mid): full-volume registration can land the CT skull # ring slightly inside the MR silhouette, which is equally concentric return bool(g) and g["std_mm"] <= 12.0 and abs(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, "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 _QA = None def _load_qa(): """Import scripts/qa_label_interpolation.py (the validated SMD + linear label-warp recipe: smd_threshold / make_tv_bin / warp_tv). Cached per process (each worker process keeps its own copy).""" global _QA if _QA is None: import importlib.util p = os.path.join(os.path.dirname(os.path.abspath(__file__)), "qa_label_interpolation.py") spec = importlib.util.spec_from_file_location("qa_label_interp", p) m = importlib.util.module_from_spec(spec) spec.loader.exec_module(m) _QA = m return _QA def register_ants_best(ct, t1c): """Best-of-ANTs CT -> T1c registration with the T1c as the fixed image. Runs the three ANTsPy pipelines from src/registration_pipelines.py (fixed = T1c, moving = CT, full volumes, no mask / crop / pad, each toolkit's native COM init): ants_rigid : Rigid[0.25] (Mattes, 500x500x200, 4x2x1) ants_affine : Affine[0.25] (Mattes, 2100x1200x1200x10, 6x4x2x1) ants_syn : SyN (Affine[0.25] + SyN[0.2,3,0]) Each returns a transform in the unified contract direction (T1c -> CT), so sitk.Resample(ct, t1c, T, ...) / sitk.Resample(tv, t1c, T, ...) place the CT / TV concentric with the MR head -- no inversion (the T1c was the fixed image). The winner is the candidate with the highest Mattes-style MI between the warped CT and the T1c (_mi_value, full volumes). The per-candidate skull-ring gap is still recorded and gates acceptance. Returns (transform, info, tissue_w) like register_ct_to_t1c. """ import time from scipy import ndimage from src import registration_pipelines as RP from src.common import head_mask_from_image t_start = time.time() ct = sitk.Cast(ct, sitk.sitkFloat32) t1c = sitk.Cast(t1c, sitk.sitkFloat32) calls = { "ants_rigid": lambda: RP.register_antspyx(ct, t1c), "ants_affine": lambda: RP.register_ants_affine_variant(ct, t1c, "antspyx"), "ants_syn": lambda: RP.register_ants_variant_deformable(ct, t1c, "ants_syn"), } ha = sitk.GetArrayFromImage(t1c) cands = {} for name, fn in calls.items(): try: T, meta = fn() except Exception as e: # noqa cands[name] = {"error": repr(e)} print(f" {name}: FAILED {e!r}", flush=True) continue warped_ct = RP._warp_ct(ct, t1c, T) mi = _mi_value(ha, warped_ct) gap = RP._skull_gap(warped_ct, ha) cands[name] = {"T": T, "meta": meta, "mi": mi, "gap": gap} print(f" {name}: mi={mi:.4f} gap={gap} ({meta.get('seconds_reg')}s)", flush=True) ok = {n: c for n, c in cands.items() if "T" in c} if not ok: raise RuntimeError("all ANTs candidates failed: " + repr({n: c.get("error") for n, c in cands.items()})) best_name = max(ok, key=lambda n: ok[n]["mi"]) T, meta, gap = ok[best_name]["T"], ok[best_name]["meta"], ok[best_name]["gap"] def _tp(x, y, z): return np.array(T.TransformPoint((float(x), float(y), float(z)))) t0 = _tp(0, 0, 0) mat = np.column_stack([_tp(1, 0, 0) - t0, _tp(0, 1, 0) - t0, _tp(0, 0, 1) - t0]) scale = [round(float(np.linalg.norm(mat[:, i])), 4) for i in range(3)] hm = head_mask_from_image(ha) warped_ct = sitk.GetArrayFromImage(sitk.Resample(ct, t1c, T, sitk.sitkLinear)) posw = warped_ct[warped_ct > 0] tissue_w = ((warped_ct > max(2.0, float(np.percentile(posw, 2)))) if posw.size else np.zeros_like(warped_ct, bool)) tissue_w = ndimage.binary_opening(tissue_w, iterations=2) coverage = float(tissue_w[hm].sum() / max(hm.sum(), 1)) info = { "pipeline": "ants_best", "fixed": "T1c", "moving": "CT", "best": best_name, "candidates": { n: ({**c["meta"], "mi": round(c["mi"], 4), "skull_ring_gap_mm": c["gap"], "selected": n == best_name} if "T" in c else c) for n, c in cands.items()}, "affine_stage": {"matrix": [[round(float(v), 4) for v in row] for row in mat.tolist()], "offset": [round(float(x), 3) for x in t0], "scale": scale, "transform_type": type(T).__name__}, "coverage_head_box": round(coverage, 3), "skull_ring_gap_mm": gap, "seconds_total": round(time.time() - t_start, 1), } print("reg done %.0fs best=%s mi=%.4f gap=%s cov %.2f" % ( info["seconds_total"], best_name, ok[best_name]["mi"], gap, coverage), flush=True) return T, info, tissue_w def reg_worker(row): reg_dir = d("data/m6reg") key = row["key"] 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_ants_best(ct, t1c) # QA files in m6reg, everything in (oriented) MR coordinates: # _ct: CT warped to the MR grid (linear), for overlaying _tv # _mr: the oriented MR volume itself # _tv: original TV warped to the MR grid (SMD + linear) — the study label # _reg.json: registration metrics + final image information qa = _load_qa() thr = qa.smd_threshold([float(s) for s in gtv.GetSpacing()])[0] z = qa.warp_tv(qa.make_tv_bin(gtv), t1c, tr, "smd_linear", thr) tv_mr = sitk.GetImageFromArray(z) tv_mr.CopyInformation(t1c) nvox = int((z > 0.5).sum()) from scipy import ndimage tv_in_ct_frac = 1.0 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 abs(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), "tv_smd_thr_mm": round(thr, 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()