scripts/01_prepare.py

"""Pre-flight checks: official assets, fold hygiene, evaluator verification.

Run this before anything else. It fails loudly if the official LUNA16 protocol
cannot be reproduced on this machine.

Usage
-----
    python scripts/01_prepare.py
"""

from __future__ import annotations

import json
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parents[1]))

from luna_rev import config as cfg
from luna_rev import evaluate, splits
from luna_rev.io_luna import all_uids, load_annotations, load_excluded, verify_scans_present


def main() -> int:
    ok = True
    report: dict = {}

    print("=" * 72)
    print("1) Official assets")
    print("=" * 72)
    for label, path in [("annotations.csv", cfg.ANNOTATIONS_CSV),
                        ("annotations_excluded.csv", cfg.ANNOTATIONS_EXCLUDED_CSV),
                        ("seriesuids.csv", cfg.SERIESUIDS_CSV)]:
        exists = path.exists()
        ok &= exists
        print(f"  {'OK  ' if exists else 'FAIL'} {label:26s} {path}")
    ann, excl = load_annotations(), load_excluded()
    uids = all_uids()
    report["n_scans"] = len(uids)
    report["n_nodules"] = len(ann)
    report["n_excluded_findings"] = len(excl)
    # Evaluation needs only the CSVs, but the pipeline as a whole needs pixels.
    verify_scans_present()
    print(f"  scans={len(uids)}  reference nodules={len(ann)}  irrelevant findings={len(excl)}")
    ok &= (len(uids) == 888 and len(ann) == 1186)

    print("\n" + "=" * 72)
    print("2) Official subset0-9 folds")
    print("=" * 72)
    subsets = splits.official_subsets()
    print("  subset sizes: " + ", ".join(f"{k}:{len(v)}" for k, v in sorted(subsets.items())))
    audit = splits.audit("official")
    report["fold_audit"] = audit
    ok &= audit["clean"]
    print(f"  covers all scans      : {audit['test_partition_covers_all_scans']}")
    print(f"  test partitions disjoint: {audit['test_partitions_disjoint']}")
    print(f"  train/val/test overlaps : "
          f"{max(audit['train_test_overlap'].values())}, "
          f"{max(audit['val_test_overlap'].values())}, "
          f"{max(audit['train_val_overlap'].values())} (all must be 0)")
    f0 = splits.official_folds()[0]
    print(f"  fold0 -> train {len(f0.train)}, val {len(f0.val)}, test {len(f0.test)}")
    splits.export(protocol="official")
    splits.export(protocol="legacy")

    print("\n" + "=" * 72)
    print("3) Evaluator verification against the official reference output")
    print("=" * 72)
    ref = evaluate.validate_against_reference("legacy_abs")
    report["evaluator_verification"] = ref
    ok &= ref["matches_reference"]
    print(f"  reproduces CADAnalysis.txt exactly: {ref['matches_reference']}")
    for k, v in ref["counters"].items():
        if isinstance(v, int):
            print(f"    {k:28s} {v}")
    if ref["mismatches"]:
        print(f"  MISMATCHES: {ref['mismatches']}")

    print("\n" + "=" * 72)
    print("4) Experiment matrix")
    print("=" * 72)
    total_runs = 0
    for group in dict.fromkeys(e.group for e in cfg.ALL_EXPERIMENTS):
        exps = cfg.experiments_for_group(group)
        n = sum(len(cfg.folds_for(e)) for e in exps)
        total_runs += n
        print(f"  {group:14s} {len(exps):2d} configs x folds -> {n:3d} training runs")
    report["total_training_runs"] = total_runs
    print(f"  {'TOTAL':14s} {len(cfg.ALL_EXPERIMENTS):2d} configs           -> {total_runs:3d} training runs")
    print(f"  imgsz={cfg.IMG_SIZE}  epochs={cfg.SCHEDULE.epochs}  batch={cfg.HW.train_batch}  "
          f"negatives/scan={cfg.NEG.per_scan}")

    out = cfg.RESULTS_DIR / "prepare_report.json"
    out.write_text(json.dumps(report, indent=1, default=str), encoding="utf-8")
    print(f"\nWrote {out}")
    print("\nRESULT:", "PASS" if ok else "FAIL")
    return 0 if ok else 1


if __name__ == "__main__":
    raise SystemExit(main())