github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-06 13:02:15 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

accurate landmarks extraction now as an option

Browse source
This commit is contained in:
iperov 2019-03-18 11:24:41 +04:00
parent bf831331e6
commit 00c6a0cdaf
1 changed files with 13 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

20
mainscripts/Extractor.py
View file

@ -35,7 +35,8 @@ class ExtractSubprocessor(Subprocessor):
self.output_path = Path(client_dict['output_dir']) if 'output_dir' in client_dict.keys() else None
self.debug_dir = client_dict['debug_dir']
self.detector = client_dict['detector']
self.accurate_landmarks_extractor = client_dict['accurate_landmarks_extractor']
self.cached_image = (None, None)
self.e = None
@ -61,7 +62,7 @@ class ExtractSubprocessor(Subprocessor):
nnlib.import_all (device_config)
self.e = facelib.LandmarksExtractor(nnlib.keras)
self.e.__enter__()
if device_config.gpu_vram_gb[0] >= 2:
if self.accurate_landmarks_extractor and device_config.gpu_vram_gb[0] >= 2:
self.second_pass_e = facelib.S3FDExtractor()
self.second_pass_e.__enter__()
else:
@ -220,7 +221,7 @@ class ExtractSubprocessor(Subprocessor):
return data[0]
#override
def __init__(self, input_data, type, image_size, face_type, debug_dir, multi_gpu=False, cpu_only=False, manual=False, manual_window_size=0, detector=None, output_path=None ):
def __init__(self, input_data, type, image_size, face_type, debug_dir, multi_gpu=False, cpu_only=False, manual=False, manual_window_size=0, detector=None, output_path=None, accurate_landmarks_extractor=False ):
self.input_data = input_data
self.type = type
self.image_size = image_size
@ -232,6 +233,7 @@ class ExtractSubprocessor(Subprocessor):
self.output_path = output_path
self.manual = manual
self.manual_window_size = manual_window_size
self.accurate_landmarks_extractor = accurate_landmarks_extractor
self.result = []
no_response_time_sec = 60 if not self.manual else 999999
@ -315,7 +317,8 @@ class ExtractSubprocessor(Subprocessor):
client_dict = base_dict.copy()
client_dict['device_idx'] = device_idx
client_dict['device_name'] = device_name
client_dict['device_type'] = device_type
client_dict['device_type'] = device_type
client_dict['accurate_landmarks_extractor'] = self.accurate_landmarks_extractor
yield client_dict['device_name'], {}, client_dict
#override
@ -610,6 +613,9 @@ def main(input_dir,
multi_gpu = device_args.get('multi_gpu', False)
cpu_only = device_args.get('cpu_only', False)
accurate_landmarks_extractor = io.input_bool("Use more accurate landmarks extraction? (y/n ?:help skip:n) : ", False, help_message="More accurate landmarks extraction, but processing is slower." )
if not input_path.exists():
raise ValueError('Input directory not found. Please ensure it exists.')
@ -648,13 +654,13 @@ def main(input_dir,
if images_found != 0:
if detector == 'manual':
io.log_info ('Performing manual extract...')
extracted_faces = ExtractSubprocessor ([ (filename,[]) for filename in input_path_image_paths ], 'landmarks', image_size, face_type, debug_dir, cpu_only=cpu_only, manual=True, manual_window_size=manual_window_size).run()
extracted_faces = ExtractSubprocessor ([ (filename,[]) for filename in input_path_image_paths ], 'landmarks', image_size, face_type, debug_dir, cpu_only=cpu_only, manual=True, manual_window_size=manual_window_size, accurate_landmarks_extractor=accurate_landmarks_extractor).run()
else:
io.log_info ('Performing 1st pass...')
extracted_rects = ExtractSubprocessor ([ (x,) for x in input_path_image_paths ], 'rects', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, detector=detector).run()
io.log_info ('Performing 2nd pass...')
extracted_faces = ExtractSubprocessor (extracted_rects, 'landmarks', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False).run()
extracted_faces = ExtractSubprocessor (extracted_rects, 'landmarks', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, accurate_landmarks_extractor=accurate_landmarks_extractor).run()
if manual_fix:
io.log_info ('Performing manual fix...')
@ -662,7 +668,7 @@ def main(input_dir,
if all ( np.array ( [ len(data[1]) > 0 for data in extracted_faces] ) == True ):
io.log_info ('All faces are detected, manual fix not needed.')
else:
extracted_faces = ExtractSubprocessor (extracted_faces, 'landmarks', image_size, face_type, debug_dir, manual=True, manual_window_size=manual_window_size).run()
extracted_faces = ExtractSubprocessor (extracted_faces, 'landmarks', image_size, face_type, debug_dir, manual=True, manual_window_size=manual_window_size, accurate_landmarks_extractor=accurate_landmarks_extractor).run()
if len(extracted_faces) > 0:
io.log_info ('Performing 3rd pass...')