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
