File size: 3,592 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 | """CT abnormality classification: a folder of .nii.gz volumes -> predictions JSON.
python predict_classification.py --weights ct_classification \
--input scans/ --output results.json
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path
import nibabel as nib
import numpy as np
import torch
from safetensors.torch import load_file
from fusion import FusionPolicy
from model import build_from_config, preprocess_ct
def case_name(path: Path) -> str:
return path.name[:-len(".nii.gz")] if path.name.endswith(".nii.gz") else path.stem
def load_models(weights_dir: Path, device):
spec = json.loads((weights_dir / "streams.json").read_text())
models, cfgs = [], []
for stream in spec["streams"]:
sd = load_file(str(weights_dir / stream["weights_file"]))
models.append(build_from_config(stream["config"], sd, len(spec["labels"])).to(device))
cfgs.append(stream["config"])
return spec["labels"], models, cfgs
@torch.inference_mode()
def predict_logits(path: Path, models, cfgs, device) -> np.ndarray:
"""(K, M) raw logits of the M models for one volume."""
img = nib.load(str(path))
volume = np.asanyarray(img.dataobj).astype(np.float32)
affine = np.asarray(img.affine, dtype=np.float32)
grids, out = {}, []
for model, cfg in zip(models, cfgs):
key = (cfg["img_size"], cfg["xy_spacing"], cfg["z_spacing"], cfg["max_slices"])
if key not in grids: # the two 1.5 mm models share one grid
grids[key] = preprocess_ct(volume, affine, img_size=cfg["img_size"], xy_spacing=cfg["xy_spacing"],
z_spacing=cfg["z_spacing"], max_slices=cfg["max_slices"])
x = grids[key].unsqueeze(0).to(device)
with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"):
out.append(model(x).float().cpu().numpy()[0])
return np.stack(out, axis=-1)
def challenge_json(names, labels, probs):
return [{"outputs": [{"value": {
"name": "Generated probabilities", "type": "Abnormality Classification",
"version": {"major": 1, "minor": 0},
"predictions": [{"input_image_name": n,
"probabilities": {l: float(p) for l, p in zip(labels, row)}}
for n, row in zip(names, probs)]}}]}]
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)
ap.add_argument("--save-logits", type=Path, help="optional .npz with raw logits per case")
args = ap.parse_args()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
labels, models, cfgs = load_models(args.weights, device)
policy = FusionPolicy(args.weights / "fusion_policy.json")
assert policy.labels == labels
paths = sorted(args.input.glob("*.nii.gz"))
names = [case_name(p) for p in paths]
logits = []
for i, p in enumerate(paths, 1):
logits.append(predict_logits(p, models, cfgs, device))
print(f"[{i}/{len(paths)}] {names[i - 1]}", flush=True)
logits = np.stack(logits)
_, emitted = policy(logits)
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(challenge_json(names, labels, emitted)))
if args.save_logits:
np.savez(args.save_logits, names=np.array(names), logits=logits, emitted=emitted)
if __name__ == "__main__":
main()
|