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
def load_config(mode=None): parser = argparse.ArgumentParser() parser.add_argument("--path", "--checkpoints", type=str, default="./checkpoints", help="model checkpoint 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" ) parser.add_argument("--train_img_path", type=str, default="./train_images") parser.add_argument("--test_img_path", type=str, default="./test_images") parser.add_argument("--eval_img_path", type=str, default="./eval_images") if mode == "test": #parser.add_argument("--input", type = str, help = "path to a test image") parser.add_argument("--output", type=str, help="path to a output folder") args = parser.parse_args() create_data_list(args.train_img_path, args.test_img_path, args.eval_img_path, "./list_folder") config_path = os.path.join(args.path, "config.yaml") if not os.path.exists(args.path): os.makedirs(args.path) if not os.path.exists(config_path): copyfile('./config.yaml', config_path) 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 if args.model: config.MODEL = args.model else: config.MODEL = 3 #eval mode elif mode == 3: config.MODE = 3 if args.model: config.MODEL = args.model else: config.MODEL = 3 return config