main.py

from IAD import IAD
import argparse
import os
import yaml
from easydict import EasyDict
import numpy as np
import pdb
import time


def parse_args():
    parser = argparse.ArgumentParser(
        description='Pytorch implementation of IAD')
    parser.add_argument('--config', default='')
    parser.add_argument('--dataset', default='')
    parser.add_argument('--logdir', default='')
    return parser.parse_args()


def main():
    # parse arguments and load config
    args = parse_args()
    with open(args.config) as f:
       config = yaml.safe_load(f)

    for k, v in vars(args).items():
       config[k] = v

    config["exp_name"] = args.config.split("/")[-1].split(".")[0]
    config["dataset"] = args.dataset


    if not config["eval_mode"]:
        if not config["logdir"]: 
            config["logdir"] = time.strftime('%Y-%m-%d-%H-%M')
    else:
        config["logdir"] = "2025-04-10-20-06"  #Test model


    config = EasyDict(config)
    agent = IAD(config)

    sampling = config.sampling
    steps = 5

    if config["eval_mode"]:
        agent.eval(sampling, config.num_steps//steps)
    else:
        agent.train(config.num_steps//steps)


if __name__ == '__main__':
    main()