max faces from image

This commit is contained in:
Colombo 2020-06-27 19:07:55 +04:00
parent 7da28b283d
commit bb8d2b5b2c
2 changed files with 28 additions and 17 deletions

View file

@ -38,6 +38,7 @@ if __name__ == "__main__":
manual_output_debug_fix = arguments.manual_output_debug_fix,
manual_window_size = arguments.manual_window_size,
face_type = arguments.face_type,
max_faces_from_image = arguments.max_faces_from_image,
image_size = arguments.image_size,
jpeg_quality = arguments.jpeg_quality,
cpu_only = arguments.cpu_only,
@ -51,6 +52,7 @@ if __name__ == "__main__":
p.add_argument('--output-debug', action="store_true", dest="output_debug", default=None, help="Writes debug images to <output-dir>_debug\ directory.")
p.add_argument('--no-output-debug', action="store_false", dest="output_debug", default=None, help="Don't writes debug images to <output-dir>_debug\ directory.")
p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'whole_face', 'head', 'mark_only'], default=None)
p.add_argument('--max-faces-from-image', type=int, dest="max_faces_from_image", default=None, help="Max faces from image.")
p.add_argument('--image-size', type=int, dest="image_size", default=None, help="Output image size.")
p.add_argument('--jpeg-quality', type=int, dest="jpeg_quality", default=None, help="Jpeg quality.")
p.add_argument('--manual-fix', action="store_true", dest="manual_fix", default=False, help="Enables manual extract only frames where faces were not recognized.")

View file

@ -151,7 +151,7 @@ class ExtractSubprocessor(Subprocessor):
if len(rects) != 0:
data.rects_rotation = rot
break
if max_faces_from_image != 0 and len(data.rects) > 1:
if max_faces_from_image is not None and len(data.rects) > 1:
data.rects = data.rects[0:max_faces_from_image]
return data
@ -730,9 +730,9 @@ def main(detector=None,
manual_output_debug_fix=False,
manual_window_size=1368,
face_type='full_face',
max_faces_from_image=None,
image_size=None,
jpeg_quality=None,
max_faces_from_image=0,
cpu_only = False,
force_gpu_idxs = None,
):
@ -741,26 +741,34 @@ def main(detector=None,
io.log_err ('Input directory not found. Please ensure it exists.')
return
if not output_path.exists():
output_path.mkdir(parents=True, exist_ok=True)
if face_type is not None:
face_type = FaceType.fromString(face_type)
if face_type is None:
if manual_output_debug_fix and output_path.exists():
if manual_output_debug_fix:
files = pathex.get_image_paths(output_path)
if len(files) != 0:
dflimg = DFLIMG.load(Path(files[0]))
if dflimg is not None and dflimg.has_data():
face_type = FaceType.fromString ( dflimg.get_face_type() )
if output_path.exists():
if not manual_output_debug_fix and input_path != output_path:
output_images_paths = pathex.get_image_paths(output_path)
if len(output_images_paths) > 0:
input_image_paths = pathex.get_image_unique_filestem_paths(input_path, verbose_print_func=io.log_info)
output_images_paths = pathex.get_image_paths(output_path)
if not manual_output_debug_fix and len(output_images_paths) > 0:
if len(output_images_paths) > 128 and io.input_bool ("Continue extraction?", True, help_message="Extraction can be continued, but you must specify the same options again."):
try:
input_image_paths = input_image_paths[ [ Path(x).stem for x in input_image_paths ].index ( Path(output_images_paths[-128]).stem.split('_')[0] ) : ]
except:
io.log_err("Error in fetching the last index. Extraction cannot be continued.")
return
elif input_path != output_path:
io.input(f"\n WARNING !!! \n {output_path} contains files! \n They will be deleted. \n Press enter to continue.\n")
for filename in output_images_paths:
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
device_config = nn.DeviceConfig.GPUIndexes( force_gpu_idxs or nn.ask_choose_device_idxs(choose_only_one=detector=='manual', suggest_all_gpu=True) ) \
if not cpu_only else nn.DeviceConfig.CPU()
@ -771,6 +779,9 @@ def main(detector=None,
'wf' : FaceType.WHOLE_FACE,
'head' : FaceType.HEAD}[face_type]
if max_faces_from_image is None:
max_faces_from_image = io.input_int(f"Max number of faces from image", 0, help_message="If you extract a src faceset that has frames with a large number of faces, it is advisable to set max faces to 3 to speed up extraction. 0 - unlimited")
if image_size is None:
image_size = io.input_int(f"Image size", 512 if face_type < FaceType.HEAD else 768, valid_range=[256,2048], help_message="Output image size. The higher image size, the worse face-enhancer works. Use higher than 512 value only if the source image is sharp enough and the face does not need to be enhanced.")
@ -788,8 +799,6 @@ def main(detector=None,
if output_debug is None:
output_debug = io.input_bool (f"Write debug images to {output_debug_path.name}?", False)
input_path_image_paths = pathex.get_image_unique_filestem_paths(input_path, verbose_print_func=io.log_info)
if manual_output_debug_fix:
if not output_debug_path.exists():
io.log_err(f'{output_debug_path} not found. Re-extract faces with "Write debug images" option.')
@ -798,9 +807,9 @@ def main(detector=None,
detector = 'manual'
io.log_info('Performing re-extract frames which were deleted from _debug directory.')
input_path_image_paths = DeletedFilesSearcherSubprocessor (input_path_image_paths, pathex.get_image_paths(output_debug_path) ).run()
input_path_image_paths = sorted (input_path_image_paths)
io.log_info('Found %d images.' % (len(input_path_image_paths)))
input_image_paths = DeletedFilesSearcherSubprocessor (input_image_paths, pathex.get_image_paths(output_debug_path) ).run()
input_image_paths = sorted (input_image_paths)
io.log_info('Found %d images.' % (len(input_image_paths)))
else:
if output_debug_path.exists():
for filename in pathex.get_image_paths(output_debug_path):
@ -808,19 +817,19 @@ def main(detector=None,
else:
output_debug_path.mkdir(parents=True, exist_ok=True)
images_found = len(input_path_image_paths)
images_found = len(input_image_paths)
faces_detected = 0
if images_found != 0:
if detector == 'manual':
io.log_info ('Performing manual extract...')
data = ExtractSubprocessor ([ ExtractSubprocessor.Data(Path(filename)) for filename in input_path_image_paths ], 'landmarks-manual', image_size, jpeg_quality, face_type, output_debug_path if output_debug else None, manual_window_size=manual_window_size, device_config=device_config).run()
data = ExtractSubprocessor ([ ExtractSubprocessor.Data(Path(filename)) for filename in input_image_paths ], 'landmarks-manual', image_size, jpeg_quality, face_type, output_debug_path if output_debug else None, manual_window_size=manual_window_size, device_config=device_config).run()
io.log_info ('Performing 3rd pass...')
data = ExtractSubprocessor (data, 'final', image_size, jpeg_quality, face_type, output_debug_path if output_debug else None, final_output_path=output_path, device_config=device_config).run()
else:
io.log_info ('Extracting faces...')
data = ExtractSubprocessor ([ ExtractSubprocessor.Data(Path(filename)) for filename in input_path_image_paths ],
data = ExtractSubprocessor ([ ExtractSubprocessor.Data(Path(filename)) for filename in input_image_paths ],
'all',
image_size,
jpeg_quality,