File size: 3,145 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 | """MR report generation: a folder of study zips -> predictions.json with one report per study.
python predict_report.py --weights mr_report_generation \
--input studies/ --output reports.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 implications import closure_matrix, select_with_implications
from model import build_model
from predict_classification import discover
from preprocess import process_files, process_zip
from report import ReportWriter
from tensors import to_model_input
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 = json.loads((args.weights / "config.json").read_text())
rc = json.loads((args.weights / "report_config.json").read_text())
net = build_model(cfg["arch"], in_ch=cfg["in_channels"], num_classes=len(cfg["labels"]))
net.load_state_dict(load_file(str(args.weights / "model.safetensors")))
net = net.to(device).eval()
order = rc["order"] # the 32 scored diagnoses
idx = [cfg["labels"].index(c) for c in order]
thr = np.array([rc["thresholds"][c] for c in order])
closure = closure_matrix(order)
writer = ReportWriter(rc["template_lines"], rc["bullets"], order,
header_findings=rc.get("header_findings", True),
header_impression=rc.get("header_impression", True),
normal_line=rc.get("normal_line", ""),
short_template=rc.get("short_template_lines"),
short_from=int(rc.get("short_from", 0)))
fallback = writer.write(set(rc.get("fallback_positives", []))) # for a study that cannot be decoded
reports = []
for i, (name, study) in enumerate(discover(args.input), 1):
geo = cfg["geometry"]
per, n_used = process_zip(study, (geo,)) if isinstance(study, Path) else process_files(study, (geo,))
if n_used == 0:
text = fallback
else:
vol, presence = per[geo]
x = to_model_input(torch.from_numpy(vol)[None].to(device),
torch.from_numpy(presence.astype(np.float32))[None].to(device))
with torch.no_grad():
probs = torch.sigmoid(net(x).float())[0].cpu().numpy().astype(np.float64)[idx]
chosen = select_with_implications(probs[None], thr, int(rc.get("max_bullets", 0)), closure)[0]
text = writer.write({c for c, k in zip(order, chosen) if k})
reports.append({"input_image_name": f"{name}.nii.gz", "report": text})
print(f"[{i}] {name}", flush=True)
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps({"generated_reports": reports}))
if __name__ == "__main__":
main()
|