diff --git a/facelib/LandmarksExtractor.py b/facelib/LandmarksExtractor.py index 9f5d6bd..810c81a 100644 --- a/facelib/LandmarksExtractor.py +++ b/facelib/LandmarksExtractor.py @@ -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 diff --git a/facelib/S3FDExtractor.py b/facelib/S3FDExtractor.py index 961b589..2054cee 100644 --- a/facelib/S3FDExtractor.py +++ b/facelib/S3FDExtractor.py @@ -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 = [] diff --git a/main.py b/main.py index 8ba5b8c..8dee298 100644 --- a/main.py +++ b/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.") diff --git a/mainscripts/Extractor.py b/mainscripts/Extractor.py index 47cf25e..e9b78bd 100644 --- a/mainscripts/Extractor.py +++ b/mainscripts/Extractor.py @@ -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,7 +170,8 @@ 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) - shutil.copy ( str(filename_path), str(output_file) ) + if str(filename_path) != str(output_file): + shutil.copy ( str(filename_path), str(output_file) ) #and transfer data source_filename = src_dflimg.get_source_filename() @@ -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 ('-------------------------') diff --git a/utils/DFLJPG.py b/utils/DFLJPG.py index 1716392..9e96807 100644 --- a/utils/DFLJPG.py +++ b/utils/DFLJPG.py @@ -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 \ No newline at end of file diff --git a/utils/DFLPNG.py b/utils/DFLPNG.py index 4fe827f..93b19bc 100644 --- a/utils/DFLPNG.py +++ b/utils/DFLPNG.py @@ -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 "".format(len(self.chunks), **self.__dict__)