File size: 3,183 Bytes
c3f98a1 | 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | """Run the detector over one CT volume and emit a LUNA16-format candidate CSV.
This is the whole-volume path: every axial slice is scanned, slice detections
are aggregated into 3D candidates, and centres are written in world coordinates
so the output can be fed straight to the official evaluator.
Usage
-----
python examples/predict_scan.py \\
--scan /data/LUNA16/subset0/1.3.6.1.4...mhd \\
--weights weights/Exp4_2p5D_Strict/fold0/best.pt \\
--out candidates.csv
Note on fold choice: fold *k* was trained without ``subset<k>``, so use the
checkpoint whose fold matches the subset the scan came from. For a scan from
outside LUNA16, any fold is valid; averaging folds is not implemented here.
"""
from __future__ import annotations
import argparse
import sys
from pathlib import Path
import pandas as pd
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
from luna_rev import config as cfg
from luna_rev.io_luna import normalize_hu, read_volume, scan_index
from luna_rev.predict import candidates_to_world, cluster_to_3d, detect_volume
def main() -> int:
ap = argparse.ArgumentParser()
ap.add_argument("--scan", required=True, help="path to a .mhd volume, or a bare series UID")
ap.add_argument("--weights", required=True)
ap.add_argument("--out", default="candidates.csv")
ap.add_argument("--representation", default="naive", choices=list(cfg.REPRESENTATIONS),
help="must match how the checkpoint was trained")
ap.add_argument("--top-k", type=int, default=150)
ap.add_argument("--imgsz", type=int, default=cfg.IMG_SIZE)
args = ap.parse_args()
from ultralytics import YOLO
path = Path(args.scan)
uid = path.stem if path.suffix else args.scan
if not path.exists():
path = scan_index()[uid]
vol_hu, meta = read_volume(uid) if path == scan_index().get(uid) else _read_direct(path, uid)
vol_u8 = normalize_hu(vol_hu)
print(f"{uid}: {vol_u8.shape[0]} slices, spacing {meta.spacing.round(3).tolist()} mm")
model = YOLO(args.weights)
det = detect_volume(model, vol_u8, args.representation)
clusters = cluster_to_3d(det)[:args.top_k]
world = candidates_to_world(clusters, meta)
print(f"{len(det)} slice detections -> {len(clusters)} 3D candidates")
df = pd.DataFrame({
"seriesuid": uid,
"coordX": world[:, 0], "coordY": world[:, 1], "coordZ": world[:, 2],
"probability": clusters[:, 5],
})
df.to_csv(args.out, index=False)
print(f"wrote {args.out}")
return 0
def _read_direct(path: Path, uid: str):
"""Read a volume that is not part of the indexed LUNA16 tree."""
import numpy as np
import SimpleITK as sitk
from luna_rev.io_luna import ScanMeta
img = sitk.ReadImage(str(path))
meta = ScanMeta(
uid=uid,
size=np.array(img.GetSize(), dtype=int),
spacing=np.array(img.GetSpacing(), dtype=float),
origin=np.array(img.GetOrigin(), dtype=float),
direction=np.array(img.GetDirection(), dtype=float).reshape(3, 3),
)
return sitk.GetArrayFromImage(img), meta
if __name__ == "__main__":
raise SystemExit(main())
|