From b72d5a3f9aef4637ca18ae09a98aedcf694eb850 Mon Sep 17 00:00:00 2001 From: iperov Date: Sun, 11 Aug 2019 11:17:22 +0400 Subject: [PATCH 1/6] fixed error "Failed to get convolution algorithm" on some systems fixed error "dll load failed" on some systems Expanded eyebrows line of face masks. It does not affect mask of FAN-x converter mode. --- converters/ConverterAvatar.py | 45 ++--- facelib/FaceType.py | 24 +-- facelib/LandmarksProcessor.py | 122 +++++++------ main.py | 4 +- mainscripts/Converter.py | 199 ++++++++++++---------- mainscripts/Extractor.py | 21 ++- mainscripts/Sorter.py | 11 +- models/ModelBase.py | 2 +- models/Model_RecycleGAN/Model.py | 1 - models/Model_SAE/Model.py | 6 + nnlib/nnlib.py | 84 +++++++++ samplelib/SampleGeneratorBase.py | 13 +- samplelib/SampleGeneratorFaceTemporal.py | 8 +- samplelib/SampleGeneratorImageTemporal.py | 5 +- samplelib/SampleProcessor.py | 44 +++-- 15 files changed, 367 insertions(+), 222 deletions(-) diff --git a/converters/ConverterAvatar.py b/converters/ConverterAvatar.py index 4a9e9b5..ca26c63 100644 --- a/converters/ConverterAvatar.py +++ b/converters/ConverterAvatar.py @@ -20,8 +20,9 @@ class ConverterAvatar(Converter): self.predictor_input_size = predictor_input_size #dummy predict and sleep, tensorflow caching kernels. If remove it, conversion speed will be x2 slower - predictor_func ( np.zeros ( (predictor_input_size,predictor_input_size,3), dtype=np.float32 ), - np.zeros ( (predictor_input_size,predictor_input_size,1), dtype=np.float32 ) ) + predictor_func ( np.zeros ( (predictor_input_size,predictor_input_size,3), dtype=np.float32 ), + np.zeros ( (predictor_input_size,predictor_input_size,3), dtype=np.float32 ), + np.zeros ( (predictor_input_size,predictor_input_size,3), dtype=np.float32 ) ) time.sleep(2) predictor_func_host, predictor_func = SubprocessFunctionCaller.make_pair(predictor_func) @@ -33,38 +34,28 @@ class ConverterAvatar(Converter): self.predictor_func_host.obj.process_messages() #override - def cli_convert_face (self, img_bgr, img_face_landmarks, debug, avaperator_face_bgr=None, **kwargs): + def cli_convert_face (self, f0, f0_lmrk, f1, f1_lmrk, f2, f2_lmrk, debug, **kwargs): if debug: - debugs = [img_bgr.copy()] - - img_size = img_bgr.shape[1], img_bgr.shape[0] - - img_face_mask_a = LandmarksProcessor.get_image_hull_mask (img_bgr.shape, img_face_landmarks) - img_face_mask_aaa = np.repeat(img_face_mask_a, 3, -1) + debugs = [] + + inp_size = self.predictor_input_size + + f0_mat = LandmarksProcessor.get_transform_mat (f0_lmrk, inp_size, face_type=FaceType.FULL_NO_ALIGN) + f1_mat = LandmarksProcessor.get_transform_mat (f1_lmrk, inp_size, face_type=FaceType.FULL_NO_ALIGN) + f2_mat = LandmarksProcessor.get_transform_mat (f2_lmrk, inp_size, face_type=FaceType.FULL_NO_ALIGN) - output_size = self.predictor_input_size - face_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, output_size, face_type=FaceType.FULL) + inp_f0 = cv2.warpAffine( f0, f0_mat, (inp_size, inp_size), flags=cv2.INTER_CUBIC ) + inp_f1 = cv2.warpAffine( f1, f1_mat, (inp_size, inp_size), flags=cv2.INTER_CUBIC ) + inp_f2 = cv2.warpAffine( f2, f2_mat, (inp_size, inp_size), flags=cv2.INTER_CUBIC ) - dst_face_mask_a_0 = cv2.warpAffine( img_face_mask_a, face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC ) + prd_f = self.predictor_func ( inp_f0, inp_f1, inp_f2 ) - predictor_input_dst_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (self.predictor_input_size,self.predictor_input_size), cv2.INTER_CUBIC ) - prd_inp_dst_face_mask_a = predictor_input_dst_face_mask_a_0[...,np.newaxis] + out_img = np.clip(prd_f, 0.0, 1.0) - prd_inp_avaperator_face_bgr = cv2.resize (avaperator_face_bgr, (self.predictor_input_size,self.predictor_input_size), cv2.INTER_CUBIC ) - - prd_face_bgr = self.predictor_func ( prd_inp_avaperator_face_bgr, prd_inp_dst_face_mask_a ) - - out_img = img_bgr.copy() - out_img = cv2.warpAffine( prd_face_bgr, face_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) - out_img = np.clip(out_img, 0.0, 1.0) + out_img = np.concatenate ( [cv2.resize ( inp_f1, (prd_f.shape[1], prd_f.shape[0]) ), + out_img], axis=1 ) if debug: debugs += [out_img.copy()] - out_img = np.clip( img_bgr*(1-img_face_mask_aaa) + (out_img*img_face_mask_aaa) , 0, 1.0 ) - - if debug: - debugs += [out_img.copy()] - - return debugs if debug else out_img diff --git a/facelib/FaceType.py b/facelib/FaceType.py index f0d5530..dba999d 100644 --- a/facelib/FaceType.py +++ b/facelib/FaceType.py @@ -4,10 +4,10 @@ class FaceType(IntEnum): HALF = 0, FULL = 1, HEAD = 2, - AVATAR = 3, #centered nose only - MARK_ONLY = 4, #no align at all, just embedded faceinfo - QTY = 5 - + + FULL_NO_ALIGN = 5, + MARK_ONLY = 10, #no align at all, just embedded faceinfo + @staticmethod def fromString (s): r = from_string_dict.get (s.lower()) @@ -17,17 +17,17 @@ class FaceType(IntEnum): @staticmethod def toString (face_type): - return to_string_list[face_type] + return to_string_dict[face_type] from_string_dict = {'half_face': FaceType.HALF, 'full_face': FaceType.FULL, 'head' : FaceType.HEAD, - 'avatar' : FaceType.AVATAR, 'mark_only' : FaceType.MARK_ONLY, + 'full_face_no_align' : FaceType.FULL_NO_ALIGN, } -to_string_list = [ 'half_face', - 'full_face', - 'head', - 'avatar', - 'mark_only' - ] +to_string_dict = { FaceType.HALF : 'half_face', + FaceType.FULL : 'full_face', + FaceType.HEAD : 'head', + FaceType.MARK_ONLY :'mark_only', + FaceType.FULL_NO_ALIGN : 'full_face_no_align' + } diff --git a/facelib/LandmarksProcessor.py b/facelib/LandmarksProcessor.py index a49076b..7c0b697 100644 --- a/facelib/LandmarksProcessor.py +++ b/facelib/LandmarksProcessor.py @@ -109,10 +109,19 @@ landmarks_68_3D = np.array( [ [0.205322 , 31.408738 , -21.903670 ], [-7.198266 , 30.844876 , -20.328022 ] ], dtype=np.float32) +def transform_points(points, mat, invert=False): + if invert: + mat = cv2.invertAffineTransform (mat) + points = np.expand_dims(points, axis=1) + points = cv2.transform(points, mat, points.shape) + points = np.squeeze(points) + return points + def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0): if not isinstance(image_landmarks, np.ndarray): image_landmarks = np.array (image_landmarks) + """ if face_type == FaceType.AVATAR: centroid = np.mean (image_landmarks, axis=0) @@ -128,76 +137,79 @@ def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0): mat = mat * scale * (output_size / 3) mat[:,2] += output_size / 2 else: - if face_type == FaceType.HALF: - padding = 0 - elif face_type == FaceType.FULL: - padding = (output_size / 64) * 12 - elif face_type == FaceType.HEAD: - padding = (output_size / 64) * 24 - else: - raise ValueError ('wrong face_type: ', face_type) + """ + remove_align = False + if face_type == FaceType.FULL_NO_ALIGN: + face_type = FaceType.FULL + remove_align = True + + if face_type == FaceType.HALF: + padding = 0 + elif face_type == FaceType.FULL: + padding = (output_size / 64) * 12 + elif face_type == FaceType.HEAD: + padding = (output_size / 64) * 24 + else: + raise ValueError ('wrong face_type: ', face_type) - mat = umeyama(image_landmarks[17:], landmarks_2D, True)[0:2] - mat = mat * (output_size - 2 * padding) - mat[:,2] += padding - mat *= (1 / scale) - mat[:,2] += -output_size*( ( (1 / scale) - 1.0 ) / 2 ) + mat = umeyama(image_landmarks[17:], landmarks_2D, True)[0:2] + mat = mat * (output_size - 2 * padding) + mat[:,2] += padding + mat *= (1 / scale) + mat[:,2] += -output_size*( ( (1 / scale) - 1.0 ) / 2 ) + + if remove_align: + bbox = transform_points ( [ (0,0), (0,output_size-1), (output_size-1, output_size-1), (output_size-1,0) ], mat, True) + area = mathlib.polygon_area(bbox[:,0], bbox[:,1] ) + side = math.sqrt(area) / 2 + center = transform_points ( [(output_size/2,output_size/2)], mat, True) + + pts1 = np.float32([ center+[-side,-side], center+[side,-side], center+[-side,side] ]) + pts2 = np.float32([[0,0],[output_size-1,0],[0,output_size-1]]) + mat = cv2.getAffineTransform(pts1,pts2) return mat -def transform_points(points, mat, invert=False): - if invert: - mat = cv2.invertAffineTransform (mat) - points = np.expand_dims(points, axis=1) - points = cv2.transform(points, mat, points.shape) - points = np.squeeze(points) - return points - - def get_image_hull_mask (image_shape, image_landmarks, ie_polys=None): if len(image_landmarks) != 68: raise Exception('get_image_hull_mask works only with 68 landmarks') - int_lmrks = np.array(image_landmarks, dtype=np.int) + int_lmrks = np.array(image_landmarks.copy(), dtype=np.int) hull_mask = np.zeros(image_shape[0:2]+(1,),dtype=np.float32) - cv2.fillConvexPoly( hull_mask, cv2.convexHull( - np.concatenate ( (int_lmrks[0:9], - int_lmrks[17:18]))) , (1,) ) + # #nose + ml_pnt = (int_lmrks[36] + int_lmrks[0]) // 2 + mr_pnt = (int_lmrks[16] + int_lmrks[45]) // 2 - cv2.fillConvexPoly( hull_mask, cv2.convexHull( - np.concatenate ( (int_lmrks[8:17], - int_lmrks[26:27]))) , (1,) ) + # mid points between the mid points and eye + ql_pnt = (int_lmrks[36] + ml_pnt) // 2 + qr_pnt = (int_lmrks[45] + mr_pnt) // 2 - cv2.fillConvexPoly( hull_mask, cv2.convexHull( - np.concatenate ( (int_lmrks[17:20], - int_lmrks[8:9]))) , (1,) ) + # Top of the eye arrays + bot_l = np.array((ql_pnt, int_lmrks[36], int_lmrks[37], int_lmrks[38], int_lmrks[39])) + bot_r = np.array((int_lmrks[42], int_lmrks[43], int_lmrks[44], int_lmrks[45], qr_pnt)) - cv2.fillConvexPoly( hull_mask, cv2.convexHull( - np.concatenate ( (int_lmrks[24:27], - int_lmrks[8:9]))) , (1,) ) + # Eyebrow arrays + top_l = int_lmrks[17:22] + top_r = int_lmrks[22:27] - cv2.fillConvexPoly( hull_mask, cv2.convexHull( - np.concatenate ( (int_lmrks[19:25], - int_lmrks[8:9], - ))) , (1,) ) + # Adjust eyebrow arrays + int_lmrks[17:22] = top_l + ((top_l - bot_l) // 2) + int_lmrks[22:27] = top_r + ((top_r - bot_r) // 2) - cv2.fillConvexPoly( hull_mask, cv2.convexHull( - np.concatenate ( (int_lmrks[17:22], - int_lmrks[27:28], - int_lmrks[31:36], - int_lmrks[8:9] - ))) , (1,) ) + r_jaw = (int_lmrks[0:9], int_lmrks[17:18]) + l_jaw = (int_lmrks[8:17], int_lmrks[26:27]) + r_cheek = (int_lmrks[17:20], int_lmrks[8:9]) + l_cheek = (int_lmrks[24:27], int_lmrks[8:9]) + nose_ridge = (int_lmrks[19:25], int_lmrks[8:9],) + r_eye = (int_lmrks[17:22], int_lmrks[27:28], int_lmrks[31:36], int_lmrks[8:9]) + l_eye = (int_lmrks[22:27], int_lmrks[27:28], int_lmrks[31:36], int_lmrks[8:9]) + nose = (int_lmrks[27:31], int_lmrks[31:36]) + parts = [r_jaw, l_jaw, r_cheek, l_cheek, nose_ridge, r_eye, l_eye, nose] - cv2.fillConvexPoly( hull_mask, cv2.convexHull( - np.concatenate ( (int_lmrks[22:27], - int_lmrks[27:28], - int_lmrks[31:36], - int_lmrks[8:9] - ))) , (1,) ) - - #nose - cv2.fillConvexPoly( hull_mask, cv2.convexHull(int_lmrks[27:36]), (1,) ) + for item in parts: + merged = np.concatenate(item) + cv2.fillConvexPoly(hull_mask, cv2.convexHull(merged), 1) if ie_polys is not None: ie_polys.overlay_mask(hull_mask) @@ -309,7 +321,7 @@ def draw_landmarks (image, image_landmarks, color=(0,255,0), transparent_mask=Fa mask = get_image_hull_mask (image.shape, image_landmarks, ie_polys) image[...] = ( image * (1-mask) + image * mask / 2 )[...] -def draw_rect_landmarks (image, rect, image_landmarks, face_size, face_type, transparent_mask=False, ie_polys=None, landmarks_color=(0,255,0) ): +def draw_rect_landmarks (image, rect, image_landmarks, face_size, face_type, transparent_mask=False, ie_polys=None, landmarks_color=(0,255,0)): draw_landmarks(image, image_landmarks, color=landmarks_color, transparent_mask=transparent_mask, ie_polys=ie_polys) imagelib.draw_rect (image, rect, (255,0,0), 2 ) diff --git a/main.py b/main.py index 8093c77..240c044 100644 --- a/main.py +++ b/main.py @@ -40,7 +40,7 @@ if __name__ == "__main__": 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-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('--face-type', dest="face_type", choices=['half_face', 'full_face', 'head', 'full_face_no_align', '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.") p.add_argument('--manual-fix', action="store_true", dest="manual_fix", default=False, help="Enables manual extract only frames where faces were not recognized.") @@ -151,7 +151,6 @@ if __name__ == "__main__": args = {'input_dir' : arguments.input_dir, 'output_dir' : arguments.output_dir, 'aligned_dir' : arguments.aligned_dir, - 'avaperator_aligned_dir' : arguments.avaperator_aligned_dir, 'model_dir' : arguments.model_dir, 'model_name' : arguments.model_name, 'debug' : arguments.debug, @@ -166,7 +165,6 @@ if __name__ == "__main__": 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 converted files will be stored.") p.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", help="Aligned directory. This is where the extracted of dst faces stored.") - p.add_argument('--avaperator-aligned-dir', action=fixPathAction, dest="avaperator_aligned_dir", help="Only for AVATAR model. Directory of aligned avatar operator faces.") p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.") p.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") p.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug converter.") diff --git a/mainscripts/Converter.py b/mainscripts/Converter.py index 5ce9b7a..53c3e4f 100644 --- a/mainscripts/Converter.py +++ b/mainscripts/Converter.py @@ -28,9 +28,9 @@ class ConvertSubprocessor(Subprocessor): self.device_idx = client_dict['device_idx'] self.device_name = client_dict['device_name'] self.converter = client_dict['converter'] + self.input_data = client_dict['input_data'] self.output_path = Path(client_dict['output_dir']) if 'output_dir' in client_dict.keys() else None self.alignments = client_dict['alignments'] - self.avatar_image_paths = client_dict['avatar_image_paths'] self.debug = client_dict['debug'] #transfer and set stdin in order to work code.interact in debug subprocess @@ -50,65 +50,30 @@ class ConvertSubprocessor(Subprocessor): #override def process_data(self, data): - idx, filename = data - filename_path = Path(filename) files_processed = 1 faces_processed = 0 + + idx, = data + filename = self.input_data[idx][0] + filename_path = Path(filename) output_filename_path = self.output_path / (filename_path.stem + '.png') + image = None + + if self.converter.type == Converter.TYPE_FACE: + if filename_path.stem not in self.alignments.keys(): + if not self.debug: + self.log_info ( 'no faces found for %s, copying without faces' % (filename_path.name) ) - if (self.converter.type == Converter.TYPE_FACE or self.converter.type == Converter.TYPE_FACE_AVATAR ) \ - and filename_path.stem not in self.alignments.keys(): - if not self.debug: - self.log_info ( 'no faces found for %s, copying without faces' % (filename_path.name) ) - - if filename_path.suffix == '.png': - shutil.copy ( str(filename_path), str(output_filename_path) ) - else: - image = cv2_imread(str(filename_path)) - cv2_imwrite ( str(output_filename_path), image ) - else: - image = (cv2_imread(str(filename_path)) / 255.0).astype(np.float32) - image = normalize_channels (image, 3) + if filename_path.suffix == '.png': + shutil.copy ( str(filename_path), str(output_filename_path) ) + else: + image = cv2_imread(str(filename_path)) + cv2_imwrite ( str(output_filename_path), image ) + else: + image = (cv2_imread(str(filename_path)) / 255.0).astype(np.float32) + image = normalize_channels (image, 3) - if self.converter.type == Converter.TYPE_IMAGE: - image = self.converter.cli_convert_image(image, None, self.debug) - - if self.debug: - return (1, image) - - faces_processed = 1 - - elif self.converter.type == Converter.TYPE_IMAGE_WITH_LANDMARKS: - #currently unused - if filename_path.suffix == '.png': - dflimg = DFLPNG.load( str(filename_path) ) - elif filename_path.suffix == '.jpg': - dflimg = DFLJPG.load ( str(filename_path) ) - else: - dflimg = None - - if dflimg is not None: - image_landmarks = dflimg.get_landmarks() - - image = self.converter.convert_image(image, image_landmarks, self.debug) - - if self.debug: - raise NotImplementedError - #for img in image: - # io.show_image ('Debug convert', img ) - # cv2.waitKey(0) - faces_processed = 1 - else: - self.log_err ("%s is not a dfl image file" % (filename_path.name) ) - - elif self.converter.type == Converter.TYPE_FACE or self.converter.type == Converter.TYPE_FACE_AVATAR: - - ava_face = None - if self.converter.type == Converter.TYPE_FACE_AVATAR: - ava_filename_path = self.avatar_image_paths[idx] - ava_face = (cv2_imread(str(ava_filename_path)) / 255.0).astype(np.float32) - ava_face = normalize_channels (ava_face, 3) faces = self.alignments[filename_path.stem] if self.debug: @@ -120,9 +85,9 @@ class ConvertSubprocessor(Subprocessor): self.log_info ( '\nConverting face_num [%d] in file [%s]' % (face_num, filename_path) ) if self.debug: - debug_images += self.converter.cli_convert_face(image, image_landmarks, self.debug, avaperator_face_bgr=ava_face) + debug_images += self.converter.cli_convert_face(image, image_landmarks, self.debug) else: - image = self.converter.cli_convert_face(image, image_landmarks, self.debug, avaperator_face_bgr=ava_face) + image = self.converter.cli_convert_face(image, image_landmarks, self.debug) except Exception as e: e_str = traceback.format_exc() @@ -135,29 +100,60 @@ class ConvertSubprocessor(Subprocessor): return (1, debug_images) faces_processed = len(faces) + elif self.converter.type == Converter.TYPE_IMAGE: + image = (cv2_imread(str(filename_path)) / 255.0).astype(np.float32) + image = normalize_channels (image, 3) + image = self.converter.cli_convert_image(image, None, self.debug) - if not self.debug: - cv2_imwrite (str(output_filename_path), (image*255).astype(np.uint8) ) + if self.debug: + return (1, image) + + faces_processed = 1 + elif self.converter.type == Converter.TYPE_FACE_AVATAR: + max_idx = len(self.input_data)-1 + + i0 = max (idx-1, 0) + i1 = idx + i2 = min (max_idx, idx+1) + + f0 = (cv2_imread( self.input_data[i0][0] ) / 255.0).astype(np.float32) + f0_lmrk = self.input_data[i0][1] + f1 = (cv2_imread( self.input_data[i1][0] ) / 255.0).astype(np.float32) + f1_lmrk = self.input_data[i1][1] + f2 = (cv2_imread( self.input_data[i2][0] ) / 255.0).astype(np.float32) + f2_lmrk = self.input_data[i2][1] + + f0, f1, f2 = [ normalize_channels (f, 3) for f in [f0,f1,f2] ] + + image = self.converter.cli_convert_face(f0, f0_lmrk, f1, f1_lmrk, f2, f2_lmrk, self.debug) + + output_filename_path = self.output_path / self.input_data[idx][2] + if self.debug: + return (1, image) + + faces_processed = 1 + + if image is not None and not self.debug: + cv2_imwrite (str(output_filename_path), (image*255).astype(np.uint8) ) return (0, files_processed, faces_processed) #overridable def get_data_name (self, data): #return string identificator of your data - idx, filename = data - return filename + idx, = data + return self.input_data[idx][0] #override - def __init__(self, converter, input_path_image_paths, output_path, alignments, avatar_image_paths=None, debug = False): + def __init__(self, converter, input_data, output_path, alignments, debug = False): super().__init__('Converter', ConvertSubprocessor.Cli, 86400 if debug == True else 60) self.converter = converter - self.input_data = self.input_path_image_paths = input_path_image_paths + self.input_data = input_data self.input_data_idxs = [ *range(len(self.input_data)) ] self.output_path = output_path self.alignments = alignments - self.avatar_image_paths = avatar_image_paths self.debug = debug self.files_processed = 0 @@ -171,9 +167,9 @@ class ConvertSubprocessor(Subprocessor): yield 'CPU%d' % (i), {}, {'device_idx': i, 'device_name': 'CPU%d' % (i), 'converter' : self.converter, + 'input_data' : self.input_data, 'output_dir' : str(self.output_path), 'alignments' : self.alignments, - 'avatar_image_paths' : self.avatar_image_paths, 'debug': self.debug, 'stdin_fd': sys.stdin.fileno() if self.debug else None } @@ -196,12 +192,12 @@ class ConvertSubprocessor(Subprocessor): def get_data(self, host_dict): if len (self.input_data_idxs) > 0: idx = self.input_data_idxs.pop(0) - return (idx, self.input_data[idx]) + return (idx, ) return None #override def on_data_return (self, host_dict, data): - idx, filename = data + idx, = data self.input_data_idxs.insert(0, idx) #override @@ -253,9 +249,9 @@ def main (args, device_args): converter = model.get_converter() input_path_image_paths = Path_utils.get_image_paths(input_path) + alignments = None - avatar_image_paths = None - if converter.type == Converter.TYPE_FACE or converter.type == Converter.TYPE_FACE_AVATAR: + if converter.type == Converter.TYPE_FACE: if aligned_dir is None: io.log_err('Aligned directory not found. Please ensure it exists.') return @@ -287,21 +283,15 @@ def main (args, device_args): alignments[ source_filename_stem ] = [] alignments[ source_filename_stem ].append (dflimg.get_source_landmarks()) - - - if converter.type == Converter.TYPE_FACE_AVATAR: - if avaperator_aligned_dir is None: - io.log_err('Avatar operator aligned directory not found. Please ensure it exists.') - return + #avatar_alignments += [ ( str(filepath), dflimg.get_source_landmarks(), dflimg.get_source_filename() ) ] + + input_data = [ (p,) for p in input_path_image_paths ] + elif converter.type == Converter.TYPE_FACE_AVATAR: + + input_data = [] + for filepath in io.progress_bar_generator(input_path_image_paths, "Collecting info"): + filepath = Path(filepath) - avaperator_aligned_path = Path(avaperator_aligned_dir) - if not avaperator_aligned_path.exists(): - io.log_err('Avatar operator aligned directory not found. Please ensure it exists.') - return - - avatar_image_paths = [] - for filename in io.progress_bar_generator( Path_utils.get_image_paths(avaperator_aligned_path) , "Sorting avaperator faces"): - filepath = Path(filename) if filepath.suffix == '.png': dflimg = DFLPNG.load( str(filepath) ) elif filepath.suffix == '.jpg': @@ -310,22 +300,19 @@ def main (args, device_args): dflimg = None if dflimg is None: - io.log_err ("Fatal error: %s is not a dfl image file" % (filepath.name) ) - return - - avatar_image_paths += [ (filename, dflimg.get_source_filename() ) ] - avatar_image_paths = [ p[0] for p in sorted(avatar_image_paths, key=operator.itemgetter(1)) ] + io.log_err ("%s is not a dfl image file" % (filepath.name) ) + continue + input_data += [ ( str(filepath), dflimg.get_landmarks(), dflimg.get_source_filename() ) ] - if len(input_path_image_paths) < len(avatar_image_paths): - io.log_err("Input faces count must be >= avatar operator faces count.") - return - + input_data = sorted(input_data, key=operator.itemgetter(2)) + else: + input_data = [ (p,) for p in input_path_image_paths ] + files_processed, faces_processed = ConvertSubprocessor ( converter = converter, - input_path_image_paths = input_path_image_paths, + input_data = input_data, output_path = output_path, alignments = alignments, - avatar_image_paths = avatar_image_paths, debug = args.get('debug',False) ).run() @@ -389,3 +376,29 @@ if model_name == 'AVATAR': # new_points = np.concatenate( [np.expand_dims(p1,-1),np.expand_dims(p2,-1)], -1 ) # # alignments[ a[i] ][0] = LandmarksProcessor.transform_points (new_points, m0, True).astype(np.int32) + +""" + elif self.converter.type == Converter.TYPE_IMAGE_WITH_LANDMARKS: + #currently unused + if filename_path.suffix == '.png': + dflimg = DFLPNG.load( str(filename_path) ) + elif filename_path.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filename_path) ) + else: + dflimg = None + + if dflimg is not None: + image_landmarks = dflimg.get_landmarks() + + image = self.converter.convert_image(image, image_landmarks, self.debug) + + if self.debug: + raise NotImplementedError + #for img in image: + # io.show_image ('Debug convert', img ) + # cv2.waitKey(0) + faces_processed = 1 + else: + self.log_err ("%s is not a dfl image file" % (filename_path.name) ) + +""" \ No newline at end of file diff --git a/mainscripts/Extractor.py b/mainscripts/Extractor.py index bb9873a..a45d59b 100644 --- a/mainscripts/Extractor.py +++ b/mainscripts/Extractor.py @@ -6,6 +6,7 @@ import multiprocessing import shutil from pathlib import Path import numpy as np +import math import mathlib import imagelib import cv2 @@ -21,6 +22,8 @@ from nnlib import nnlib from joblib import Subprocessor from interact import interact as io +DEBUG = False + class ExtractSubprocessor(Subprocessor): class Data(object): def __init__(self, filename=None, rects=None, landmarks = None, landmarks_accurate=True, pitch_yaw_roll=None, final_output_files = None): @@ -44,6 +47,11 @@ class ExtractSubprocessor(Subprocessor): self.cpu_only = client_dict['device_type'] == 'CPU' self.final_output_path = Path(client_dict['final_output_dir']) if 'final_output_dir' in client_dict.keys() else None self.debug_dir = client_dict['debug_dir'] + + #transfer and set stdin in order to work code.interact in debug subprocess + stdin_fd = client_dict['stdin_fd'] + if stdin_fd is not None and DEBUG: + sys.stdin = os.fdopen(stdin_fd) self.cached_image = (None, None) @@ -224,10 +232,12 @@ class ExtractSubprocessor(Subprocessor): rect = np.array(rect) if self.face_type == FaceType.MARK_ONLY: + image_to_face_mat = None face_image = image face_image_landmarks = image_landmarks else: image_to_face_mat = LandmarksProcessor.get_transform_mat (image_landmarks, self.image_size, self.face_type) + face_image = cv2.warpAffine(image, image_to_face_mat, (self.image_size, self.image_size), cv2.INTER_LANCZOS4) face_image_landmarks = LandmarksProcessor.transform_points (image_landmarks, image_to_face_mat) @@ -239,8 +249,8 @@ 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_dir is not None: - LandmarksProcessor.draw_rect_landmarks (debug_image, rect, image_landmarks, self.image_size, self.face_type, transparent_mask=True) + 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) if src_dflimg is not None and filename_path.suffix == '.jpg': #if extracting from dflimg and jpg copy it in order not to lose quality @@ -296,7 +306,7 @@ class ExtractSubprocessor(Subprocessor): self.devices = ExtractSubprocessor.get_devices_for_config(self.manual, self.type, multi_gpu, cpu_only) - no_response_time_sec = 60 if not self.manual else 999999 + no_response_time_sec = 60 if not self.manual and not DEBUG else 999999 super().__init__('Extractor', ExtractSubprocessor.Cli, no_response_time_sec) #override @@ -342,7 +352,8 @@ class ExtractSubprocessor(Subprocessor): 'image_size': self.image_size, 'face_type': self.face_type, 'debug_dir': self.debug_dir, - 'final_output_dir': str(self.final_output_path)} + 'final_output_dir': str(self.final_output_path), + 'stdin_fd': sys.stdin.fileno() } for (device_idx, device_type, device_name, device_total_vram_gb) in self.devices: @@ -620,7 +631,7 @@ class ExtractSubprocessor(Subprocessor): return [ (i, 'CPU', 'CPU%d' % (i), 0 ) for i in range( min(8, multiprocessing.cpu_count() // 2) ) ] elif type == 'final': - return [ (i, 'CPU', 'CPU%d' % (i), 0 ) for i in range(min(8, multiprocessing.cpu_count())) ] + return [ (i, 'CPU', 'CPU%d' % (i), 0 ) for i in (range(min(8, multiprocessing.cpu_count())) if not DEBUG else [0]) ] class DeletedFilesSearcherSubprocessor(Subprocessor): class Cli(Subprocessor.Cli): diff --git a/mainscripts/Sorter.py b/mainscripts/Sorter.py index 9ff2588..8ff0ee9 100644 --- a/mainscripts/Sorter.py +++ b/mainscripts/Sorter.py @@ -423,14 +423,11 @@ def sort_by_hist_dissim(input_path): else: dflimg = None - if dflimg is None: - io.log_err ("%s is not a dfl image file" % (filepath.name) ) - trash_img_list.append ([str(filepath)]) - continue - image = cv2_imread(str(filepath)) - face_mask = LandmarksProcessor.get_image_hull_mask (image.shape, dflimg.get_landmarks()) - image = (image*face_mask).astype(np.uint8) + + if dflimg is not None: + face_mask = LandmarksProcessor.get_image_hull_mask (image.shape, dflimg.get_landmarks()) + image = (image*face_mask).astype(np.uint8) img_list.append ([str(filepath), cv2.calcHist([cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)], [0], None, [256], [0, 256]), 0 ]) diff --git a/models/ModelBase.py b/models/ModelBase.py index b823a74..0897dd2 100644 --- a/models/ModelBase.py +++ b/models/ModelBase.py @@ -466,7 +466,7 @@ class ModelBase(object): return imagelib.equalize_and_stack_square (images) def generate_next_sample(self): - return [next(generator) for generator in self.generator_list] + return [ generator.generate_next() for generator in self.generator_list] def train_one_iter(self): sample = self.generate_next_sample() diff --git a/models/Model_RecycleGAN/Model.py b/models/Model_RecycleGAN/Model.py index 6f9d1c8..9533ab1 100644 --- a/models/Model_RecycleGAN/Model.py +++ b/models/Model_RecycleGAN/Model.py @@ -378,7 +378,6 @@ class RecycleGANModel(ModelBase): return x return func - nnlib.UNet = UNet @staticmethod def UNetTemporalPredictor(output_nc, use_batch_norm, ngf=64, use_dropout=False): diff --git a/models/Model_SAE/Model.py b/models/Model_SAE/Model.py index f4fd2b6..dc2f0a2 100644 --- a/models/Model_SAE/Model.py +++ b/models/Model_SAE/Model.py @@ -545,6 +545,12 @@ class SAEModel(ModelBase): return Norm(norm)( Act(act) (Conv2D(dim, kernel_size=5, strides=2, padding=padding)(x)) ) return func SAEModel.downscale = downscale + + #def downscale (dim, padding='zero', norm='', act='', **kwargs): + # def func(x): + # return BlurPool()( Norm(norm)( Act(act) (Conv2D(dim, kernel_size=5, strides=1, padding=padding)(x)) ) ) + # return func + #SAEModel.downscale = downscale def upscale (dim, padding='zero', norm='', act='', **kwargs): def func(x): diff --git a/nnlib/nnlib.py b/nnlib/nnlib.py index edcb201..97e98b9 100644 --- a/nnlib/nnlib.py +++ b/nnlib/nnlib.py @@ -89,6 +89,8 @@ dssim = nnlib.dssim PixelShuffler = nnlib.PixelShuffler SubpixelUpscaler = nnlib.SubpixelUpscaler Scale = nnlib.Scale +BlurPool = nnlib.BlurPool +SelfAttention = nnlib.SelfAttention CAInitializerMP = nnlib.CAInitializerMP @@ -455,6 +457,51 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator nnlib.PixelShuffler = PixelShuffler nnlib.SubpixelUpscaler = PixelShuffler + + class BlurPool(KL.Layer): + """ + https://arxiv.org/abs/1904.11486 https://github.com/adobe/antialiased-cnns + """ + def __init__(self, filt_size=3, stride=2, **kwargs): + self.strides = (stride,stride) + self.filt_size = filt_size + self.padding = ( (int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)) ), (int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)) ) ) + if(self.filt_size==1): + self.a = np.array([1.,]) + elif(self.filt_size==2): + self.a = np.array([1., 1.]) + elif(self.filt_size==3): + self.a = np.array([1., 2., 1.]) + elif(self.filt_size==4): + self.a = np.array([1., 3., 3., 1.]) + elif(self.filt_size==5): + self.a = np.array([1., 4., 6., 4., 1.]) + elif(self.filt_size==6): + self.a = np.array([1., 5., 10., 10., 5., 1.]) + elif(self.filt_size==7): + self.a = np.array([1., 6., 15., 20., 15., 6., 1.]) + + super(BlurPool, self).__init__(**kwargs) + + def compute_output_shape(self, input_shape): + height = input_shape[1] // self.strides[0] + width = input_shape[2] // self.strides[1] + channels = input_shape[3] + return (input_shape[0], height, width, channels) + + def call(self, x): + k = self.a + k = k[:,None]*k[None,:] + k = k / np.sum(k) + k = np.tile (k[:,:,None,None], (1,1,K.int_shape(x)[-1],1) ) + k = K.constant (k, dtype=K.floatx() ) + + x = K.spatial_2d_padding(x, padding=self.padding) + x = K.depthwise_conv2d(x, k, strides=self.strides, padding='valid') + return x + + nnlib.BlurPool = BlurPool + class Scale(KL.Layer): """ @@ -487,6 +534,43 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator return dict(list(base_config.items()) + list(config.items())) nnlib.Scale = Scale + class SelfAttention(KL.Layer): + def __init__(self, nc, squeeze_factor=8, **kwargs): + assert nc//squeeze_factor > 0, f"Input channels must be >= {squeeze_factor}, recieved nc={nc}" + + self.nc = nc + self.squeeze_factor = squeeze_factor + super(SelfAttention, self).__init__(**kwargs) + + def compute_output_shape(self, input_shape): + return (input_shape[0], input_shape[1], input_shape[2], self.nc) + + def call(self, inp): + x = inp + shape_x = x.get_shape().as_list() + + f = Conv2D(self.nc//self.squeeze_factor, 1, kernel_regularizer=keras.regularizers.l2(1e-4))(x) + g = Conv2D(self.nc//self.squeeze_factor, 1, kernel_regularizer=keras.regularizers.l2(1e-4))(x) + h = Conv2D(self.nc, 1, kernel_regularizer=keras.regularizers.l2(1e-4))(x) + + shape_f = f.get_shape().as_list() + shape_g = g.get_shape().as_list() + shape_h = h.get_shape().as_list() + flat_f = Reshape( (-1, shape_f[-1]) )(f) + flat_g = Reshape( (-1, shape_g[-1]) )(g) + flat_h = Reshape( (-1, shape_h[-1]) )(h) + + s = Lambda(lambda x: K.batch_dot(x[0], keras.layers.Permute((2,1))(x[1]) ))([flat_g, flat_f]) + beta = keras.layers.Softmax(axis=-1)(s) + o = Lambda(lambda x: K.batch_dot(x[0], x[1]))([beta, flat_h]) + + o = Reshape(shape_x[1:])(o) + o = Scale()(o) + + out = Add()([o, inp]) + return out + nnlib.SelfAttention = SelfAttention + class Adam(keras.optimizers.Optimizer): """Adam optimizer. diff --git a/samplelib/SampleGeneratorBase.py b/samplelib/SampleGeneratorBase.py index dec741e..b89c506 100644 --- a/samplelib/SampleGeneratorBase.py +++ b/samplelib/SampleGeneratorBase.py @@ -13,7 +13,18 @@ class SampleGeneratorBase(object): self.samples_path = Path(samples_path) self.debug = debug self.batch_size = 1 if self.debug else batch_size - + self.last_generation = None + self.active = True + + def set_active(self, is_active): + self.active = is_active + + def generate_next(self): + if not self.active and self.last_generation is not None: + return self.last_generation + self.last_generation = next(self) + return self.last_generation + #overridable def __iter__(self): #implement your own iterator diff --git a/samplelib/SampleGeneratorFaceTemporal.py b/samplelib/SampleGeneratorFaceTemporal.py index cf670d2..add15ae 100644 --- a/samplelib/SampleGeneratorFaceTemporal.py +++ b/samplelib/SampleGeneratorFaceTemporal.py @@ -8,7 +8,7 @@ from samplelib import SampleType, SampleProcessor, SampleLoader, SampleGenerator ''' output_sample_types = [ - [SampleProcessor.TypeFlags, size, (optional)random_sub_size] , + [SampleProcessor.TypeFlags, size, (optional) {} opts ] , ... ] ''' @@ -46,9 +46,9 @@ class SampleGeneratorFaceTemporal(SampleGeneratorBase): raise ValueError('No training data provided.') mult_max = 1 - l = samples_len - (self.temporal_image_count-1)*mult_max + 1 + l = samples_len - ( (self.temporal_image_count)*mult_max - (mult_max-1) ) - samples_idxs = [ *range(l) ] [generator_id::self.generators_count] + samples_idxs = [ *range(l+1) ] [generator_id::self.generators_count] if len(samples_idxs) - self.temporal_image_count < 0: raise ValueError('Not enough samples to fit temporal line.') @@ -67,7 +67,7 @@ class SampleGeneratorFaceTemporal(SampleGeneratorBase): idx = shuffle_idxs.pop() temporal_samples = [] - mult = np.random.randint(mult_max) + mult = np.random.randint(mult_max)+1 for i in range( self.temporal_image_count ): sample = samples[ idx+i*mult ] try: diff --git a/samplelib/SampleGeneratorImageTemporal.py b/samplelib/SampleGeneratorImageTemporal.py index 190f98d..227c09e 100644 --- a/samplelib/SampleGeneratorImageTemporal.py +++ b/samplelib/SampleGeneratorImageTemporal.py @@ -43,7 +43,8 @@ class SampleGeneratorImageTemporal(SampleGeneratorBase): raise ValueError('No training data provided.') mult_max = 4 - samples_sub_len = samples_len - (self.temporal_image_count-1)*mult_max + samples_sub_len = samples_len - ( (self.temporal_image_count)*mult_max - (mult_max-1) ) + if samples_sub_len <= 0: raise ValueError('Not enough samples to fit temporal line.') @@ -61,7 +62,7 @@ class SampleGeneratorImageTemporal(SampleGeneratorBase): idx = shuffle_idxs.pop() temporal_samples = [] - mult = np.random.randint(mult_max) + mult = np.random.randint(mult_max)+1 for i in range( self.temporal_image_count ): sample = samples[ idx+i*mult ] try: diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py index 0921e52..338dc65 100644 --- a/samplelib/SampleProcessor.py +++ b/samplelib/SampleProcessor.py @@ -61,6 +61,7 @@ class SampleProcessor(object): FACE_TYPE_FULL = 11 FACE_TYPE_HEAD = 12 #currently unused FACE_TYPE_AVATAR = 13 #currently unused + FACE_TYPE_FULL_NO_ALIGN = 14 FACE_TYPE_END = 20 MODE_BEGIN = 40 @@ -103,7 +104,7 @@ class SampleProcessor(object): SPTF_FACETYPE_TO_FACETYPE = { SPTF.FACE_TYPE_HALF : FaceType.HALF, SPTF.FACE_TYPE_FULL : FaceType.FULL, SPTF.FACE_TYPE_HEAD : FaceType.HEAD, - SPTF.FACE_TYPE_AVATAR : FaceType.AVATAR } + SPTF.FACE_TYPE_FULL_NO_ALIGN : FaceType.FULL_NO_ALIGN } outputs = [] for opts in output_sample_types: @@ -157,6 +158,20 @@ class SampleProcessor(object): if mode_type == SPTF.NONE: raise ValueError ('expected MODE_ type') + def do_transform(img, mask): + warp = (img_type==SPTF.IMG_WARPED or img_type==SPTF.IMG_WARPED_TRANSFORMED) + transform = (img_type==SPTF.IMG_WARPED_TRANSFORMED or img_type==SPTF.IMG_TRANSFORMED) + flip = img_type != SPTF.IMG_WARPED + + img = imagelib.warp_by_params (params, img, warp, transform, flip, True) + if mask is not None: + mask = imagelib.warp_by_params (params, mask, warp, transform, flip, False) + if len(mask.shape) == 2: + mask = mask[...,np.newaxis] + + img = np.concatenate( (img, mask ), -1 ) + return img + img = cached_images.get(img_type, None) if img is None: @@ -181,15 +196,12 @@ class SampleProcessor(object): if cur_sample.ie_polys is not None: cur_sample.ie_polys.overlay_mask(mask) - - warp = (img_type==SPTF.IMG_WARPED or img_type==SPTF.IMG_WARPED_TRANSFORMED) - transform = (img_type==SPTF.IMG_WARPED_TRANSFORMED or img_type==SPTF.IMG_TRANSFORMED) - flip = img_type != SPTF.IMG_WARPED - - img = imagelib.warp_by_params (params, img, warp, transform, flip, True) - if mask is not None: - mask = imagelib.warp_by_params (params, mask, warp, transform, flip, False)[...,np.newaxis] - img = np.concatenate( (img, mask ), -1 ) + + if sample.face_type == FaceType.MARK_ONLY: + if mask is not None: + img = np.concatenate( (img, mask), -1 ) + else: + img = do_transform (img, mask) cached_images[img_type] = img @@ -197,7 +209,17 @@ class SampleProcessor(object): ft = SPTF_FACETYPE_TO_FACETYPE[target_face_type] if ft > sample.face_type: raise Exception ('sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.' % (sample.filename, sample.face_type, ft) ) - img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, resolution, ft), (resolution,resolution), flags=cv2.INTER_CUBIC ) + + if sample.face_type == FaceType.MARK_ONLY: + img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, sample.shape[0], ft), (sample.shape[0],sample.shape[0]), flags=cv2.INTER_CUBIC ) + + mask = img[...,3:4] if img.shape[2] > 3 else None + img = img[...,0:3] + img = do_transform (img, mask) + img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC ) + else: + img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, resolution, ft), (resolution,resolution), flags=cv2.INTER_CUBIC ) + else: img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC ) From d16759bcd3d3b8f3eedc83cf81f0dd4257b688c7 Mon Sep 17 00:00:00 2001 From: iperov Date: Sun, 11 Aug 2019 13:22:01 +0400 Subject: [PATCH 2/6] ConverterMasked: added mask gradient of bottom area, same as side gradient --- converters/ConverterMasked.py | 1 + 1 file changed, 1 insertion(+) diff --git a/converters/ConverterMasked.py b/converters/ConverterMasked.py index 34e494c..d080e49 100644 --- a/converters/ConverterMasked.py +++ b/converters/ConverterMasked.py @@ -260,6 +260,7 @@ class ConverterMasked(Converter): prd_border_size = int ( prd_hborder_rect_mask_a.shape[1] * self.clip_hborder_mask_per ) prd_hborder_rect_mask_a[:,0:prd_border_size,:] = 0 prd_hborder_rect_mask_a[:,-prd_border_size:,:] = 0 + prd_hborder_rect_mask_a[-prd_border_size:,:,:] = 0 prd_hborder_rect_mask_a = np.expand_dims(cv2.blur(prd_hborder_rect_mask_a, (prd_border_size, prd_border_size) ),-1) img_prd_hborder_rect_mask_a = cv2.warpAffine( prd_hborder_rect_mask_a, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4 ) From 625bcc212da9d1bae8396e8cac50d488af29302f Mon Sep 17 00:00:00 2001 From: Auroir <36361058+Auroir@users.noreply.github.com> Date: Mon, 12 Aug 2019 07:18:55 -0700 Subject: [PATCH 3/6] Return sorted filenames in path_utils (#340) Linux does not guarantee filenames are returned in any specific order. This leads to exporting frames in random order, sorting them here makes the export run sequentially. Other portions of the program should remain unaffected, if not behave more consistently (E.G. get_first_file_by_stem). This mostly helpful during exporting. Say you are expecting to not have faces for frames 1000-2000, during your export all the "no faces for..." messages will appear in random order. Since you are expecting to see this you ignore them. If you are also (unexpectedly) missing a face for frame 3000 you will not head the warning since it's mixed up in all the warnings that you are expecting. With this patch export runs in sequential order, you'll see the messages all in a row for frames 1000-2000, then again at 3000. The user is much more likely to see and head the warning this way. This also allows you to force stop the export midway though and have a contiguous set of frames to encode and preview. --- utils/Path_utils.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/utils/Path_utils.py b/utils/Path_utils.py index c9ed448..d34e348 100644 --- a/utils/Path_utils.py +++ b/utils/Path_utils.py @@ -11,7 +11,7 @@ def get_image_paths(dir_path, image_extensions=image_extensions): for x in list(scandir(str(dir_path))): if any([x.name.lower().endswith(ext) for ext in image_extensions]): result.append(x.path) - return result + return sorted(result) def get_image_unique_filestem_paths(dir_path, verbose_print_func=None): result = get_image_paths(dir_path) @@ -26,7 +26,7 @@ def get_image_unique_filestem_paths(dir_path, verbose_print_func=None): continue result_dup.add(f_stem) - return result + return sorted(result) def get_file_paths(dir_path): dir_path = Path (dir_path) @@ -34,7 +34,7 @@ def get_file_paths(dir_path): result = [] if dir_path.exists(): return [ x.path for x in list(scandir(str(dir_path))) if x.is_file() ] - return result + return sorted(result) def get_all_dir_names (dir_path): dir_path = Path (dir_path) @@ -43,7 +43,7 @@ def get_all_dir_names (dir_path): if dir_path.exists(): return [ x.name for x in list(scandir(str(dir_path))) if x.is_dir() ] - return result + return sorted(result) def get_all_dir_names_startswith (dir_path, startswith): dir_path = Path (dir_path) @@ -54,14 +54,14 @@ def get_all_dir_names_startswith (dir_path, startswith): for x in list(scandir(str(dir_path))): if x.name.lower().startswith(startswith): result.append ( x.name[len(startswith):] ) - return result + return sorted(result) def get_first_file_by_stem (dir_path, stem, exts=None): dir_path = Path (dir_path) stem = stem.lower() if dir_path.exists(): - for x in list(scandir(str(dir_path))): + for x in sorted(list(scandir(str(dir_path)))): if not x.is_file(): continue xp = Path(x.path) @@ -80,4 +80,4 @@ def delete_all_files (dir_path): paths = get_file_paths(dir_path) for p in paths: p = Path(p) - p.unlink() \ No newline at end of file + p.unlink() From c4e68ef539aaf77f3864d2fe5401ac6ddf3590db Mon Sep 17 00:00:00 2001 From: Auroir <36361058+Auroir@users.noreply.github.com> Date: Fri, 16 Aug 2019 07:35:27 -0700 Subject: [PATCH 4/6] Formatted Model Summary (#348) * Formatted Model Summary Aligns the model summary output using f-string formatting. The logic structure of the base class has not been changed, only the lines put into `model_summary_text`. Output width is calculated from keys & values and will scale to show a clean summary for any model/platform. GPU VRAM has been added as an output. Incorrect detection of VRAM is possible in broken environments and GPUs of different sizes can report the same name. Showing it here adds clarity for the user and for issue tickets. Concatenation changed from "\r\n" to "\n", CRLF end of lines for Windows are handled transparently so using it here caused extra blank lines in the summary txt file. **Examples:** Using CUDA + SAE-LIAE ``` ============= Model Summary ============== == == == Model name: SAE == == == == Current iteration: 16 == == == ==----------- Model Options ------------== == == == batch_size: 4 == == sort_by_yaw: False == == random_flip: True == == resolution: 128 == == face_type: f == == learn_mask: True == == optimizer_mode: 1 == == archi: liae == == ae_dims: 256 == == e_ch_dims: 42 == == d_ch_dims: 21 == == multiscale_decoder: False == == ca_weights: False == == pixel_loss: False == == face_style_power: 0.0 == == bg_style_power: 0.0 == == apply_random_ct: False == == clipgrad: False == == == ==------------- Running On -------------== == == == Device index: 0 == == Name: GeForce GTX 1080 == == VRAM: 8.00GB == == == ========================================== ``` Colab ``` ========== Model Summary ========== == == == Model name: SAE == == == == Current iteration: 39822 == == == ==-------- Model Options --------== == == == batch_size: 24 == == sort_by_yaw: True == == random_flip: False == == resolution: 128 == == face_type: f == == learn_mask: True == == optimizer_mode: 2 == == archi: liae == == ae_dims: 222 == == e_ch_dims: 34 == == d_ch_dims: 16 == == multiscale_decoder: True == == ca_weights: True == == pixel_loss: False == == face_style_power: 2.0 == == bg_style_power: 1.5 == == apply_random_ct: False == == clipgrad: True == == == ==--------- Running On ----------== == == == Device index: 0 == == Name: Tesla K80 == == VRAM: 11.00GB == == == =================================== ``` Using OpenCL + H128 ``` =========================== Model Summary =========================== == == == Model name: H128 == == == == Current iteration: 0 == == == ==------------------------- Model Options -------------------------== == == == batch_size: 4 == == sort_by_yaw: False == == random_flip: True == == lighter_ae: False == == pixel_loss: False == == == ==-------------------------- Running On ---------------------------== == == == Device index: 0 == == Name: Advanced Micro Devices, Inc. gfx900 (OpenCL) == == VRAM: 7.98GB == == == ===================================================================== ``` Using CPU (output trimmed) ``` ==------- Running On --------== == == == Using device: CPU == == == =============================== ``` multi_gpu support is retained (output trimmed) ``` ==------------- Running On -------------== == == == Using multi_gpu: True == == == == Device index: 1 == == Name: Geforce GTX 1080 == == VRAM: 8.00GB == == Device index: 2 == == Name: Geforce GTX 1080 == == VRAM: 8.00GB == == == ========================================== ``` Low VRAM warning (output trimmed) ``` ==------------- Running On -------------== == == == Device index: 0 == == Name: Geforce GTX 1050 == == VRAM: 2.00GB == == == ========================================== /!\ /!\ WARNING: /!\ You are using a GPU with 2GB or less VRAM. This may significantly reduce the quality of your result! /!\ If training does not start, close all programs and try again. /!\ Also you can disable Windows Aero Desktop to increase available VRAM. /!\ ``` * Fix indent --- models/ModelBase.py | 66 ++++++++++++++++++++++++++++----------------- 1 file changed, 42 insertions(+), 24 deletions(-) diff --git a/models/ModelBase.py b/models/ModelBase.py index f841ca6..e7ad843 100644 --- a/models/ModelBase.py +++ b/models/ModelBase.py @@ -231,36 +231,54 @@ class ModelBase(object): else: self.sample_for_preview = self.generate_next_sample() self.last_sample = self.sample_for_preview + + ###Generate text summary of model hyperparameters + #Find the longest key name and value string. Used as column widths. + width_name = max([len(k) for k in self.options.keys()] + [17]) + 1 # Single space buffer to left edge. Minimum of 17, the length of the longest static string used "Current iteration" + width_value = max([len(str(x)) for x in self.options.values()] + [len(str(self.iter)), len(self.get_model_name())]) + 1 # Single space buffer to right edge + if not self.device_config.cpu_only: #Check length of GPU names + width_value = max([len(nnlib.device.getDeviceName(idx))+1 for idx in self.device_config.gpu_idxs] + [width_value]) + width_total = width_name + width_value + 2 #Plus 2 for ": " + model_summary_text = [] - - model_summary_text += ["===== Model summary ====="] - model_summary_text += ["== Model name: " + self.get_model_name()] - model_summary_text += ["=="] - model_summary_text += ["== Current iteration: " + str(self.iter)] - model_summary_text += ["=="] - model_summary_text += ["== Model options:"] + model_summary_text += [f'=={" Model Summary ":=^{width_total}}=='] # Model/status summary + model_summary_text += [f'=={" "*width_total}=='] + model_summary_text += [f'=={"Model name": >{width_name}}: {self.get_model_name(): <{width_value}}=='] # Name + model_summary_text += [f'=={" "*width_total}=='] + model_summary_text += [f'=={"Current iteration": >{width_name}}: {str(self.iter): <{width_value}}=='] # Iter + model_summary_text += [f'=={" "*width_total}=='] + + model_summary_text += [f'=={" Model Options ":-^{width_total}}=='] # Model options + model_summary_text += [f'=={" "*width_total}=='] for key in self.options.keys(): - model_summary_text += ["== |== %s : %s" % (key, self.options[key])] - + model_summary_text += [f'=={key: >{width_name}}: {str(self.options[key]): <{width_value}}=='] # self.options key/value pairs + model_summary_text += [f'=={" "*width_total}=='] + + model_summary_text += [f'=={" Running On ":-^{width_total}}=='] # Training hardware info + model_summary_text += [f'=={" "*width_total}=='] if self.device_config.multi_gpu: - model_summary_text += ["== |== multi_gpu : True "] - - model_summary_text += ["== Running on:"] + model_summary_text += [f'=={"Using multi_gpu": >{width_name}}: {"True": <{width_value}}=='] # multi_gpu + model_summary_text += [f'=={" "*width_total}=='] if self.device_config.cpu_only: - model_summary_text += ["== |== [CPU]"] + model_summary_text += [f'=={"Using device": >{width_name}}: {"CPU": <{width_value}}=='] # cpu_only else: for idx in self.device_config.gpu_idxs: - model_summary_text += ["== |== [%d : %s]" % (idx, nnlib.device.getDeviceName(idx))] - - if not self.device_config.cpu_only and self.device_config.gpu_vram_gb[0] == 2: - model_summary_text += ["=="] - model_summary_text += ["== WARNING: You are using 2GB GPU. Result quality may be significantly decreased."] - model_summary_text += ["== If training does not start, close all programs and try again."] - model_summary_text += ["== Also you can disable Windows Aero Desktop to get extra free VRAM."] - model_summary_text += ["=="] - - model_summary_text += ["========================="] - model_summary_text = "\r\n".join (model_summary_text) + model_summary_text += [f'=={"Device index": >{width_name}}: {idx: <{width_value}}=='] # GPU hardware device index + model_summary_text += [f'=={"Name": >{width_name}}: {nnlib.device.getDeviceName(idx): <{width_value}}=='] # GPU name + vram_str = f'{nnlib.device.getDeviceVRAMTotalGb(idx):.2f}GB' # GPU VRAM - Formated as #.## (or ##.##) + model_summary_text += [f'=={"VRAM": >{width_name}}: {vram_str: <{width_value}}=='] + model_summary_text += [f'=={" "*width_total}=='] + model_summary_text += [f'=={"="*width_total}=='] + + if not self.device_config.cpu_only and self.device_config.gpu_vram_gb[0] <= 2: # Low VRAM warning + model_summary_text += ["/!\\"] + model_summary_text += ["/!\\ WARNING:"] + model_summary_text += ["/!\\ You are using a GPU with 2GB or less VRAM. This may significantly reduce the quality of your result!"] + model_summary_text += ["/!\\ If training does not start, close all programs and try again."] + model_summary_text += ["/!\\ Also you can disable Windows Aero Desktop to increase available VRAM."] + model_summary_text += ["/!\\"] + + model_summary_text = "\n".join (model_summary_text) self.model_summary_text = model_summary_text io.log_info(model_summary_text) From a906f24a4d89dd0895a13f04fb96c051d176c0f3 Mon Sep 17 00:00:00 2001 From: Auroir <36361058+Auroir@users.noreply.github.com> Date: Sat, 17 Aug 2019 00:20:10 -0700 Subject: [PATCH 5/6] Fix log_info/log_err printing with progress bars (#349) --- interact/interact.py | 22 +++++++++++++++++++--- 1 file changed, 19 insertions(+), 3 deletions(-) diff --git a/interact/interact.py b/interact/interact.py index c9b22f1..bda74af 100644 --- a/interact/interact.py +++ b/interact/interact.py @@ -33,6 +33,7 @@ class InteractBase(object): self.key_events = {} self.pg_bar = None self.focus_wnd_name = None + self.error_log_line_prefix = '/!\\ ' def is_support_windows(self): return False @@ -65,10 +66,22 @@ class InteractBase(object): raise NotImplemented def log_info(self, msg, end='\n'): + if self.pg_bar is not None: + try: # Attempt print before the pb + tqdm.write(msg) + return + except: + pass #Fallback to normal print upon failure print (msg, end=end) def log_err(self, msg, end='\n'): - print (msg, end=end) + if self.pg_bar is not None: + try: # Attempt print before the pb + tqdm.write(f'{self.error_log_line_prefix}{msg}') + return + except: + pass #Fallback to normal print upon failure + print (f'{self.error_log_line_prefix}{msg}', end=end) def named_window(self, wnd_name): if wnd_name not in self.named_windows: @@ -150,9 +163,12 @@ class InteractBase(object): else: print("progress_bar not set.") def progress_bar_generator(self, data, desc, leave=True): - for x in tqdm( data, desc=desc, leave=leave, ascii=True ): + self.pg_bar = tqdm( data, desc=desc, leave=leave, ascii=True ) + for x in self.pg_bar: yield x - + self.pg_bar.close() + self.pg_bar = None + def process_messages(self, sleep_time=0): self.on_process_messages(sleep_time) From e7562054d06e2482f274359d7dd9904d3c871c55 Mon Sep 17 00:00:00 2001 From: Auroir <36361058+Auroir@users.noreply.github.com> Date: Sat, 17 Aug 2019 21:01:44 -0700 Subject: [PATCH 6/6] Fix sorting pathlib objects (#353) Fixing an issue caused by attempting to sort Path objects. Directly using `<` is unsupported between these, so `sorted()` needs a key specified. "PurePath" objects support `>` while normal paths do not, causing the confusion. https://docs.python.org/3/library/pathlib.html --- utils/Path_utils.py | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) diff --git a/utils/Path_utils.py b/utils/Path_utils.py index d34e348..e753834 100644 --- a/utils/Path_utils.py +++ b/utils/Path_utils.py @@ -31,19 +31,18 @@ def get_image_unique_filestem_paths(dir_path, verbose_print_func=None): def get_file_paths(dir_path): dir_path = Path (dir_path) - result = [] if dir_path.exists(): - return [ x.path for x in list(scandir(str(dir_path))) if x.is_file() ] - return sorted(result) + return sorted([ x.path for x in list(scandir(str(dir_path))) if x.is_file() ]) + else: + return [] def get_all_dir_names (dir_path): dir_path = Path (dir_path) - result = [] if dir_path.exists(): - return [ x.name for x in list(scandir(str(dir_path))) if x.is_dir() ] - - return sorted(result) + return sorted([ x.name for x in list(scandir(str(dir_path))) if x.is_dir() ]) + else: + return [] def get_all_dir_names_startswith (dir_path, startswith): dir_path = Path (dir_path) @@ -61,7 +60,7 @@ def get_first_file_by_stem (dir_path, stem, exts=None): stem = stem.lower() if dir_path.exists(): - for x in sorted(list(scandir(str(dir_path)))): + for x in sorted(list(scandir(str(dir_path))), key=lambda x: x.name): if not x.is_file(): continue xp = Path(x.path)