File size: 2,183 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 | """CT report generation: a folder of .nii.gz volumes -> challenge report JSON.
python predict_report.py --weights ct_classification \
--writer ct_report_generation/report_writer.json \
--input scans/ --output reports.json
"""
from __future__ import annotations
import argparse
import json
import re
from pathlib import Path
import numpy as np
import torch
from fusion import FusionPolicy
from predict_classification import case_name, load_models, predict_logits
from report_writer import ReportWriter
UUID = re.compile(r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}", re.I)
def report_name(path: Path) -> str:
"""The report track names a case by the trailing UUID in its file name, if there is one."""
stem = case_name(path)
return next((part for part in reversed(stem.split("_")) if UUID.fullmatch(part)), stem)
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--weights", type=Path, required=True)
ap.add_argument("--writer", 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")
labels, models, cfgs = load_models(args.weights, device)
policy = FusionPolicy(args.weights / "fusion_policy.json")
writer = ReportWriter(args.writer)
assert policy.labels == labels == writer.labels
paths = sorted(args.input.rglob("*.nii.gz"))
names = [report_name(p) for p in paths]
assert len(set(names)) == len(names), "duplicate case names"
logits = np.stack([predict_logits(p, models, cfgs, device) for p in paths])
_, emitted = policy(logits)
reports = [{"input_image_name": n, "report": writer.render(row)} for n, row in zip(names, emitted)]
out = [{"outputs": [{"value": {"name": "Generated reports", "type": "Report generation",
"version": "1.0", "generated_reports": reports}}]}]
args.output.parent.mkdir(parents=True, exist_ok=True)
args.output.write_text(json.dumps(out, ensure_ascii=False))
if __name__ == "__main__":
main()
|