Spaces:
Sleeping
Sleeping
File size: 3,861 Bytes
f559cc0 | 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 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | """
Train the archive conjunctiva screening model.
Usage::
python scripts/train_archive_model.py [--dataset PATH] [--output-dir PATH] [--quiet]
The script trains the model, writes the artefact and a human-readable
training report, then exits with code 0 on success or 1 on failure.
"""
from __future__ import annotations
import argparse
import json
import sys
import time
from pathlib import Path
ROOT = Path(__file__).resolve().parents[2]
BACKEND_ROOT = ROOT / "backend"
sys.path.insert(0, str(BACKEND_ROOT))
from app.config import DEFAULT_ARCHIVE_MODEL_PATH, DEFAULT_TRAINING_REPORT_PATH # noqa: E402
from app.ml.archive_model import save_archive_model, train_archive_model # noqa: E402
DEFAULT_DATASET = ROOT / "archive" / "dataset anemia"
# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------
def _build_parser() -> argparse.ArgumentParser:
p = argparse.ArgumentParser(
description="Train the AnemiaLens archive conjunctiva screening model.",
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
)
p.add_argument(
"--dataset",
type=Path,
default=DEFAULT_DATASET,
help="Root directory of the labelled anemia dataset.",
)
p.add_argument(
"--output-dir",
type=Path,
default=DEFAULT_ARCHIVE_MODEL_PATH.parent,
help="Directory where the model artefact and report are written.",
)
p.add_argument(
"--quiet",
action="store_true",
help="Suppress progress output (report still written to disk).",
)
return p
# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def main(argv: list[str] | None = None) -> int:
args = _build_parser().parse_args(argv)
if not args.dataset.exists():
print(
f"ERROR: Dataset directory not found: {args.dataset}\n"
" Download the anemia dataset and place it there, or pass --dataset PATH.",
file=sys.stderr,
)
return 1
if not args.quiet:
print(f"Dataset : {args.dataset}")
print(f"Output : {args.output_dir}")
print()
t0 = time.perf_counter()
try:
artifact, report = train_archive_model(args.dataset)
except Exception as exc:
print(f"ERROR: Training failed — {exc}", file=sys.stderr)
return 1
elapsed = time.perf_counter() - t0
# --- Write artefacts ---------------------------------------------------
args.output_dir.mkdir(parents=True, exist_ok=True)
model_path = args.output_dir / DEFAULT_ARCHIVE_MODEL_PATH.name
report_path = args.output_dir / DEFAULT_TRAINING_REPORT_PATH.name
save_archive_model(artifact, model_path)
report_path.write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
# --- Summary -----------------------------------------------------------
if not args.quiet:
metrics = report.get("metrics", {})
print(json.dumps(report, indent=2))
print()
print("=" * 56)
print(f" Model : {report.get('primary_model', '?')}")
print(f" Subjects : {report.get('subject_count', '?')}")
print(f" Records : {report.get('record_count', '?')}")
print(f" Accuracy : {metrics.get('accuracy', 0):.3f}")
print(f" F1 : {metrics.get('f1', 0):.3f}")
print(f" Val size : {metrics.get('validation_size', '?')}")
print(f" Elapsed : {elapsed:.1f}s")
print("=" * 56)
print(f" Saved model → {model_path}")
print(f" Saved report → {report_path}")
return 0
if __name__ == "__main__":
sys.exit(main())
|