TrainLUT.py

"""TrainLUT - root entrypoint for LUT training."""
from __future__ import annotations

import argparse
from typing import Sequence

from LUT_Train import Driver


def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="启动 ColorMatch LUT 训练流程"
    )
    warmup_group = parser.add_mutually_exclusive_group()
    warmup_group.add_argument(
        "--warmup",
        dest="use_warmup",
        action="store_true",
        help="启用 PSNR 预热（默认值）",
    )
    warmup_group.add_argument(
        "--no-warmup",
        dest="use_warmup",
        action="store_false",
        help="跳过 PSNR 预热以加快启动速度",
    )
    parser.set_defaults(use_warmup=True)
    return parser.parse_args(args=argv)


def launch_training(use_warmup: bool = True) -> None:
    driver = Driver(use_warmup=use_warmup)
    driver.main()


def main(argv: Sequence[str] | None = None) -> None:
    args = parse_args(argv)
    launch_training(use_warmup=args.use_warmup)


if __name__ == "__main__":
    main()