| """MR-RATE atlas study -> 4-channel uint8 volume + 4 presence flags. |
| |
| Every series of a study is registered to the MNI152 1 mm grid (193x229x193), so |
| the series are voxel-aligned. Each MR type becomes one channel: the mean of up |
| to four of its series, each scaled to [0, 1] first. The same code builds the |
| training cache and runs at inference. |
| """ |
| from __future__ import annotations |
|
|
| import gzip |
| import io |
| import re |
| import zipfile |
|
|
| import nibabel as nib |
| import numpy as np |
| from scipy.ndimage import zoom |
|
|
| MODALITIES = ("t1w", "t2w", "flair", "swi") |
| SRC_SHAPE = (193, 229, 193) |
| MAX_PER_MODALITY = 4 |
|
|
| |
| |
| GEOMETRIES = { |
| "d15": {"offset": (12, 16, 6), "crop": (168, 192, 156), "out": (112, 128, 104), "mode": "zoom"}, |
| "d2": {"offset": (0, 2, 0), "crop": (192, 224, 192), "out": (96, 112, 96), "mode": "block"}, |
| } |
|
|
| _MOD_RE = re.compile(r"_atlas_([a-z0-9]+)-", re.IGNORECASE) |
| _MOD_ALT = {m: re.compile(rf"(?:^|[^a-z]){m}(?:[^a-z]|$)", re.IGNORECASE) for m in MODALITIES} |
|
|
|
|
| def modality_of(name: str) -> str | None: |
| stem = name.rsplit("/", 1)[-1] |
| m = _MOD_RE.search(stem) |
| tok = m.group(1).lower() if m else None |
| if tok in MODALITIES: |
| return tok |
| return next((mod for mod, rx in _MOD_ALT.items() if rx.search(stem)), None) |
|
|
|
|
| def normalize(vol: np.ndarray) -> np.ndarray: |
| """Min-max to [0, 1] between the 1st and 99.5th percentile of non-zero voxels. |
| |
| Registration leaves exact zeros outside the head, so non-zero voxels stand in for a head mask. |
| """ |
| fg = vol[vol > 0] |
| if fg.size < 1000: |
| return np.zeros_like(vol) |
| lo, hi = np.percentile(fg, (1.0, 99.5)) |
| if not np.isfinite(hi) or hi <= lo: |
| lo, hi = float(fg.min()), float(fg.max()) |
| if hi <= lo: |
| return np.zeros_like(vol) |
| return np.clip((vol - lo) / (hi - lo), 0.0, 1.0) |
|
|
|
|
| def _crop(vol: np.ndarray, geometry: str) -> np.ndarray: |
| g = GEOMETRIES[geometry] |
| offset, shape = g["offset"], g["crop"] |
| if vol.shape == SRC_SHAPE: |
| return vol[tuple(slice(o, o + c) for o, c in zip(offset, shape))] |
| if vol.shape == shape: |
| return vol |
| out = np.zeros(shape, dtype=vol.dtype) |
| src, dst = [], [] |
| for s, t in zip(vol.shape, shape): |
| o = abs(s - t) // 2 |
| src.append(slice(o, o + t) if s >= t else slice(0, s)) |
| dst.append(slice(0, t) if s >= t else slice(o, o + s)) |
| out[tuple(dst)] = vol[tuple(src)] |
| return out |
|
|
|
|
| def resample(vol: np.ndarray, geometry: str) -> np.ndarray: |
| g = GEOMETRIES[geometry] |
| v = _crop(vol, geometry) |
| d, h, w = g["out"] |
| if g["mode"] == "block": |
| return v.reshape(d, 2, h, 2, w, 2).mean(axis=(1, 3, 5)) |
| out = zoom(v, tuple(o / c for o, c in zip(g["out"], v.shape)), order=1, prefilter=False) |
| if out.shape != g["out"]: |
| pad = np.zeros(g["out"], dtype=out.dtype) |
| sl = tuple(slice(0, min(a, b)) for a, b in zip(out.shape, g["out"])) |
| pad[sl] = out[sl] |
| out = pad |
| return out |
|
|
|
|
| def _read_normalized(read, name): |
| try: |
| vol = read(name) |
| except Exception: |
| return None |
| if vol is None: |
| return None |
| if vol.ndim == 4: |
| vol = vol[..., 0] |
| return normalize(vol) if vol.ndim == 3 else None |
|
|
|
|
| def process_study(names, read, geometries=("d15",)): |
| """Series names + reader -> ({geometry: (volume uint8 [4, *out], presence uint8 [4])}, n_series_used).""" |
| cache = {} |
|
|
| def load(n): |
| if n not in cache: |
| cache[n] = _read_normalized(read, n) |
| return cache[n] |
|
|
| per_mod, unmatched, n_used = {}, [], 0 |
| for name in sorted(names): |
| mod = modality_of(name) |
| if mod is None: |
| unmatched.append(name) |
| continue |
| bucket = per_mod.setdefault(mod, []) |
| if len(bucket) < MAX_PER_MODALITY and load(name) is not None: |
| bucket.append(name) |
| n_used += 1 |
| if n_used == 0 and unmatched: |
| for name in unmatched[:MAX_PER_MODALITY]: |
| if load(name) is not None: |
| per_mod.setdefault(MODALITIES[0], []).append(name) |
| n_used += 1 |
|
|
| out = {} |
| for geo in geometries: |
| vol = np.zeros((len(MODALITIES), *GEOMETRIES[geo]["out"]), dtype=np.uint8) |
| presence = np.zeros(len(MODALITIES), dtype=np.uint8) |
| for i, mod in enumerate(MODALITIES): |
| series = [resample(cache[n], geo) for n in per_mod.get(mod, [])] |
| if series: |
| acc = series[0] if len(series) == 1 else np.mean(series, axis=0) |
| vol[i] = np.clip(np.rint(acc * 255.0), 0, 255).astype(np.uint8) |
| presence[i] = 1 |
| out[geo] = (vol, presence) |
| return out, n_used |
|
|
|
|
| def _load_nii_bytes(raw: bytes) -> np.ndarray: |
| fh = nib.FileHolder(fileobj=gzip.GzipFile(fileobj=io.BytesIO(raw))) |
| img = nib.Nifti1Image.from_file_map({"header": fh, "image": fh}) |
| return np.asanyarray(img.dataobj, dtype=np.float32) |
|
|
|
|
| def process_zip(path, geometries=("d15",)): |
| """A study zip (training layout <id>/atlas_img/, challenge layout <id>/atlas/).""" |
| with zipfile.ZipFile(path) as zf: |
| names = [n for n in zf.namelist() if n.endswith((".nii.gz", ".nii")) and "/atlas" in n |
| and "/atlas_seg/" not in n and "mask" not in n] |
| return process_study(names, lambda n: _load_nii_bytes(zf.read(n)), geometries) |
|
|
|
|
| def process_files(paths, geometries=("d15",)): |
| return process_study([str(p) for p in paths], |
| lambda p: np.asanyarray(nib.load(p).dataobj, dtype=np.float32), geometries) |
|
|
