| """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"] |
| 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", []))) |
|
|
| 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() |
|
|
