fix --aligned-dir arg

This commit is contained in:
iperov 2018-12-20 04:14:42 +04:00
parent b2d28e0b95
commit 4ff67ad26b
2 changed files with 19 additions and 12 deletions

View file

@ -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.")

View file

@ -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) )