File size: 4,571 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 | """MR abnormality classification: a folder of study zips -> predictions JSON.
python predict_classification.py --weights mr_classification \
--input studies/ --output results.json
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import numpy as np
import torch
from safetensors.torch import load_file
from model import build_model
from preprocess import process_files, process_zip
from tensors import to_model_input
def discover(root: Path):
"""Study zips, or pre-extracted <study>/**/atlas*/*.nii.gz folders."""
zips = sorted(root.glob("*.zip"))
if zips:
return [(z.stem, z) for z in zips]
out = []
for d in sorted(p for p in root.iterdir() if p.is_dir()):
vols = sorted(p for p in d.rglob("*.nii.gz")
if "/atlas" in str(p) and "/atlas_seg/" not in str(p) and "mask" not in str(p))
if vols:
out.append((d.name, vols))
return out
def load_ensemble(weights_dir: Path, device):
cfg = json.loads((weights_dir / "config.json").read_text())
models = []
for m in cfg["members"]:
net = build_model(m["arch"], in_ch=m["in_channels"], num_classes=len(cfg["labels"]))
net.load_state_dict(load_file(str(weights_dir / m["weights_file"])))
models.append((net.to(device).eval(), m["geometry"]))
return cfg, models
@torch.no_grad()
def ensemble_probabilities(study, models, device) -> np.ndarray | None:
"""Mean over members of sigmoid(mean of logits for the volume and its left-right mirror)."""
geos = sorted({g for _, g in models})
per, n_used = process_zip(study, geos) if isinstance(study, Path) else process_files(study, geos)
if n_used == 0:
return None
inputs = {}
for g, (vol, presence) in per.items():
x = to_model_input(torch.from_numpy(vol)[None].to(device),
torch.from_numpy(presence.astype(np.float32))[None].to(device))
inputs[g] = (x, torch.flip(x, dims=[2]))
acc = 0.0
for net, g in models:
x, xf = inputs[g]
with torch.autocast("cuda", dtype=torch.bfloat16, enabled=device.type == "cuda"):
logits = (net(x).float() + net(xf).float()) / 2.0
acc = acc + torch.sigmoid(logits)[0].cpu().numpy().astype(np.float64)
return acc / len(models)
def move_threshold_to_half(p: np.ndarray, t: np.ndarray) -> np.ndarray:
p = np.clip(p, 1e-9, 1 - 1e-9)
t = np.clip(t, 1e-6, 1 - 1e-6)
num = p * (1.0 - t)
return np.clip(num / np.maximum(num + (1.0 - p) * t, 1e-15), 1e-9, 1 - 1e-9)
def apply_positive_rates(probs: np.ndarray, rates: np.ndarray) -> np.ndarray:
n = probs.shape[0]
thr = np.empty(probs.shape[1])
for c in range(probs.shape[1]):
k = max(1, min(n - 1, int(round(float(rates[c]) * n))))
s = np.sort(probs[:, c])[::-1]
thr[c] = (s[k - 1] + s[k]) / 2.0
return move_threshold_to_half(probs, np.clip(thr, 1e-6, 1 - 1e-6))
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--weights", type=Path, required=True)
ap.add_argument("--input", type=Path, required=True)
ap.add_argument("--output", type=Path, required=True)
args = ap.parse_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
cfg, models = load_ensemble(args.weights, device)
labels = cfg["labels"]
rates = np.array([cfg["positive_rates"][l] for l in labels])
names, rows, failed = [], [], []
for i, (name, study) in enumerate(discover(args.input), 1):
p = ensemble_probabilities(study, models, device)
names.append(f"{name}.nii.gz")
rows.append(p)
if p is None:
failed.append(len(rows) - 1)
print(f"[{i}] {name}{' (no usable series)' if p is None else ''}", flush=True)
real = [i for i in range(len(rows)) if rows[i] is not None]
if len(real) < 2:
raise RuntimeError("need at least 2 decoded studies")
probs = np.full((len(rows), len(labels)), 0.5)
probs[real] = apply_positive_rates(np.stack([rows[i] for i in real]), rates)
preds = [{"input_image_name": n, "probabilities": {l: float(v) for l, v in zip(labels, row)}}
for n, row in zip(names, probs)]
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps([{"outputs": [{"value": {
"name": "Generated probabilities", "type": "Abnormality Classification",
"version": {"major": 1, "minor": 0}, "predictions": preds}}]}]))
if __name__ == "__main__":
main()
|