diff --git a/facelib/FANSegmentator.py b/facelib/FANSegmentator.py index 531b829..4ec47c4 100644 --- a/facelib/FANSegmentator.py +++ b/facelib/FANSegmentator.py @@ -12,6 +12,9 @@ from nnlib import nnlib """ FANSegmentator is designed to segment faces aligned by 2DFAN-4 landmarks extractor. +Dataset used to train located in official DFL mega.nz folder +https://mega.nz/#F!b9MzCK4B!zEAG9txu7uaRUjXz9PtBqg + using https://github.com/ternaus/TernausNet TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentation """ @@ -33,18 +36,18 @@ class FANSegmentator(object): if load_weights: self.model.load_weights (str(self.weights_path)) else: - if training: + if training: try: with open( Path(__file__).parent / 'vgg11_enc_weights.npy', 'rb' ) as f: d = pickle.loads (f.read()) for i in [0,3,6,8,11,13,16,18]: s = 'features.%d' % i - + self.model.get_layer (s).set_weights ( d[s] ) except: io.log_err("Unable to load VGG11 pretrained weights from vgg11_enc_weights.npy") - + if training: self.model.compile(loss='mse', optimizer=Adam(tf_cpu_mode=2)) @@ -64,14 +67,14 @@ class FANSegmentator(object): input_shape_len = len(input_image.shape) if input_shape_len == 3: input_image = input_image[np.newaxis,...] - + result = np.clip ( self.model.predict( [input_image] ), 0, 1.0 ) - + if input_shape_len == 3: result = result[0] - + return result - + @staticmethod def BuildModel ( resolution, ngf=64, norm='', act='lrelu'): exec( nnlib.import_all(), locals(), globals() ) @@ -90,45 +93,45 @@ class FANSegmentator(object): x0 = x = Conv2D(ngf, kernel_size=3, strides=1, padding='same', activation='relu', name='features.0')(x) x = MaxPooling2D()(x) - + x1 = x = Conv2D(ngf*2, kernel_size=3, strides=1, padding='same', activation='relu', name='features.3')(x) x = MaxPooling2D()(x) - + x = Conv2D(ngf*4, kernel_size=3, strides=1, padding='same', activation='relu', name='features.6')(x) x2 = x = Conv2D(ngf*4, kernel_size=3, strides=1, padding='same', activation='relu', name='features.8')(x) x = MaxPooling2D()(x) - + x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.11')(x) x3 = x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.13')(x) x = MaxPooling2D()(x) - + x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.16')(x) x4 = x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.18')(x) x = MaxPooling2D()(x) - + x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same')(x) - - x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu') (x) - x = Concatenate(axis=3)([ x, x4]) - x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu') (x) - - x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu') (x) - x = Concatenate(axis=3)([ x, x3]) + + x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu') (x) + x = Concatenate(axis=3)([ x, x4]) x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu') (x) - - x = Conv2DTranspose (ngf*2, 3, strides=2, padding='same', activation='relu') (x) - x = Concatenate(axis=3)([ x, x2]) - x = Conv2D (ngf*4, 3, strides=1, padding='same', activation='relu') (x) - - x = Conv2DTranspose (ngf, 3, strides=2, padding='same', activation='relu') (x) - x = Concatenate(axis=3)([ x, x1]) - x = Conv2D (ngf*2, 3, strides=1, padding='same', activation='relu') (x) - - x = Conv2DTranspose (ngf // 2, 3, strides=2, padding='same', activation='relu') (x) - x = Concatenate(axis=3)([ x, x0]) + + x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu') (x) + x = Concatenate(axis=3)([ x, x3]) + x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu') (x) + + x = Conv2DTranspose (ngf*2, 3, strides=2, padding='same', activation='relu') (x) + x = Concatenate(axis=3)([ x, x2]) + x = Conv2D (ngf*4, 3, strides=1, padding='same', activation='relu') (x) + + x = Conv2DTranspose (ngf, 3, strides=2, padding='same', activation='relu') (x) + x = Concatenate(axis=3)([ x, x1]) + x = Conv2D (ngf*2, 3, strides=1, padding='same', activation='relu') (x) + + x = Conv2DTranspose (ngf // 2, 3, strides=2, padding='same', activation='relu') (x) + x = Concatenate(axis=3)([ x, x0]) x = Conv2D (ngf, 3, strides=1, padding='same', activation='relu') (x) - + return Conv2D(1, 3, strides=1, padding='same', activation='sigmoid')(x) - - + + return func diff --git a/main.py b/main.py index 5f210d6..d2f56b7 100644 --- a/main.py +++ b/main.py @@ -49,6 +49,21 @@ if __name__ == "__main__": p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU. Forces to use MT extractor.") p.set_defaults (func=process_extract) + def process_extract_fanseg(arguments): + os_utils.set_process_lowest_prio() + from mainscripts import Extractor + Extractor.extract_fanseg( arguments.input_dir, + device_args={'cpu_only' : arguments.cpu_only, + 'multi_gpu' : arguments.multi_gpu, + } + ) + + p = subparsers.add_parser( "extract_fanseg", help="Extract fanseg mask from faces.") + 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('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="Enables multi GPU.") + p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU.") + p.set_defaults (func=process_extract_fanseg) + def process_sort(arguments): os_utils.set_process_lowest_prio() from mainscripts import Sorter @@ -71,12 +86,17 @@ if __name__ == "__main__": if arguments.recover_original_aligned_filename: Util.recover_original_aligned_filename (input_path=arguments.input_dir) - + + if arguments.remove_fanseg: + Util.remove_fanseg_folder (input_path=arguments.input_dir) + p = subparsers.add_parser( "util", help="Utilities.") 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('--convert-png-to-jpg', action="store_true", dest="convert_png_to_jpg", default=False, help="Convert DeepFaceLAB PNG files to JPEG.") p.add_argument('--add-landmarks-debug-images', action="store_true", dest="add_landmarks_debug_images", default=False, help="Add landmarks debug image for aligned faces.") p.add_argument('--recover-original-aligned-filename', action="store_true", dest="recover_original_aligned_filename", default=False, help="Recover original aligned filename.") + p.add_argument('--remove-fanseg', action="store_true", dest="remove_fanseg", default=False, help="Remove fanseg mask from aligned faces.") + p.set_defaults (func=process_util) def process_train(arguments): diff --git a/mainscripts/Extractor.py b/mainscripts/Extractor.py index eaee690..1797a88 100644 --- a/mainscripts/Extractor.py +++ b/mainscripts/Extractor.py @@ -10,11 +10,13 @@ import mathlib import imagelib import cv2 from utils import Path_utils +from utils.DFLPNG import DFLPNG from utils.DFLJPG import DFLJPG from utils.cv2_utils import * import facelib from facelib import FaceType from facelib import LandmarksProcessor +from facelib import FANSegmentator from nnlib import nnlib from joblib import Subprocessor from interact import interact as io @@ -79,7 +81,12 @@ class ExtractSubprocessor(Subprocessor): self.second_pass_e.__enter__() else: self.second_pass_e = None - + + elif self.type == 'fanseg': + nnlib.import_all (device_config) + self.e = facelib.FANSegmentator(256, FaceType.toString(FaceType.FULL) ) + self.e.__enter__() + elif self.type == 'final': pass @@ -124,6 +131,8 @@ class ExtractSubprocessor(Subprocessor): h, w, ch = image.shape if h == w: #extracting from already extracted jpg image? + if filename_path.suffix == '.png': + src_dflimg = DFLPNG.load ( str(filename_path) ) if filename_path.suffix == '.jpg': src_dflimg = DFLJPG.load ( str(filename_path) ) @@ -253,15 +262,22 @@ class ExtractSubprocessor(Subprocessor): cv2_imwrite(debug_output_file, debug_image, [int(cv2.IMWRITE_JPEG_QUALITY), 50] ) return data - - + + elif self.type == 'fanseg': + if src_dflimg is not None: + fanseg_mask = self.e.extract( image / 255.0 ) + src_dflimg.embed_and_set( filename_path_str, + fanseg_mask=fanseg_mask, + #fanseg_mask_ver=FANSegmentator.VERSION, + ) + #overridable def get_data_name (self, data): #return string identificator of your data return data.filename #override - def __init__(self, input_data, type, image_size, face_type, debug_dir=None, multi_gpu=False, cpu_only=False, manual=False, manual_window_size=0, final_output_path=None): + def __init__(self, input_data, type, image_size=None, face_type=None, debug_dir=None, multi_gpu=False, cpu_only=False, manual=False, manual_window_size=0, final_output_path=None): self.input_data = input_data self.type = type self.image_size = image_size @@ -561,7 +577,7 @@ class ExtractSubprocessor(Subprocessor): if 'cpu' in backend: cpu_only = True - if 'rects' in type or type == 'landmarks': + if 'rects' in type or type == 'landmarks' or type == 'fanseg': if not cpu_only and type == 'rects-mt' and backend == "plaidML": #plaidML works with MT very slowly cpu_only = True @@ -583,7 +599,7 @@ class ExtractSubprocessor(Subprocessor): dev_name = nnlib.device.getDeviceName(idx) dev_vram = nnlib.device.getDeviceVRAMTotalGb(idx) - if not manual and (type == 'rects-dlib' or type == 'rects-mt'): + if not manual and (type == 'rects-dlib' or type == 'rects-mt' ): for i in range ( int (max (1, dev_vram / 2) ) ): result += [ (idx, 'GPU', '%s #%d' % (dev_name,i) , dev_vram) ] else: @@ -657,8 +673,33 @@ class DeletedFilesSearcherSubprocessor(Subprocessor): def get_result(self): return self.result +def extract_fanseg(input_dir, device_args={} ):#ignore_extracted + multi_gpu = device_args.get('multi_gpu', False) + cpu_only = device_args.get('cpu_only', False) + + input_path = Path(input_dir) + if not input_path.exists(): + raise ValueError('Input directory not found. Please ensure it exists.') + + paths_to_extract = [] + for filename in Path_utils.get_image_paths(input_path) : + filepath = Path(filename) + if filepath.suffix == '.png': + dflimg = DFLPNG.load( str(filepath) ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filepath) ) + else: + dflimg = None + if dflimg is not None: + paths_to_extract.append (filepath) + + paths_to_extract_len = len(paths_to_extract) + if paths_to_extract_len > 0: + io.log_info ("Performing extract fanseg for %d files..." % (paths_to_extract_len) ) + data = ExtractSubprocessor ([ ExtractSubprocessor.Data(filename) for filename in paths_to_extract ], 'fanseg', multi_gpu=multi_gpu, cpu_only=cpu_only).run() + def main(input_dir, output_dir, debug_dir=None, diff --git a/mainscripts/MaskEditorTool.py b/mainscripts/MaskEditorTool.py index 9c48a7a..1d9750f 100644 --- a/mainscripts/MaskEditorTool.py +++ b/mainscripts/MaskEditorTool.py @@ -397,13 +397,17 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None): else: lmrks = dflimg.get_landmarks() ie_polys = dflimg.get_ie_polys() + fanseg_mask = dflimg.get_fanseg_mask() if filepath.name in cached_images: img = cached_images[filepath.name] else: img = cached_images[filepath.name] = cv2_imread(str(filepath)) / 255.0 - mask = LandmarksProcessor.get_image_hull_mask( img.shape, lmrks) + if fanseg_mask is not None: + mask = fanseg_mask + else: + mask = LandmarksProcessor.get_image_hull_mask( img.shape, lmrks) else: img = np.zeros ( (target_wh,target_wh,3) ) mask = np.ones ( (target_wh,target_wh,3) ) diff --git a/mainscripts/Util.py b/mainscripts/Util.py index a7c4a28..0d3e7c4 100644 --- a/mainscripts/Util.py +++ b/mainscripts/Util.py @@ -7,6 +7,33 @@ from utils.cv2_utils import * from facelib import LandmarksProcessor from interact import interact as io +def remove_fanseg_file (filepath): + filepath = Path(filepath) + + if filepath.suffix == '.png': + dflimg = DFLPNG.load( str(filepath) ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filepath) ) + else: + return + + if dflimg is None: + io.log_err ("%s is not a dfl image file" % (filepath.name) ) + return + + dflimg.remove_fanseg_mask() + dflimg.embed_and_set( str(filepath) ) + + +def remove_fanseg_folder(input_path): + input_path = Path(input_path) + + io.log_info ("Removing fanseg mask...\r\n") + + for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Removing"): + filepath = Path(filepath) + remove_fanseg_file(filepath) + def convert_png_to_jpg_file (filepath): filepath = Path(filepath) diff --git a/models/Model_SAE/Model.py b/models/Model_SAE/Model.py index 3ddbe7e..c46ef02 100644 --- a/models/Model_SAE/Model.py +++ b/models/Model_SAE/Model.py @@ -385,20 +385,20 @@ class SAEModel(ModelBase): #override def onGetPreview(self, sample): - test_A = sample[0][1][0:4] #first 4 samples - test_A_m = sample[0][2][0:4] #first 4 samples - test_B = sample[1][1][0:4] - test_B_m = sample[1][2][0:4] + test_S = sample[0][1][0:4] #first 4 samples + test_S_m = sample[0][2][0:4] #first 4 samples + test_D = sample[1][1][0:4] + test_D_m = sample[1][2][0:4] if self.options['learn_mask']: - S, D, SS, DD, DDM, SD, SDM = [ np.clip(x, 0.0, 1.0) for x in ([test_A,test_B] + self.AE_view ([test_A, test_B]) ) ] + S, D, SS, DD, DDM, SD, SDM = [ np.clip(x, 0.0, 1.0) for x in ([test_S,test_D] + self.AE_view ([test_S, test_D]) ) ] DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ] else: - S, D, SS, DD, SD, = [ np.clip(x, 0.0, 1.0) for x in ([test_A,test_B] + self.AE_view ([test_A, test_B]) ) ] + S, D, SS, DD, SD, = [ np.clip(x, 0.0, 1.0) for x in ([test_S,test_D] + self.AE_view ([test_S, test_D]) ) ] result = [] st = [] - for i in range(0, len(test_A)): + for i in range(0, len(test_S)): ar = S[i], SS[i], D[i], DD[i], SD[i] st.append ( np.concatenate ( ar, axis=1) ) @@ -406,8 +406,15 @@ class SAEModel(ModelBase): if self.options['learn_mask']: st_m = [] - for i in range(0, len(test_A)): - ar = S[i], SS[i], D[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i]) + for i in range(0, len(test_S)): + ar = S[i]*test_S_m[i], SS[i], D[i]*test_D_m[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i]) + st_m.append ( np.concatenate ( ar, axis=1) ) + + result += [ ('SAE masked', np.concatenate (st_m, axis=0 )), ] + else: + st_m = [] + for i in range(0, len(test_S)): + ar = S[i]*test_S_m[i], SS[i], D[i]*test_D_m[i], DD[i], SD[i] st_m.append ( np.concatenate ( ar, axis=1) ) result += [ ('SAE masked', np.concatenate (st_m, axis=0 )), ] diff --git a/samplelib/Sample.py b/samplelib/Sample.py index 1c311e2..3c59896 100644 --- a/samplelib/Sample.py +++ b/samplelib/Sample.py @@ -1,7 +1,13 @@ from enum import IntEnum +from pathlib import Path + import cv2 import numpy as np + from utils.cv2_utils import * +from utils.DFLJPG import DFLJPG +from utils.DFLPNG import DFLPNG + class SampleType(IntEnum): IMAGE = 0 #raw image @@ -16,7 +22,7 @@ class SampleType(IntEnum): QTY = 5 class Sample(object): - def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, mirror=None, close_target_list=None): + def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, mirror=None, close_target_list=None, fanseg_mask_exist=False): self.sample_type = sample_type if sample_type is not None else SampleType.IMAGE self.filename = filename self.face_type = face_type @@ -27,8 +33,9 @@ class Sample(object): self.yaw = yaw self.mirror = mirror self.close_target_list = close_target_list + self.fanseg_mask_exist = fanseg_mask_exist - def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, mirror=None, close_target_list=None): + def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, mirror=None, close_target_list=None, fanseg_mask=None, fanseg_mask_exist=None): return Sample( sample_type=sample_type if sample_type is not None else self.sample_type, filename=filename if filename is not None else self.filename, @@ -39,7 +46,8 @@ class Sample(object): pitch=pitch if pitch is not None else self.pitch, yaw=yaw if yaw is not None else self.yaw, mirror=mirror if mirror is not None else self.mirror, - close_target_list=close_target_list if close_target_list is not None else self.close_target_list) + close_target_list=close_target_list if close_target_list is not None else self.close_target_list, + fanseg_mask_exist=fanseg_mask_exist if fanseg_mask_exist is not None else self.fanseg_mask_exist) def load_bgr(self): img = cv2_imread (self.filename).astype(np.float32) / 255.0 @@ -47,6 +55,19 @@ class Sample(object): img = img[:,::-1].copy() return img + def load_fanseg_mask(self): + if self.fanseg_mask_exist: + filepath = Path(self.filename) + if filepath.suffix == '.png': + dflimg = DFLPNG.load ( str(filepath) ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filepath) ) + else: + dflimg = None + return dflimg.get_fanseg_mask() + + return None + def get_random_close_target_sample(self): if self.close_target_list is None: return None diff --git a/samplelib/SampleLoader.py b/samplelib/SampleLoader.py index 8184409..667acbc 100644 --- a/samplelib/SampleLoader.py +++ b/samplelib/SampleLoader.py @@ -77,7 +77,8 @@ class SampleLoader: landmarks=dflimg.get_landmarks(), ie_polys=dflimg.get_ie_polys(), pitch=pitch, - yaw=yaw) ) + yaw=yaw, + fanseg_mask_exist=dflimg.get_fanseg_mask() is not None, ) ) except: print ("Unable to load %s , error: %s" % (str(s_filename_path), traceback.format_exc() ) ) diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py index 10aefd7..93f30be 100644 --- a/samplelib/SampleProcessor.py +++ b/samplelib/SampleProcessor.py @@ -13,18 +13,18 @@ class SampleProcessor(object): WARPED_TRANSFORMED = 0x00000004, TRANSFORMED = 0x00000008, LANDMARKS_ARRAY = 0x00000010, #currently unused - + RANDOM_CLOSE = 0x00000020, #currently unused MORPH_TO_RANDOM_CLOSE = 0x00000040, #currently unused - + FACE_TYPE_HALF = 0x00000100, FACE_TYPE_FULL = 0x00000200, FACE_TYPE_HEAD = 0x00000400, #currently unused FACE_TYPE_AVATAR = 0x00000800, #currently unused - + FACE_MASK_FULL = 0x00001000, FACE_MASK_EYES = 0x00002000, #currently unused - + MODE_BGR = 0x00010000, #BGR MODE_G = 0x00020000, #Grayscale MODE_GGG = 0x00040000, #3xGrayscale @@ -35,7 +35,7 @@ class SampleProcessor(object): class Options(object): #motion_blur = [chance_int, range] - chance 0..100 to apply to face (not mask), and range [1..3] where 3 is highest power of motion blur - + def __init__(self, random_flip = True, normalize_tanh = False, rotation_range=[-10,10], scale_range=[-0.05, 0.05], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05], motion_blur=None ): self.random_flip = random_flip self.normalize_tanh = normalize_tanh @@ -49,11 +49,11 @@ class SampleProcessor(object): chance = np.clip(chance, 0, 100) range = [3,5,7,9][ : np.clip(range, 0, 3)+1 ] self.motion_blur = (chance, range) - + @staticmethod def process (sample, sample_process_options, output_sample_types, debug): SPTF = SampleProcessor.TypeFlags - + sample_bgr = sample.load_bgr() h,w,c = sample_bgr.shape @@ -113,7 +113,7 @@ class SampleProcessor(object): target_face_type = FaceType.HEAD elif f & SPTF.FACE_TYPE_AVATAR != 0: target_face_type = FaceType.AVATAR - + apply_motion_blur = f & SPTF.OPT_APPLY_MOTION_BLUR != 0 if img_type == 4: @@ -170,9 +170,16 @@ class SampleProcessor(object): if np.random.randint(100) < chance : dim = mb_range[ np.random.randint(len(mb_range) ) ] img = imagelib.LinearMotionBlur (img, dim, np.random.randint(180) ) - + if face_mask_type == 1: - mask = LandmarksProcessor.get_image_hull_mask (img.shape, cur_sample.landmarks, cur_sample.ie_polys) + mask = cur_sample.load_fanseg_mask() #using fanseg_mask if exist + + if mask is None: + mask = LandmarksProcessor.get_image_hull_mask (img.shape, cur_sample.landmarks) + + if cur_sample.ie_polys is not None: + cur_sample.ie_polys.overlay_mask(mask) + img = np.concatenate( (img, mask ), -1 ) elif face_mask_type == 2: mask = LandmarksProcessor.get_image_eye_mask (img.shape, cur_sample.landmarks) diff --git a/utils/DFLJPG.py b/utils/DFLJPG.py index 7aa31b8..a4557c3 100644 --- a/utils/DFLJPG.py +++ b/utils/DFLJPG.py @@ -1,9 +1,13 @@ -import struct import pickle +import struct + +import cv2 import numpy as np + from facelib import FaceType from imagelib import IEPolys from utils.struct_utils import * +from interact import interact as io class DFLJPG(object): def __init__(self): @@ -137,8 +141,15 @@ class DFLJPG(object): if type(chunk['data']) == bytes: inst.dfl_dict = pickle.loads(chunk['data']) - if (inst.dfl_dict is not None) and ('face_type' not in inst.dfl_dict.keys()): - inst.dfl_dict['face_type'] = FaceType.toString (FaceType.FULL) + if (inst.dfl_dict is not None): + if 'face_type' not in inst.dfl_dict: + inst.dfl_dict['face_type'] = FaceType.toString (FaceType.FULL) + + if 'fanseg_mask' in inst.dfl_dict: + fanseg_mask = inst.dfl_dict['fanseg_mask'] + if fanseg_mask is not None: + numpyarray = np.asarray( inst.dfl_dict['fanseg_mask'], dtype=np.uint8) + inst.dfl_dict['fanseg_mask'] = cv2.imdecode(numpyarray, cv2.IMREAD_UNCHANGED) if inst.dfl_dict == None: return None @@ -155,9 +166,21 @@ class DFLJPG(object): source_filename=None, source_rect=None, source_landmarks=None, - image_to_face_mat=None + image_to_face_mat=None, + fanseg_mask=None, **kwargs ): + if fanseg_mask is not None: + fanseg_mask = np.clip ( (fanseg_mask*255).astype(np.uint8), 0, 255 ) + + ret, buf = cv2.imencode( '.jpg', fanseg_mask, [int(cv2.IMWRITE_JPEG_QUALITY), 85] ) + + if ret and len(buf) < 60000: + fanseg_mask = buf + else: + io.log_err("Unable to encode fanseg_mask for %s" % (filename) ) + fanseg_mask = None + inst = DFLJPG.load_raw (filename) inst.setDFLDictData ({ 'face_type': face_type, @@ -166,7 +189,8 @@ class DFLJPG(object): 'source_filename': source_filename, 'source_rect': source_rect, 'source_landmarks': source_landmarks, - 'image_to_face_mat': image_to_face_mat + 'image_to_face_mat': image_to_face_mat, + 'fanseg_mask' : fanseg_mask, }) try: @@ -181,7 +205,8 @@ class DFLJPG(object): source_filename=None, source_rect=None, source_landmarks=None, - image_to_face_mat=None + image_to_face_mat=None, + fanseg_mask=None, **kwargs ): if face_type is None: face_type = self.get_face_type() if landmarks is None: landmarks = self.get_landmarks() @@ -190,14 +215,18 @@ class DFLJPG(object): if source_rect is None: source_rect = self.get_source_rect() if source_landmarks is None: source_landmarks = self.get_source_landmarks() if image_to_face_mat is None: image_to_face_mat = self.get_image_to_face_mat() + if fanseg_mask is None: fanseg_mask = self.get_fanseg_mask() DFLJPG.embed_data (filename, face_type=face_type, landmarks=landmarks, ie_polys=ie_polys, source_filename=source_filename, source_rect=source_rect, source_landmarks=source_landmarks, - image_to_face_mat=image_to_face_mat) - + image_to_face_mat=image_to_face_mat, + fanseg_mask=fanseg_mask) + def remove_fanseg_mask(self): + self.dfl_dict['fanseg_mask'] = None + def dump(self): data = b"" @@ -252,8 +281,13 @@ 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): + def get_image_to_face_mat(self): mat = self.dfl_dict.get ('image_to_face_mat', None) if mat is not None: return np.array (mat) - return None \ No newline at end of file + return None + def get_fanseg_mask(self): + fanseg_mask = self.dfl_dict.get ('fanseg_mask', None) + if fanseg_mask is not None: + return np.clip ( np.array (fanseg_mask) / 255.0, 0.0, 1.0 )[...,np.newaxis] + return None diff --git a/utils/DFLPNG.py b/utils/DFLPNG.py index 79ec9c1..a4c668c 100644 --- a/utils/DFLPNG.py +++ b/utils/DFLPNG.py @@ -1,13 +1,16 @@ -PNG_HEADER = b"\x89PNG\r\n\x1a\n" - +import pickle import string import struct import zlib -import pickle + +import cv2 import numpy as np + from facelib import FaceType from imagelib import IEPolys +PNG_HEADER = b"\x89PNG\r\n\x1a\n" + class Chunk(object): def __init__(self, name=None, data=None): self.length = 0 @@ -219,7 +222,7 @@ class DFLPNG(object): self.data = b"" self.length = 0 self.chunks = [] - self.fcwp_dict = None + self.dfl_dict = None @staticmethod def load_raw(filename): @@ -252,12 +255,19 @@ class DFLPNG(object): def load(filename): try: inst = DFLPNG.load_raw (filename) - inst.fcwp_dict = inst.getDFLDictData() + inst.dfl_dict = inst.getDFLDictData() - if (inst.fcwp_dict is not None) and ('face_type' not in inst.fcwp_dict.keys()): - inst.fcwp_dict['face_type'] = FaceType.toString (FaceType.FULL) + if inst.dfl_dict is not None: + if 'face_type' not in inst.dfl_dict: + inst.dfl_dict['face_type'] = FaceType.toString (FaceType.FULL) - if inst.fcwp_dict == None: + if 'fanseg_mask' in inst.dfl_dict: + fanseg_mask = inst.dfl_dict['fanseg_mask'] + if fanseg_mask is not None: + numpyarray = np.asarray( inst.dfl_dict['fanseg_mask'], dtype=np.uint8) + inst.dfl_dict['fanseg_mask'] = cv2.imdecode(numpyarray, cv2.IMREAD_UNCHANGED) + + if inst.dfl_dict == None: return None return inst @@ -272,9 +282,21 @@ class DFLPNG(object): source_filename=None, source_rect=None, source_landmarks=None, - image_to_face_mat=None + image_to_face_mat=None, + fanseg_mask=None, **kwargs ): + if fanseg_mask is not None: + fanseg_mask = np.clip ( (fanseg_mask*255).astype(np.uint8), 0, 255 ) + + ret, buf = cv2.imencode( '.jpg', fanseg_mask, [int(cv2.IMWRITE_JPEG_QUALITY), 85] ) + + if ret and len(buf) < 60000: + fanseg_mask = buf + else: + io.log_err("Unable to encode fanseg_mask for %s" % (filename) ) + fanseg_mask = None + inst = DFLPNG.load_raw (filename) inst.setDFLDictData ({ 'face_type': face_type, @@ -283,7 +305,8 @@ class DFLPNG(object): 'source_filename': source_filename, 'source_rect': source_rect, 'source_landmarks': source_landmarks, - 'image_to_face_mat':image_to_face_mat + 'image_to_face_mat':image_to_face_mat, + 'fanseg_mask' : fanseg_mask, }) try: @@ -292,13 +315,14 @@ class DFLPNG(object): except: raise Exception( 'cannot save %s' % (filename) ) - def embed_and_set(self, filename, face_type=None, - landmarks=None, - ie_polys=None, - source_filename=None, - source_rect=None, - source_landmarks=None, - image_to_face_mat=None + def embed_and_set(self, filename, face_type=None, + landmarks=None, + ie_polys=None, + source_filename=None, + source_rect=None, + source_landmarks=None, + image_to_face_mat=None, + fanseg_mask=None, **kwargs ): if face_type is None: face_type = self.get_face_type() if landmarks is None: landmarks = self.get_landmarks() @@ -307,13 +331,18 @@ class DFLPNG(object): if source_rect is None: source_rect = self.get_source_rect() if source_landmarks is None: source_landmarks = self.get_source_landmarks() if image_to_face_mat is None: image_to_face_mat = self.get_image_to_face_mat() + if fanseg_mask is None: fanseg_mask = self.get_fanseg_mask() DFLPNG.embed_data (filename, face_type=face_type, landmarks=landmarks, ie_polys=ie_polys, source_filename=source_filename, source_rect=source_rect, source_landmarks=source_landmarks, - image_to_face_mat=image_to_face_mat) + image_to_face_mat=image_to_face_mat, + fanseg_mask=fanseg_mask) + + def remove_fanseg_mask(self): + self.dfl_dict['fanseg_mask'] = None def dump(self): data = PNG_HEADER @@ -352,17 +381,21 @@ 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_ie_polys(self): return IEPolys.load(self.fcwp_dict.get('ie_polys',None)) - 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.dfl_dict['face_type'] + def get_landmarks(self): return np.array ( self.dfl_dict['landmarks'] ) + def get_ie_polys(self): return IEPolys.load(self.dfl_dict.get('ie_polys',None)) + 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.fcwp_dict.get ('image_to_face_mat', None) + mat = self.dfl_dict.get ('image_to_face_mat', None) if mat is not None: return np.array (mat) return None - + def get_fanseg_mask(self): + fanseg_mask = self.dfl_dict.get ('fanseg_mask', None) + if fanseg_mask is not None: + return np.clip ( np.array (fanseg_mask) / 255.0, 0.0, 1.0 )[...,np.newaxis] + return None def __str__(self): return "".format(len(self.chunks), **self.__dict__)