| """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() |
|
|
