mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-06 13:02:15 -07:00
fix ability for auto re-extract dst faces, after this update DFL torrent extractor scripts will not work.
This commit is contained in:
iperov
parent
b03d62a29b
commit
1d56585f33
6 changed files with 48 additions and 49 deletions
|
|
@ -1,3 +1,4 @@
|
|||
import traceback
|
||||
import numpy as np
|
||||
import os
|
||||
import cv2
|
||||
|
|
@ -41,7 +42,8 @@ class LandmarksExtractor(object):
|
|||
pts_img = self.get_pts_from_predict ( predicted[-1], center, scale)
|
||||
pts_img = [ ( int(pt[0]), int(pt[1]) ) for pt in pts_img ]
|
||||
landmarks.append ( ( (left, top, right, bottom),pts_img ) )
|
||||
except:
|
||||
except Exception as e:
|
||||
print ("extract_from_bgr: ", traceback.format_exc() )
|
||||
landmarks.append ( ( (left, top, right, bottom), [ (0,0) for _ in range(68) ] ) )
|
||||
|
||||
return landmarks
|
||||
|
|
|
|||
|
|
@ -32,20 +32,16 @@ class S3FDExtractor(object):
|
|||
|
||||
olist = self.model.predict( np.expand_dims(input_image,0) )
|
||||
|
||||
detected_faces = self.refine (olist)
|
||||
|
||||
#filtering faces < 40pix by any side
|
||||
#enlarging bottom line a bit for 2DFAN-4, because default is not enough covering a chin
|
||||
new_detected_faces = []
|
||||
for ltrb in detected_faces:
|
||||
detected_faces = []
|
||||
for ltrb in self.refine (olist):
|
||||
l,t,r,b = [ x*input_scale for x in ltrb]
|
||||
bt = b-t
|
||||
if min(r-l,bt) < 40:
|
||||
if min(r-l,bt) < 40: #filtering faces < 40pix by any side
|
||||
continue
|
||||
b += bt*0.1
|
||||
new_detected_faces.append ( [int(x) for x in (l,t,r,b) ] )
|
||||
b += bt*0.1 #enlarging bottom line a bit for 2DFAN-4, because default is not enough covering a chin
|
||||
detected_faces.append ( [int(x) for x in (l,t,r,b) ] )
|
||||
|
||||
return new_detected_faces
|
||||
return detected_faces
|
||||
|
||||
def refine(self, olist):
|
||||
bboxlist = []
|
||||
|
|
|
|||
4
main.py
4
main.py
|
|
@ -26,7 +26,7 @@ if __name__ == "__main__":
|
|||
from mainscripts import Extractor
|
||||
Extractor.main( arguments.input_dir,
|
||||
arguments.output_dir,
|
||||
arguments.debug,
|
||||
arguments.debug_dir,
|
||||
arguments.detector,
|
||||
arguments.manual_fix,
|
||||
arguments.manual_output_debug_fix,
|
||||
|
|
@ -40,7 +40,7 @@ if __name__ == "__main__":
|
|||
p = subparsers.add_parser( "extract", help="Extract the faces from a pictures.")
|
||||
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
|
||||
p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the extracted files will be stored.")
|
||||
p.add_argument('--debug', action="store_true", dest="debug", default=False, help="Writes debug images to [output_dir]_debug\ directory.")
|
||||
p.add_argument('--debug-dir', action=fixPathAction, dest="debug_dir", help="Writes debug images to this directory.")
|
||||
p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'head', 'avatar', 'mark_only'], default='full_face', help="Default 'full_face'. Don't change this option, currently all models uses 'full_face'")
|
||||
p.add_argument('--detector', dest="detector", choices=['dlib','mt','s3fd','manual'], default='dlib', help="Type of detector. Default 'dlib'. 'mt' (MTCNNv1) - faster, better, almost no jitter, perfect for gathering thousands faces for src-set. It is also good for dst-set, but can generate false faces in frames where main face not recognized! In this case for dst-set use either 'dlib' with '--manual-fix' or '--detector manual'. Manual detector suitable only for dst-set.")
|
||||
p.add_argument('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="Enables multi GPU.")
|
||||
|
|
|
|||
|
|
@ -33,7 +33,7 @@ class ExtractSubprocessor(Subprocessor):
|
|||
self.device_idx = client_dict['device_idx']
|
||||
self.cpu_only = client_dict['device_type'] == 'CPU'
|
||||
self.output_path = Path(client_dict['output_dir']) if 'output_dir' in client_dict.keys() else None
|
||||
self.debug = client_dict['debug']
|
||||
self.debug_dir = client_dict['debug_dir']
|
||||
self.detector = client_dict['detector']
|
||||
|
||||
self.cached_image = (None, None)
|
||||
|
|
@ -132,8 +132,8 @@ class ExtractSubprocessor(Subprocessor):
|
|||
result = []
|
||||
faces = data[1]
|
||||
|
||||
if self.debug:
|
||||
debug_output_file = '{}{}'.format( str(Path(str(self.output_path) + '_debug') / filename_path.stem), '.jpg')
|
||||
if self.debug_dir is not None:
|
||||
debug_output_file = str( Path(self.debug_dir) / (filename_path.stem+'.jpg') )
|
||||
debug_image = image.copy()
|
||||
|
||||
face_idx = 0
|
||||
|
|
@ -157,7 +157,7 @@ class ExtractSubprocessor(Subprocessor):
|
|||
if landmarks_area > 4*rect_area: #get rid of faces which umeyama-landmark-area > 4*detector-rect-area
|
||||
continue
|
||||
|
||||
if self.debug:
|
||||
if self.debug_dir is not None:
|
||||
LandmarksProcessor.draw_rect_landmarks (debug_image, rect, image_landmarks, self.image_size, self.face_type, transparent_mask=True)
|
||||
|
||||
output_file = '{}_{}{}'.format(str(self.output_path / filename_path.stem), str(face_idx), '.jpg')
|
||||
|
|
@ -170,6 +170,7 @@ class ExtractSubprocessor(Subprocessor):
|
|||
if src_dflimg is not None:
|
||||
#if extracting from dflimg copy it in order not to lose quality
|
||||
output_file = str(self.output_path / filename_path.name)
|
||||
if str(filename_path) != str(output_file):
|
||||
shutil.copy ( str(filename_path), str(output_file) )
|
||||
|
||||
#and transfer data
|
||||
|
|
@ -191,7 +192,7 @@ class ExtractSubprocessor(Subprocessor):
|
|||
|
||||
result.append (output_file)
|
||||
|
||||
if self.debug:
|
||||
if self.debug_dir is not None:
|
||||
cv2_imwrite(debug_output_file, debug_image, [int(cv2.IMWRITE_JPEG_QUALITY), 50] )
|
||||
|
||||
return result
|
||||
|
|
@ -203,12 +204,12 @@ class ExtractSubprocessor(Subprocessor):
|
|||
return data[0]
|
||||
|
||||
#override
|
||||
def __init__(self, input_data, type, image_size, face_type, debug, 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 ):
|
||||
self.input_data = input_data
|
||||
self.type = type
|
||||
self.image_size = image_size
|
||||
self.face_type = face_type
|
||||
self.debug = debug
|
||||
self.debug_dir = debug_dir
|
||||
self.multi_gpu = multi_gpu
|
||||
self.cpu_only = cpu_only
|
||||
self.detector = detector
|
||||
|
|
@ -290,7 +291,7 @@ class ExtractSubprocessor(Subprocessor):
|
|||
base_dict = {'type' : self.type,
|
||||
'image_size': self.image_size,
|
||||
'face_type': self.face_type,
|
||||
'debug': self.debug,
|
||||
'debug_dir': self.debug_dir,
|
||||
'output_dir': str(self.output_path),
|
||||
'detector': self.detector}
|
||||
|
||||
|
|
@ -577,7 +578,7 @@ class DeletedFilesSearcherSubprocessor(Subprocessor):
|
|||
|
||||
def main(input_dir,
|
||||
output_dir,
|
||||
debug=False,
|
||||
debug_dir=None,
|
||||
detector='mt',
|
||||
manual_fix=False,
|
||||
manual_output_debug_fix=False,
|
||||
|
|
@ -597,20 +598,21 @@ def main(input_dir,
|
|||
raise ValueError('Input directory not found. Please ensure it exists.')
|
||||
|
||||
if output_path.exists():
|
||||
if not manual_output_debug_fix:
|
||||
if not manual_output_debug_fix and input_path != output_path:
|
||||
for filename in Path_utils.get_image_paths(output_path):
|
||||
Path(filename).unlink()
|
||||
else:
|
||||
output_path.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if manual_output_debug_fix:
|
||||
debug = True
|
||||
if debug_dir is None:
|
||||
raise ValueError('debug-dir must be specified')
|
||||
detector = 'manual'
|
||||
io.log_info('Performing re-extract frames which were deleted from _debug directory.')
|
||||
|
||||
input_path_image_paths = Path_utils.get_image_unique_filestem_paths(input_path, verbose_print_func=io.log_info)
|
||||
if debug:
|
||||
debug_output_path = Path(str(output_path) + '_debug')
|
||||
if debug_dir is not None:
|
||||
debug_output_path = Path(debug_dir)
|
||||
|
||||
if manual_output_debug_fix:
|
||||
if not debug_output_path.exists():
|
||||
|
|
@ -630,13 +632,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, 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).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, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, detector=detector).run()
|
||||
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, 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).run()
|
||||
|
||||
if manual_fix:
|
||||
io.log_info ('Performing manual fix...')
|
||||
|
|
@ -644,11 +646,11 @@ 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, 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).run()
|
||||
|
||||
if len(extracted_faces) > 0:
|
||||
io.log_info ('Performing 3rd pass...')
|
||||
final_imgs_paths = ExtractSubprocessor (extracted_faces, 'final', image_size, face_type, debug, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, output_path=output_path).run()
|
||||
final_imgs_paths = ExtractSubprocessor (extracted_faces, 'final', image_size, face_type, debug_dir, multi_gpu=multi_gpu, cpu_only=cpu_only, manual=False, output_path=output_path).run()
|
||||
faces_detected = len(final_imgs_paths)
|
||||
|
||||
io.log_info ('-------------------------')
|
||||
|
|
|
|||
|
|
@ -152,7 +152,8 @@ class DFLJPG(object):
|
|||
landmarks=None,
|
||||
source_filename=None,
|
||||
source_rect=None,
|
||||
source_landmarks=None
|
||||
source_landmarks=None,
|
||||
image_to_face_mat=None
|
||||
):
|
||||
|
||||
inst = DFLJPG.load_raw (filename)
|
||||
|
|
@ -161,7 +162,8 @@ class DFLJPG(object):
|
|||
'landmarks': landmarks,
|
||||
'source_filename': source_filename,
|
||||
'source_rect': source_rect,
|
||||
'source_landmarks': source_landmarks
|
||||
'source_landmarks': source_landmarks,
|
||||
'image_to_face_mat': image_to_face_mat
|
||||
})
|
||||
|
||||
try:
|
||||
|
|
@ -223,3 +225,6 @@ class DFLJPG(object):
|
|||
def get_source_filename(self): return self.dfl_dict['source_filename']
|
||||
def get_source_rect(self): return self.dfl_dict['source_rect']
|
||||
def get_source_landmarks(self): return np.array ( self.dfl_dict['source_landmarks'] )
|
||||
def get_image_to_face_mat(self):
|
||||
mat = self.dfl_dict.get('image_to_face_mat', None)
|
||||
return np.array(mat) if mat is not None else None
|
||||
|
|
@ -324,20 +324,14 @@ class DFLPNG(object):
|
|||
chunk = DFLChunk(dict_data)
|
||||
self.chunks.insert(-1, chunk)
|
||||
|
||||
def get_face_type(self):
|
||||
return self.fcwp_dict['face_type']
|
||||
|
||||
def get_landmarks(self):
|
||||
return np.array ( self.fcwp_dict['landmarks'] )
|
||||
|
||||
def get_source_filename(self):
|
||||
return self.fcwp_dict['source_filename']
|
||||
|
||||
def get_source_rect(self):
|
||||
return self.fcwp_dict['source_rect']
|
||||
|
||||
def get_source_landmarks(self):
|
||||
return np.array ( self.fcwp_dict['source_landmarks'] )
|
||||
def get_face_type(self): return self.fcwp_dict['face_type']
|
||||
def get_landmarks(self): return np.array ( self.fcwp_dict['landmarks'] )
|
||||
def get_source_filename(self): return self.fcwp_dict['source_filename']
|
||||
def get_source_rect(self): return self.fcwp_dict['source_rect']
|
||||
def get_source_landmarks(self): return np.array ( self.fcwp_dict['source_landmarks'] )
|
||||
def get_image_to_face_mat(self):
|
||||
mat = self.fcwp_dict.get('image_to_face_mat', None)
|
||||
return np.array(mat) if mat is not None else None
|
||||
|
||||
def __str__(self):
|
||||
return "<PNG length={length} chunks={}>".format(len(self.chunks), **self.__dict__)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue