"""Prep utilities to standardize T1 + lesion masks into MNI and normalized form. Defaults are self-contained under the project root: - ANTs binaries are resolved from `/tools/ants/(bin)/` first. - TemplateFlow cache defaults to `/data/templateflow`. Environment overrides still work: - `ANTS_REG` / `ANTS_APPLY` - `TEMPLATEFLOW_HOME` """ from __future__ import annotations from dataclasses import dataclass from pathlib import Path import os, shutil, subprocess, re, json, csv, hashlib, platform, stat, tempfile, time, zipfile import urllib.request import numpy as np import nibabel as nib import scipy.ndimage as ndi from nibabel.processing import resample_from_to from typing import Iterable from collections import defaultdict try: import resource except Exception: # pragma: no cover - resource is Unix-only resource = None def _project_root() -> Path: # .../ARC_ATLAS_Train_v4/src/data_prep/prep_utils.py -> .../ARC_ATLAS_Train_v4 return Path(__file__).resolve().parents[2] def _first_existing(paths: list[Path]) -> Path | None: for p in paths: if p and p.exists(): return p return None def _ensure_templateflow_home() -> Path: tf_home = os.environ.get("TEMPLATEFLOW_HOME") if tf_home: return Path(tf_home).expanduser() local_tf_home = _project_root() / "data" / "templateflow" local_tf_home.mkdir(parents=True, exist_ok=True) os.environ["TEMPLATEFLOW_HOME"] = str(local_tf_home) return local_tf_home def _ants_candidates(binary_name: str) -> list[Path]: root = _project_root() return [ root / "tools" / "ants" / "bin" / binary_name, root / "tools" / "ants" / binary_name, root / "tools" / binary_name, ] _ANTS_RELEASE_API = "https://api.github.com/repos/ANTsX/ANTs/releases/latest" _ANTS_FALLBACK_URLS = [ "https://github.com/ANTsX/ANTs/releases/download/v2.6.5/ants-2.6.5-ubuntu-22.04-X64-gcc.zip", "https://github.com/ANTsX/ANTs/releases/download/v2.6.5/ants-2.6.5-ubuntu20.04-X64-gcc.zip", "https://github.com/ANTsX/ANTs/releases/download/v2.6.5/ants-2.6.5-ubuntu18.04-X64-gcc.zip", ] _TF_TEMPLATE = "tpl-MNI152NLin2009cAsym" def _template_relpath(resolution: int) -> Path: res_tag = f"{int(resolution):02d}" return Path(_TF_TEMPLATE) / f"{_TF_TEMPLATE}_res-{res_tag}_desc-brain_T1w.nii.gz" def _template_url(resolution: int) -> str: return f"https://templateflow.s3.amazonaws.com/{_template_relpath(resolution).as_posix()}" def _ants_asset_preferences() -> list[str]: system = platform.system().lower() machine = platform.machine().lower() if system == "linux" and machine in {"x86_64", "amd64"}: return [ "ubuntu-22.04-X64-gcc.zip", "ubuntu20.04-X64-gcc.zip", "ubuntu18.04-X64-gcc.zip", "ubuntu-24.04-X64-gcc.zip", "almalinux9-X64-gcc.zip", "almalinux8-X64-gcc.zip", "centos7-X64-gcc.zip", ] if system == "darwin" and machine in {"arm64", "aarch64"}: return ["macos-14-ARM64-clang.zip"] if system == "darwin" and machine in {"x86_64", "amd64"}: return ["macos-15-intel-X64-clang.zip"] return [] def _resolve_latest_ants_zip_url() -> str | None: req = urllib.request.Request(_ANTS_RELEASE_API, headers={"User-Agent": "chronic-stroke-segmentation"}) with urllib.request.urlopen(req, timeout=30) as resp: payload = json.load(resp) assets = payload.get("assets", []) if not assets: return None prefs = _ants_asset_preferences() for suffix in prefs: for asset in assets: name = str(asset.get("name", "")) if name.endswith(suffix): url = asset.get("browser_download_url") if url: return str(url) for asset in assets: name = str(asset.get("name", "")) if name.endswith(".zip"): url = asset.get("browser_download_url") if url: return str(url) return None def _download_file(url: str, destination: Path): req = urllib.request.Request(url, headers={"User-Agent": "chronic-stroke-segmentation"}) with urllib.request.urlopen(req, timeout=180) as resp, open(destination, "wb") as out_f: shutil.copyfileobj(resp, out_f) def _install_ants_from_zip(zip_path: Path, install_root: Path): with tempfile.TemporaryDirectory(prefix="ants_extract_") as td: extract_root = Path(td) with zipfile.ZipFile(zip_path) as zf: zf.extractall(extract_root) candidates = [] for p in extract_root.rglob("*"): if not p.is_dir(): continue reg = p / "bin" / "antsRegistration" app = p / "bin" / "antsApplyTransforms" if reg.exists() and app.exists(): candidates.append(p) if not candidates: raise RuntimeError("Downloaded archive did not contain ANTs binaries.") src_root = min(candidates, key=lambda p: len(str(p))) tmp_install = install_root.parent / f"{install_root.name}.tmp" if tmp_install.exists(): shutil.rmtree(tmp_install, ignore_errors=True) shutil.copytree(src_root, tmp_install) if install_root.exists(): shutil.rmtree(install_root, ignore_errors=True) tmp_install.rename(install_root) reg = install_root / "bin" / "antsRegistration" app = install_root / "bin" / "antsApplyTransforms" for binary in (reg, app): if binary.exists(): binary.chmod(binary.stat().st_mode | stat.S_IXUSR | stat.S_IXGRP | stat.S_IXOTH) def _auto_install_ants() -> tuple[Path, Path]: install_root = _project_root() / "tools" / "ants" install_root.parent.mkdir(parents=True, exist_ok=True) urls: list[str] = [] env_url = os.environ.get("ANTS_ZIP_URL") if env_url: urls.append(env_url) try: latest = _resolve_latest_ants_zip_url() if latest: urls.append(latest) except Exception as exc: print(f"[setup] unable to resolve ANTs latest release automatically: {exc}") urls.extend(_ANTS_FALLBACK_URLS) # preserve order but remove duplicates deduped: list[str] = [] seen = set() for u in urls: if u and u not in seen: deduped.append(u) seen.add(u) last_exc: Exception | None = None with tempfile.TemporaryDirectory(prefix="ants_download_") as td: archive = Path(td) / "ants.zip" for url in deduped: print(f"[setup] downloading ANTs binaries from: {url}") try: _download_file(url, archive) _install_ants_from_zip(archive, install_root) reg = install_root / "bin" / "antsRegistration" app = install_root / "bin" / "antsApplyTransforms" if reg.exists() and app.exists(): print(f"[setup] ANTs installed under {install_root}") return reg, app except Exception as exc: last_exc = exc print(f"[setup] failed from {url}: {exc}") continue if last_exc is not None: raise RuntimeError(f"Failed to install ANTs binaries automatically: {last_exc}") raise RuntimeError("Failed to install ANTs binaries automatically.") def _ants_bins(): reg_env = Path(os.environ["ANTS_REG"]).expanduser() if os.environ.get("ANTS_REG") else None app_env = Path(os.environ["ANTS_APPLY"]).expanduser() if os.environ.get("ANTS_APPLY") else None reg = reg_env or _first_existing(_ants_candidates("antsRegistration")) app = app_env or _first_existing(_ants_candidates("antsApplyTransforms")) if (reg is None or not reg.exists()) and shutil.which("antsRegistration"): reg = Path(shutil.which("antsRegistration")) if (app is None or not app.exists()) and shutil.which("antsApplyTransforms"): app = Path(shutil.which("antsApplyTransforms")) if reg is None or app is None or not reg.exists() or not app.exists(): auto_install = os.environ.get("ANTS_AUTO_INSTALL", "1").lower() not in {"0", "false", "no"} if auto_install: print("[setup] ANTs binaries missing; attempting automatic install to tools/ants") try: reg, app = _auto_install_ants() except Exception as exc: print(f"[setup] automatic ANTs install failed: {exc}") if reg is None or app is None or not reg.exists() or not app.exists(): exp_reg = _ants_candidates("antsRegistration")[0] exp_app = _ants_candidates("antsApplyTransforms")[0] raise FileNotFoundError( "ANTs binaries not found. Expected local binaries under " f"'{exp_reg.parent}' (e.g. {exp_reg.name}, {exp_app.name}), " "or set ANTS_REG / ANTS_APPLY. You can also call ensure_prep_runtime() " "or set ANTS_ZIP_URL to an ANTs release zip." ) return reg, app def ensure_prep_runtime(prefetch_template: bool = True) -> dict[str, str]: """Ensure ANTs binaries and TemplateFlow MNI template are available locally.""" reg, app = _ants_bins() info = { "ants_registration": str(reg), "ants_apply_transforms": str(app), } if prefetch_template: tpl = _tpl_path(resolution=1) info["template_mni152_1mm_t1w"] = str(tpl) return info def _command_timeout_seconds() -> int | None: """Timeout for ANTs subprocesses; <=0 disables timeout.""" raw = os.environ.get("PREP_CMD_TIMEOUT_SEC", "").strip() if not raw: raw = os.environ.get("ANTS_CMD_TIMEOUT_SEC", "").strip() if not raw: return 900 try: seconds = int(float(raw)) except ValueError: print(f"[setup] invalid PREP_CMD_TIMEOUT_SEC={raw!r}; falling back to 900s") return 900 return None if seconds <= 0 else seconds def _affinity_cpu_count() -> int: try: return max(1, len(os.sched_getaffinity(0))) except Exception: return max(1, os.cpu_count() or 1) def _cpu_quota_count() -> int | None: cpu_max = Path("/sys/fs/cgroup/cpu.max") if cpu_max.exists(): try: quota_raw, period_raw = cpu_max.read_text().strip().split()[:2] if quota_raw != "max": quota = int(quota_raw) period = int(period_raw) if quota > 0 and period > 0: return max(1, int(quota / period)) except Exception: pass quota_path = Path("/sys/fs/cgroup/cpu/cpu.cfs_quota_us") period_path = Path("/sys/fs/cgroup/cpu/cpu.cfs_period_us") if quota_path.exists() and period_path.exists(): try: quota = int(quota_path.read_text().strip()) period = int(period_path.read_text().strip()) if quota > 0 and period > 0: return max(1, int(quota / period)) except Exception: pass return None def _available_cpu_count() -> int: counts = [_affinity_cpu_count()] quota_cpus = _cpu_quota_count() if quota_cpus is not None: counts.append(quota_cpus) return max(1, min(counts)) def _ants_thread_count() -> int: raw = os.environ.get("PREP_ANTS_THREADS", "").strip() if raw: try: return max(1, int(float(raw))) except ValueError: print(f"[setup] invalid PREP_ANTS_THREADS={raw!r}; using automatic thread count") cap_raw = os.environ.get("PREP_ANTS_THREAD_CAP", "8").strip() try: cap = max(1, int(float(cap_raw))) except ValueError: print(f"[setup] invalid PREP_ANTS_THREAD_CAP={cap_raw!r}; using 8") cap = 8 return max(1, min(_available_cpu_count(), cap)) def _subprocess_env() -> dict[str, str]: env = os.environ.copy() threads = str(_ants_thread_count()) env["ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS"] = threads env["OMP_NUM_THREADS"] = threads env["OMP_THREAD_LIMIT"] = threads return env def _child_cpu_snapshot() -> tuple[float, float] | None: if resource is None: return None usage = resource.getrusage(resource.RUSAGE_CHILDREN) return float(usage.ru_utime), float(usage.ru_stime) def _child_cpu_summary(before: tuple[float, float] | None, elapsed: float) -> str: after = _child_cpu_snapshot() if before is None or after is None: return "child_cpu=unknown" user = max(0.0, after[0] - before[0]) system = max(0.0, after[1] - before[1]) total = user + system effective = total / elapsed if elapsed > 0 else 0.0 return f"child_cpu={total:.1f}s user={user:.1f}s sys={system:.1f}s effective_cores={effective:.2f}" def _run(cmd: list[str]): cmd_str = " ".join(map(str, cmd)) timeout_sec = _command_timeout_seconds() timeout_label = f"{timeout_sec}s" if timeout_sec is not None else "disabled" env = _subprocess_env() threads = env["ITK_GLOBAL_DEFAULT_NUMBER_OF_THREADS"] print(">>", cmd_str) print(f"[cmd] timeout={timeout_label}") print( f"[cmd] ants_threads={threads} " f"available_cpus={_available_cpu_count()} " f"affinity_cpus={_affinity_cpu_count()} " f"cpu_quota_cpus={_cpu_quota_count() or 'unknown'} " f"thread_cap={os.environ.get('PREP_ANTS_THREAD_CAP', '8')}" ) t0 = time.monotonic() child_cpu_before = _child_cpu_snapshot() try: res = subprocess.run( cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True, timeout=timeout_sec, env=env, ) except subprocess.TimeoutExpired as exc: elapsed = time.monotonic() - t0 partial = exc.stdout or "" if isinstance(partial, bytes): partial = partial.decode(errors="replace") tail = "\n".join(partial.splitlines()[-120:]) cpu_summary = _child_cpu_summary(child_cpu_before, elapsed) raise RuntimeError( f"Command timed out after {elapsed:.1f}s (limit={timeout_sec}s): {cmd_str}\n" f"[cmd] {cpu_summary}\n" f"---- output tail ----\n{tail}" ) from exc elapsed = time.monotonic() - t0 print(f"[cmd] completed in {elapsed:.1f}s rc={res.returncode} {_child_cpu_summary(child_cpu_before, elapsed)}") if res.returncode != 0: output = res.stdout or "" tail = "\n".join(output.splitlines()[-200:]) raise RuntimeError( f"Command failed with exit code {res.returncode}: {cmd_str}\n" f"---- output tail ----\n{tail}" ) return res.stdout def _key_from_name(name: str) -> str | None: sub = re.search(r"(sub-[^_]+)", name) ses = re.search(r"(ses-[^_]+)", name) parts = [m.group(1) for m in (sub, ses) if m] return "_".join(parts) if parts else None def _slug_from_raw(raw_root: Path, name: str) -> str: stem = raw_root.name safe = re.sub(r"[^a-zA-Z0-9]+", "-", f"{name}-{stem}").strip("-") digest = hashlib.md5(str(raw_root.resolve()).encode()).hexdigest()[:8] return f"{safe}-{digest}" def _tpl_path(resolution: int = 1) -> Path: """Return MNI T1 path for given resolution, using local cache with robust fallback.""" tf_home = _ensure_templateflow_home() local_tpl = tf_home / _template_relpath(resolution) if local_tpl.exists(): return local_tpl # Prefer TemplateFlow client when available. try: from templateflow.api import get as tf_get except Exception as exc: print(f"[setup] templateflow package unavailable ({exc}); using direct template download.") else: try: tpl = tf_get( "MNI152NLin2009cAsym", resolution=resolution, suffix="T1w", desc="brain", extension="nii.gz", ) tpl_path = Path(tpl[0]) if isinstance(tpl, (list, tuple)) else Path(tpl) if tpl_path.exists(): return tpl_path except Exception as exc: print(f"[setup] templateflow API lookup failed ({exc}); using direct template download.") # Direct download fallback avoids optional templateflow dependency issues. local_tpl.parent.mkdir(parents=True, exist_ok=True) _download_file(_template_url(resolution), local_tpl) if not local_tpl.exists(): raise FileNotFoundError(f"Template download failed: expected {local_tpl}") return local_tpl def _find_existing_output(out_root: Path, raw_root: Path) -> Path | None: """Look for an existing prep folder whose marker matches raw_root.""" raw_root = raw_root.resolve() for m in out_root.glob("*/source.json"): try: info = json.loads(m.read_text()) if Path(info.get("raw_root", "")).resolve() == raw_root: return m.parent except Exception: continue return None def list_pairs(raw_root: Path, t1_glob: str, mask_glob: str): t1s = list(raw_root.glob(t1_glob)) masks = list(raw_root.glob(mask_glob)) return _match_pairs(t1s, masks) def _norm_key_from_name(name: str) -> str: base = name if base.endswith(".nii.gz"): base = base[:-7] elif base.endswith(".nii"): base = base[:-4] drop = [ "_T1w_MNI_norm", "_T1w_MNI", "_T1w_brain", "_T1w", "_T1", "_lesion_mask_MNI_clean", "_lesion_mask_MNI", "_lesion_mask", "_desc-lesion_mask", "_mask", "mask" ] for sfx in drop: if base.endswith(sfx): base = base[: -len(sfx)] return base def _match_pairs(t1s: list[Path], masks: list[Path]): """Match T1 and mask by the most specific key available. Priority: exact normalized basename -> unique; else sub/ses key unique. Ambiguous cases are skipped. """ mask_by_base = defaultdict(list) mask_by_subses = defaultdict(list) for m in masks: base = _norm_key_from_name(m.name) mask_by_base[base].append(m) key = _key_from_name(m.name) or base mask_by_subses[key].append(m) pairs = [] for t1 in t1s: base = _norm_key_from_name(t1.name) key = _key_from_name(t1.name) or base chosen = None if mask_by_base.get(base): if len(mask_by_base[base]) == 1: chosen = mask_by_base[base][0] else: print(f"[warn] multiple masks share base {base}; skipping") continue elif mask_by_subses.get(key): if len(mask_by_subses[key]) == 1: chosen = mask_by_subses[key][0] else: print(f"[warn] ambiguous masks for {t1.name} (key {key}): {len(mask_by_subses[key])}; skipping") continue if not chosen: print(f"[warn] no mask for {t1.name} (key {key}); skipping") continue pairs.append((t1, chosen, base)) return pairs def resample_mask_to_t1(mask: Path, t1: Path, out_path: Path): mi = nib.load(str(mask)) ti = nib.load(str(t1)) if mi.shape[:3] != ti.shape[:3] or not np.allclose(mi.affine, ti.affine, atol=1e-4): rs = resample_from_to(mi, (ti.shape, ti.affine), order=0) data = (rs.get_fdata() > 0.5).astype(np.uint8) else: data = (mi.get_fdata() > 0.5).astype(np.uint8) out_path.parent.mkdir(parents=True, exist_ok=True) _save_nifti_like(out_path, data, ti, np.uint8) def _save_nifti_like(path: Path, data: np.ndarray, ref_img: nib.Nifti1Image, dtype) -> Path: path.parent.mkdir(parents=True, exist_ok=True) header = ref_img.header.copy() header.set_data_dtype(dtype) nib.save(nib.Nifti1Image(np.asarray(data, dtype=dtype), ref_img.affine, header), str(path)) return path def _ants_registration_profile() -> dict[str, str]: """Return ANTs registration schedule based on PREP_ANTS_PROFILE.""" profile = os.environ.get("PREP_ANTS_PROFILE", "balanced").strip().lower() presets: dict[str, dict[str, str]] = { "quick": { "metric_sampling": "Random,0.1", "rigid_affine_convergence": "120x60x20", "rigid_affine_smoothing": "3x2x1vox", "rigid_affine_shrink": "8x4x2", "syn_convergence": "8x4x2", "syn_smoothing": "2x1x0vox", "syn_shrink": "4x2x1", "syn_transform": "SyN[0.05,3,0]", }, "fast": { "metric_sampling": "Random,0.2", "rigid_affine_convergence": "300x120x40", "rigid_affine_smoothing": "3x2x1vox", "rigid_affine_shrink": "6x4x2", "syn_convergence": "24x12x6", "syn_smoothing": "2x1x0vox", "syn_shrink": "4x2x1", "syn_transform": "SyN[0.08,3,0]", }, "balanced": { "metric_sampling": "Regular,0.2", "rigid_affine_convergence": "600x250x100", "rigid_affine_smoothing": "3x2x1vox", "rigid_affine_shrink": "4x2x1", "syn_convergence": "40x20x10", "syn_smoothing": "2x1x0vox", "syn_shrink": "4x2x1", "syn_transform": "SyN[0.1,3,0]", }, "accurate": { "metric_sampling": "Regular,0.25", "rigid_affine_convergence": "1000x500x250", "rigid_affine_smoothing": "3x2x1vox", "rigid_affine_shrink": "4x2x1", "syn_convergence": "60x40x20", "syn_smoothing": "2x1x0vox", "syn_shrink": "4x2x1", "syn_transform": "SyN[0.1,3,0]", }, } if profile not in presets: print(f"[setup] unknown PREP_ANTS_PROFILE={profile!r}; using 'balanced'") profile = "balanced" cfg = dict(presets[profile]) cfg["profile"] = profile return cfg def ants_register(t1: Path, prefix: Path, tpl: Path, reg_bin: Path, use_2mm: bool = True): tpl_reg = tpl if use_2mm: try: from templateflow.api import get as tf_get tpl2 = tf_get("MNI152NLin2009cAsym", resolution=2, suffix="T1w", desc="brain", extension="nii.gz") tpl_reg = Path(tpl2[0]) if isinstance(tpl2, (list, tuple)) else Path(tpl2) except Exception: tpl_reg = tpl cfg = _ants_registration_profile() print( "[ants] profile=" f"{cfg['profile']} rigid_affine={cfg['rigid_affine_convergence']} " f"syn={cfg['syn_convergence']}" ) _run([ str(reg_bin), '-d','3', '-r', f'[{tpl_reg},{t1},1]', '-m', f"Mattes[{tpl_reg},{t1},1,32,{cfg['metric_sampling']}]", '-t','Rigid[0.1]', '-c', cfg['rigid_affine_convergence'], '-s', cfg['rigid_affine_smoothing'], '-f', cfg['rigid_affine_shrink'], '-m', f"Mattes[{tpl_reg},{t1},1,32,{cfg['metric_sampling']}]", '-t','Affine[0.1]', '-c', cfg['rigid_affine_convergence'], '-s', cfg['rigid_affine_smoothing'], '-f', cfg['rigid_affine_shrink'], '-m', f'CC[{tpl_reg},{t1},1,4]', '-t', cfg['syn_transform'], '-c', cfg['syn_convergence'], '-s', cfg['syn_smoothing'], '-f', cfg['syn_shrink'], '-o', f'[{prefix},{prefix}warped.nii.gz,{prefix}invwarped.nii.gz]' ]) def ants_apply(img_in: Path, ref: Path, xfm_prefix: Path, out_path: Path, apply_bin: Path, nn: bool=False): args = [str(apply_bin), '-d','3', '-i', str(img_in), '-r', str(ref), '-o', str(out_path)] if nn: args += ['-n','NearestNeighbor'] args += ['-t', str(xfm_prefix)+'1Warp.nii.gz', '-t', str(xfm_prefix)+'0GenericAffine.mat'] _run(args) def normalize_t1(vol: np.ndarray) -> np.ndarray: nz = vol[vol>0] if nz.size == 0: return np.zeros_like(vol, np.float32) p1,p99 = np.percentile(nz,[1,99]) vol = np.clip(vol, p1, p99) mu, sd = nz.mean(), nz.std() vol = (vol - mu)/(sd+1e-8) mn, mx = vol.min(), vol.max() return ((vol - mn)/(mx - mn + 1e-8)).astype(np.float32) def largest_component(mask: np.ndarray) -> np.ndarray: labeled, nlab = ndi.label(mask) if nlab <= 1: return mask.astype(np.uint8) sizes = np.bincount(labeled.ravel()) keep = sizes[1:].argmax() + 1 return (labeled == keep).astype(np.uint8) def _bbox(mask: np.ndarray): coords = np.argwhere(mask > 0) if coords.size == 0: return None mins = coords.min(axis=0) maxs = coords.max(axis=0) return mins, maxs def _overlap_score(mask: np.ndarray, brain: np.ndarray) -> float: inter = np.logical_and(mask > 0, brain > 0).sum() return inter / max((mask > 0).sum(), 1) def _try_flips(mask: np.ndarray, brain: np.ndarray): base_score = _overlap_score(mask, brain) best = (mask, base_score, "none") flips = [ (np.flip(mask, axis=0), "flip_x"), (np.flip(mask, axis=1), "flip_y"), (np.flip(mask, axis=2), "flip_z"), ] for flipped, tag in flips: score = _overlap_score(flipped, brain) if score > best[1]: best = (flipped, score, tag) return best @dataclass class DatasetConfig: name: str raw_root: Path | None = None images_dir: Path | None = None masks_dir: Path | None = None t1_glob: str = "**/*_T1w.nii.gz" mask_glob: str = "**/*mask*.nii.gz" overwrite: bool = False already_mni: bool = False def run_prep(datasets: Iterable[DatasetConfig], out_root: Path, force_overwrite: bool = False): tpl = None reg_bin = apply_bin = None out_root.mkdir(parents=True, exist_ok=True) qc_rows = [] outputs = [] for ds in datasets: raw = Path(ds.raw_root).expanduser() if ds.raw_root else None img_root = Path(ds.images_dir).expanduser() if ds.images_dir else raw msk_root = Path(ds.masks_dir).expanduser() if ds.masks_dir else raw needs_ants = not ds.already_mni print(f"[{ds.name}] image root: {img_root} | mask root: {msk_root}") print(f"[{ds.name}] globs: t1={ds.t1_glob} masks={ds.mask_glob}") if not img_root or not img_root.exists(): print(f"[skip] {ds.name}: images root missing {img_root}") continue if not msk_root or not msk_root.exists(): print(f"[skip] {ds.name}: masks root missing {msk_root}") continue t1s = list(img_root.glob(ds.t1_glob)) mks = list(msk_root.glob(ds.mask_glob)) pairs = _match_pairs(t1s, mks) print(f"[{ds.name}] images: {len(t1s)} masks: {len(mks)} pairs found: {len(pairs)}") slug_base = raw or img_root slug = _slug_from_raw(slug_base, ds.name) overwrite = force_overwrite or ds.overwrite out_ds_existing = None if overwrite else _find_existing_output(out_root, slug_base) out_ds = out_ds_existing or (out_root / slug) marker = out_ds / "source.json" if overwrite and out_ds.exists(): shutil.rmtree(out_ds, ignore_errors=True) print(f"[{ds.name}] overwrite=True -> cleared {out_ds}") elif marker.exists() and not overwrite: print(f"[{ds.name}] already processed -> {out_ds}, skipping (overwrite=True to redo).") outputs.append(out_ds) continue if not pairs: print(f"[warn] {ds.name}: no matched pairs; skipping dataset.") continue out_nat = out_ds / 'native_resampled_masks' out_mni = out_ds / 'mni_1mm_ants_fixed' out_norm = out_mni / 't1_norm' out_clean = out_mni / 'masks_clean' out_xfm = out_ds / 'xfm' for d in (out_nat, out_mni, out_norm, out_clean, out_xfm): d.mkdir(parents=True, exist_ok=True) for t1, mask, key in pairs: mask_t1 = out_nat / f"{key}_lesion_mask_T1w_native.nii.gz" prefix = out_xfm / f"{key}_t1_to_mni_" t1_mni = out_mni / f"{key}_T1w_MNI.nii.gz" mask_mni = out_mni / f"{key}_lesion_mask_MNI.nii.gz" t1_norm = out_norm / f"{key}_T1w_MNI_norm.nii.gz" mask_clean = out_clean / f"{key}_lesion_mask_MNI_clean.nii.gz" if ds.already_mni: if not mask_t1.exists(): resample_mask_to_t1(mask, t1, mask_t1) if not t1_mni.exists(): shutil.copy2(t1, t1_mni) if not mask_mni.exists(): shutil.copy2(mask_t1, mask_mni) else: if tpl is None: tpl = _tpl_path(resolution=1) if reg_bin is None or apply_bin is None: reg_bin, apply_bin = _ants_bins() if not mask_t1.exists(): resample_mask_to_t1(mask, t1, mask_t1) if not (prefix.with_name(prefix.name+'0GenericAffine.mat')).exists(): ants_register(t1, prefix, tpl, reg_bin, use_2mm=True) if not t1_mni.exists(): ants_apply(t1, tpl, prefix, t1_mni, apply_bin, nn=False) if not mask_mni.exists(): ants_apply(mask_t1, tpl, prefix, mask_mni, apply_bin, nn=True) mi = nib.load(str(mask_mni)); data=(mi.get_fdata()>0.5).astype(np.uint8) _save_nifti_like(mask_mni, data, mi, np.uint8) if not t1_norm.exists(): t1_mni_img = nib.load(str(t1_mni)) vol = t1_mni_img.get_fdata().astype(np.float32) norm = normalize_t1(vol) _save_nifti_like(t1_norm, norm, t1_mni_img, np.float32) if not mask_clean.exists(): data = (nib.load(str(mask_mni)).get_fdata()>0.5).astype(np.uint8) data = largest_component(data) # alignment sanity check: ensure mask overlaps brain; try flips if not brain = (nib.load(str(t1_mni)).get_fdata()>0).astype(np.uint8) best_mask, best_score, tag = _try_flips(data, brain) if best_score < 0.1: print(f"[warn] low overlap for {mask_mni.name} (score {best_score:.3f}); keeping original") elif tag != "none": print(f"[info] flipped {tag} for {mask_mni.name} (overlap {best_score:.3f})") data = best_mask mask_mni_img = nib.load(str(mask_mni)) _save_nifti_like(mask_clean, data, mask_mni_img, np.uint8) ti = nib.load(str(t1_mni)); mi = nib.load(str(mask_clean)) qc_rows.append(dict( dataset=ds.name, slug=slug, key=key, t1_mni=str(t1_mni), mask_mni=str(mask_clean), t1_shape=str(ti.shape[:3]), t1_zooms=str(tuple(round(z,3) for z in ti.header.get_zooms()[:3])), mask_nonzero=int(np.count_nonzero(mi.get_fdata()>0)), )) marker.parent.mkdir(parents=True, exist_ok=True) marker.write_text(json.dumps({"raw_root": str(slug_base), "slug": slug}, indent=2)) outputs.append(out_ds) if qc_rows: qc_csv = out_root / 'prep_qc.csv' with open(qc_csv, 'w', newline='') as f: writer = csv.DictWriter(f, fieldnames=qc_rows[0].keys()) writer.writeheader(); writer.writerows(qc_rows) print('QC written', qc_csv) else: print('No QC rows written (no datasets processed).') return outputs def combine_standardized(dataset_roots: Iterable[Path], dest: Path): dest_t1 = dest / "t1" dest_mk = dest / "masks" dest_t1.mkdir(parents=True, exist_ok=True) dest_mk.mkdir(parents=True, exist_ok=True) manifest = [] for ds_root in dataset_roots: slug = ds_root.name t1_dir = ds_root / "mni_1mm_ants_fixed" / "t1_norm" msk_dir = ds_root / "mni_1mm_ants_fixed" / "masks_clean" if not t1_dir.exists() or not msk_dir.exists(): print(f"[combine] missing t1_norm or masks_clean in {ds_root}, skipping") continue for t1 in sorted(t1_dir.glob("*.nii.gz")): base = t1.name mask = msk_dir / base.replace("_T1w_MNI_norm", "_lesion_mask_MNI_clean") if not mask.exists(): continue out_t1 = dest_t1 / f"{slug}__{base}" out_mk = dest_mk / f"{slug}__{mask.name}" shutil.copy2(t1, out_t1) shutil.copy2(mask, out_mk) manifest.append({"slug": slug, "key": base, "t1": str(out_t1.resolve()), "mask": str(out_mk.resolve())}) if manifest: mf = dest / "manifest.csv" with open(mf, "w", newline="") as f: writer = csv.DictWriter(f, fieldnames=manifest[0].keys()) writer.writeheader(); writer.writerows(manifest) print("Combined manifest ->", mf) else: print("No combined manifest written (no pairs).") def run_prep_images_only( images_dir: Path, out_root: Path, name: str = "TEST", t1_glob: str = "**/*.nii.gz", already_mni: bool = False, overwrite: bool = False, ) -> Path: """Prep images to MNI/normalized space without requiring lesion masks.""" images_dir = Path(images_dir).expanduser() out_root = Path(out_root).expanduser() if not images_dir.exists(): raise FileNotFoundError(f"Images root missing: {images_dir}") out_root.mkdir(parents=True, exist_ok=True) t1s = sorted(images_dir.glob(t1_glob)) if not t1s: raise RuntimeError(f"No T1 images found in {images_dir} using glob {t1_glob}") slug = _slug_from_raw(images_dir, name) out_ds = out_root / slug marker = out_ds / "source.json" if overwrite and out_ds.exists(): shutil.rmtree(out_ds, ignore_errors=True) elif marker.exists() and not overwrite: print(f"[{name}] already processed -> {out_ds}, skipping (overwrite=True to redo).") return out_ds out_mni = out_ds / "mni_1mm_ants_fixed" out_norm = out_mni / "t1_norm" out_xfm = out_ds / "xfm" for d in (out_mni, out_norm, out_xfm): d.mkdir(parents=True, exist_ok=True) tpl = None reg_bin = apply_bin = None seen_keys = set() qc_rows = [] def _unique_key(base: str) -> str: key = base or "case" if key not in seen_keys: seen_keys.add(key) return key i = 2 while f"{key}_{i}" in seen_keys: i += 1 new_key = f"{key}_{i}" seen_keys.add(new_key) return new_key for t1 in t1s: key = _unique_key(_norm_key_from_name(t1.name)) prefix = out_xfm / f"{key}_t1_to_mni_" t1_mni = out_mni / f"{key}_T1w_MNI.nii.gz" t1_norm = out_norm / f"{key}_T1w_MNI_norm.nii.gz" if already_mni: if not t1_mni.exists(): shutil.copy2(t1, t1_mni) else: if tpl is None: tpl = _tpl_path(resolution=1) if reg_bin is None or apply_bin is None: reg_bin, apply_bin = _ants_bins() if not (prefix.with_name(prefix.name + "0GenericAffine.mat")).exists(): ants_register(t1, prefix, tpl, reg_bin, use_2mm=True) if not t1_mni.exists(): ants_apply(t1, tpl, prefix, t1_mni, apply_bin, nn=False) if not t1_norm.exists(): img_mni = nib.load(str(t1_mni)) norm = normalize_t1(img_mni.get_fdata().astype(np.float32)) _save_nifti_like(t1_norm, norm, img_mni, np.float32) ti = nib.load(str(t1_mni)) qc_rows.append( dict( dataset=name, slug=slug, key=key, t1_mni=str(t1_mni), t1_shape=str(ti.shape[:3]), t1_zooms=str(tuple(round(z, 3) for z in ti.header.get_zooms()[:3])), ) ) marker.parent.mkdir(parents=True, exist_ok=True) marker.write_text(json.dumps({"raw_root": str(images_dir), "slug": slug}, indent=2)) if qc_rows: qc_csv = out_root / "prep_qc_images_only.csv" with open(qc_csv, "w", newline="") as f: writer = csv.DictWriter(f, fieldnames=qc_rows[0].keys()) writer.writeheader() writer.writerows(qc_rows) print("QC written", qc_csv) return out_ds def combine_standardized_images_only(dataset_roots: Iterable[Path], dest: Path): """Combine standardized image-only outputs into a single test_input root.""" dest_t1 = dest / "t1" dest_t1.mkdir(parents=True, exist_ok=True) manifest = [] for ds_root in dataset_roots: slug = ds_root.name t1_dir = ds_root / "mni_1mm_ants_fixed" / "t1_norm" if not t1_dir.exists(): print(f"[combine] missing t1_norm in {ds_root}, skipping") continue for t1 in sorted(t1_dir.glob("*.nii.gz")): out_t1 = dest_t1 / f"{slug}__{t1.name}" shutil.copy2(t1, out_t1) manifest.append({"slug": slug, "key": t1.name, "t1": str(out_t1.resolve())}) if manifest: mf = dest / "manifest.csv" with open(mf, "w", newline="") as f: writer = csv.DictWriter(f, fieldnames=manifest[0].keys()) writer.writeheader() writer.writerows(manifest) print("Combined image-only manifest ->", mf) else: print("No image-only manifest written (no images).") __all__ = [ 'DatasetConfig', 'ensure_prep_runtime', 'run_prep', 'combine_standardized', 'run_prep_images_only', 'combine_standardized_images_only', ]