def load_config(mode=None): r"""loads model config """ 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, 4], help= '1: edge model, 2: inpaint model, 3: edge-inpaint model, 4: joint model' ) # test mode parser.add_argument( '--input', type=str, help='path to the input images directory or an input image') parser.add_argument('--edge', type=str, help='path to the edges directory or an edge file') parser.add_argument('--output', type=str, help='path to the output directory') parser.add_argument('--remove', nargs='*', type=int, help='objects to remove') parser.add_argument('--cpu', type=str, help='machine to run segmentation model on') args = parser.parse_args() #if path for checkpoint not given if args.path is None: args.path = './checkpoints' 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.example', config_path) # load config file config = Config(config_path) # test mode config.MODE = 2 config.MODEL = args.model if args.model is not None else 3 config.OBJECTS = args.remove if args.remove is not None else [3, 15] config.SEG_DEVICE = 'cpu' if args.cpu is not None else 'cuda' config.INPUT_SIZE = 256 if args.input is not None: config.TEST_FLIST = args.input if args.edge is not None: config.TEST_EDGE_FLIST = args.edge if args.output is not None: config.RESULTS = args.output else: if not os.path.exists('./results_images'): os.makedirs('./results_images') config.RESULTS = './results_images' return config