Furen Xiao
16b9828d15
- Added `--clip-ct LO HI` argument to clip CT intensities before registration and scoring, improving A/B comparisons. - Implemented `_clip_image` function to handle intensity clipping of images. - Updated registration functions to utilize CT clipping where beneficial, specifically in `register_niftyreg` and `register_synthmorph`. - Enhanced `fixed_moving` function to ensure CT is always treated as the fixed image. - Introduced Docker support for NiftyReg v2.1.1, including a Dockerfile for building the image. - Updated documentation and help messages to reflect new features and usage instructions. - Adjusted evaluation criteria for registration quality to account for intensity clipping.
447 lines
No EOL
20 KiB
Python
447 lines
No EOL
20 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 the full
|
|
volumes -- no mask, no crop / pad; 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.
|
|
|
|
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, 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 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
|
|
|
|
|
|
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 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),
|
|
"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() |