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