examples/predict_scan.py

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