File size: 5,868 Bytes
b360028 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | """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") # channel order is part of the model
SRC_SHAPE = (193, 229, 193)
MAX_PER_MODALITY = 4
# d15: head box cropped from the atlas, linear zoom to 1.5 mm
# d2: near-full field of view, exact 2x2x2 block average to 2 mm
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) # unexpected grid: centre crop or pad
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"]: # zoom can be off by one
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: # unrecognised names: use them as T1w rather than nothing
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)
|