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