scripts/02_build_datasets.py

"""Render the shared slice pools and every label variant.

Single HDD pass over the 888 LUNA16 volumes: all three input representations
are written while each volume is resident, then the ``(r, w_min)`` label
variants are materialised with hard-linked images.

Usage
-----
    python scripts/02_build_datasets.py [--overwrite] [--workers N]
"""

from __future__ import annotations

import argparse
import json
import sys
import time
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 dataset_builder as db


def required_variants() -> list[tuple[str, float, float]]:
    """Every ``(representation, r, w_min)`` triple the experiment matrix needs."""
    seen, out = set(), []
    for exp in cfg.ALL_EXPERIMENTS:
        key = (exp.representation, exp.r_sample, exp.w_min_px)
        if key not in seen:
            seen.add(key)
            out.append(key)
    return out


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--overwrite", action="store_true", help="re-render existing slices/labels")
    ap.add_argument("--workers", type=int, default=cfg.HW.build_workers)
    ap.add_argument("--skip-pools", action="store_true", help="only rebuild label variants")
    args = ap.parse_args()

    t0 = time.time()
    if args.skip_pools:
        manifest = db.load_manifest()
        print(f"[pools] reusing manifest with {len(manifest)} scans")
    else:
        print(f"[pools] rendering {len(cfg.REPRESENTATIONS)} representations "
              f"with {args.workers} workers ...")
        manifest = db.build_pools(overwrite=args.overwrite, workers=args.workers)

    n_pos = sum(len(e["positive"]) for e in manifest.values())
    n_neg = sum(len(e["negative"]) for e in manifest.values())
    n_nodule_free = sum(1 for e in manifest.values() if not e["positive"])
    print(f"[pools] scans={len(manifest)} nodule-free={n_nodule_free} "
          f"positive_slices={n_pos} negative_slices={n_neg} total={n_pos + n_neg} "
          f"({time.time() - t0:.0f}s)")

    stats = []
    for rep, r, w in required_variants():
        t1 = time.time()
        s = db.build_variant(rep, r, w, manifest=manifest, overwrite=args.overwrite)
        s["seconds"] = round(time.time() - t1, 1)
        stats.append(s)
        print(f"[variant] {s['variant']:26s} images={s['n_images']:6d} "
              f"boxes={s['n_boxes']:6d} clamped_at_w_min={s['frac_clamped'] * 100:5.1f}% "
              f"({s['seconds']}s)")

    out = cfg.RESULTS_DIR / "dataset_stats.json"
    out.write_text(json.dumps({
        "n_scans": len(manifest),
        "n_nodule_free_scans": n_nodule_free,
        "n_positive_slices": n_pos,
        "n_negative_slices": n_neg,
        "negatives_per_scan": cfg.NEG.per_scan,
        "img_size": cfg.IMG_SIZE,
        "variants": stats,
    }, indent=1), encoding="utf-8")
    print(f"\nWrote {out}  (total {time.time() - t0:.0f}s)")
    return 0


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