scripts/train_phase2.py

"""Phase 2 training entrypoint."""

from __future__ import annotations

import argparse

from dipauglib.utils.io import load_yaml
from dipaugsevernet.models.dipaugsevernet import DIPAugSeverNet
from dipaugsevernet.training.engine import fit_phase2


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Train DIPAug-SeverNet phase 2 scaffold.")
    parser.add_argument("--config", required=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    config = load_yaml(args.config)
    model = DIPAugSeverNet(num_classes=config["dataset"]["num_classes"])
    result = fit_phase2(model=model)
    print(result)


if __name__ == "__main__":
    main()