scripts/main.py

"""
main.py
========
اسکریپت اصلی برای اجرای کل pipeline از خط فرمان:
    python scripts/main.py --mode train
    python scripts/main.py --mode evaluate --ckpt models/saved_models/omnial_epoch50.pt
"""

import os
import sys
import argparse
import yaml
import torch

sys.path.append(os.path.join(os.path.dirname(__file__), ".."))
from data.data_loader import build_dataloader
from scripts.train import build_model_and_optimizer, train_model
from scripts.evaluate import evaluate_dataset


def load_config(path: str = "config/config.yaml") -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return yaml.safe_load(f)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--mode", choices=["train", "evaluate"], default="train")
    parser.add_argument("--dataset", choices=["mvtec", "visa"], default="mvtec")
    parser.add_argument("--config", default="config/config.yaml")
    parser.add_argument("--ckpt", default=None)
    args = parser.parse_args()

    cfg = load_config(args.config)
    torch.manual_seed(cfg["train"]["seed"])

    model, optimizer, criterion, device = build_model_and_optimizer(cfg)
    print(f"Using device: {device}")

    if args.mode == "train":
        train_loader, train_dataset = build_dataloader(cfg, dataset_name=args.dataset, shuffle=True)
        print(f"تعداد نمونه‌های آموزشی ({args.dataset}): {len(train_dataset)}")
        history = train_model(model, train_loader, optimizer, criterion, device, cfg)
        print("آموزش تمام شد.")

    elif args.mode == "evaluate":
        if args.ckpt:
            model.load_state_dict(torch.load(args.ckpt, map_location=device))
            print(f"وزن‌ها از {args.ckpt} بارگذاری شد.")
        evaluate_dataset(model, cfg, dataset_name=args.dataset, device=device)


if __name__ == "__main__":
    main()