def load_config(mode=None):
r"""loads model config
Args:
mode (int): 1: train, 2: test, 3: eval, reads from config file if not specified
"""
parser = argparse.ArgumentParser()
parser.add_argument('--path', '--checkpoints', type=str, default='./checkpoints', help='model checkpoints path (default: ./checkpoints)')
parser.add_argument('--model', type=int, choices=[1, 2, 3], help='1: edge model, 2: SR model, 3: joint SR model with edge enhancer')
#import pdb;pdb.set_trace()
# test mode
if mode == 2:
parser.add_argument('--input', type=str, help='path to the input images directory or an input image')
parser.add_argument('--output', type=str, help='path to the output directory')
args = parser.parse_args()
config_path = os.path.join(args.path, 'config.yml')
# create checkpoints path if does't exist
if not os.path.exists(args.path):
os.makedirs(args.path)
# copy config template if does't exist
if not os.path.exists(config_path):
copyfile('./config.yml', config_path)
# load config file
config = Config(config_path)
# train mode
if mode == 1:
config.MODE = 1
if args.model:
config.MODEL = args.model
# test mode
elif mode == 2:
config.MODE = 2
config.HR_SIZE = 0
config.MODEL = args.model if args.model is not None else 3
if args.input is not None:
config.TEST_FLIST_LR = args.input
if args.output is not None:
config.RESULTS = args.output
# eval mode
elif mode == 3:
config.MODE = 3
config.MODEL = args.model if args.model is not None else 3
return config
