From 4ff67ad26bfda428d0d247e1e86cc3cd42c1351c Mon Sep 17 00:00:00 2001 From: iperov Date: Thu, 20 Dec 2018 04:14:42 +0400 Subject: [PATCH] fix --aligned-dir arg --- main.py | 4 +--- mainscripts/Converter.py | 27 ++++++++++++++++++--------- 2 files changed, 19 insertions(+), 12 deletions(-) diff --git a/main.py b/main.py index f2f9fc8..37dfe6b 100644 --- a/main.py +++ b/main.py @@ -166,8 +166,6 @@ if __name__ == "__main__": arguments.alpha = bool ( {"1":True,"0":False}[input("Export png with alpha channel? [0..1] (default 0) : ").lower()] ) except: arguments.alpha = False - - arguments.erode_mask_modifier = np.clip ( int(arguments.erode_mask_modifier), -200, 200) arguments.blur_mask_modifier = np.clip ( int(arguments.blur_mask_modifier), -200, 200) @@ -197,7 +195,7 @@ if __name__ == "__main__": convert_parser = subparsers.add_parser( "convert", help="Converter") convert_parser.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") convert_parser.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the converted files will be stored.") - convert_parser.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", help="Aligned directory. This is where the aligned files stored. Not used in AVATAR model.") + convert_parser.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", help="Aligned directory. This is where the extracted of dst faces stored. Not used in AVATAR model.") convert_parser.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.") convert_parser.add_argument('--model', required=True, dest="model_name", choices=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Type of model") convert_parser.add_argument('--ask-for-params', action="store_true", dest="ask_for_params", default=False, help="Ask for params.") diff --git a/mainscripts/Converter.py b/mainscripts/Converter.py index cda413f..9bd092b 100644 --- a/mainscripts/Converter.py +++ b/mainscripts/Converter.py @@ -202,15 +202,14 @@ class ConvertSubprocessor(SubprocessorBase): def get_start_return(self): return self.files_processed, self.faces_processed -def main (input_dir, output_dir, aligned_dir, model_dir, model_name, **in_options): +def main (input_dir, output_dir, model_dir, model_name, aligned_dir=None, **in_options): print ("Running converter.\r\n") debug = in_options['debug'] try: input_path = Path(input_dir) - output_path = Path(output_dir) - aligned_path = Path(aligned_dir) + output_path = Path(output_dir) model_path = Path(model_dir) if not input_path.exists(): @@ -223,9 +222,7 @@ def main (input_dir, output_dir, aligned_dir, model_dir, model_name, **in_option else: output_path.mkdir(parents=True, exist_ok=True) - if not aligned_path.exists(): - print('Aligned directory not found. Please ensure it exists.') - return + if not model_path.exists(): print('Model directory not found. Please ensure it exists.') @@ -244,9 +241,21 @@ def main (input_dir, output_dir, aligned_dir, model_dir, model_name, **in_option if obj_op == 'init': converter = obj['converter'] break - - alignments = {} - if converter.get_mode() == ConverterBase.MODE_FACE: + + alignments = None + + if converter.get_mode() == ConverterBase.MODE_FACE: + if aligned_dir is None: + print('Aligned directory not found. Please ensure it exists.') + return + + aligned_path = Path(aligned_dir) + if not aligned_path.exists(): + print('Aligned directory not found. Please ensure it exists.') + return + + alignments = {} + aligned_path_image_paths = Path_utils.get_image_paths(aligned_path) for filename in tqdm(aligned_path_image_paths, desc= "Collecting alignments" ): a_png = AlignedPNG.load( str(filename) )