added FANseg extractor for src and dst faces to use it in training

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 = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu') (x)
-            x = Concatenate(axis=3)([ x, x4])            
+            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 = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu') (x)
-            
+
-            x = Conv2DTranspose (ngf*2, 3, strides=2, padding='same', activation='relu') (x)            
+            x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu') (x)
-            x = Concatenate(axis=3)([ x, x2])            
+            x = Concatenate(axis=3)([ x, x3])
-            x = Conv2D (ngf*4, 3, strides=1, padding='same', activation='relu') (x)     
+            x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu') (x)
-                   
+
-            x = Conv2DTranspose (ngf, 3, strides=2, padding='same', activation='relu') (x)            
+            x = Conv2DTranspose (ngf*2, 3, strides=2, padding='same', activation='relu') (x)
-            x = Concatenate(axis=3)([ x, x1])            
+            x = Concatenate(axis=3)([ x, x2])
-            x = Conv2D (ngf*2, 3, strides=1, padding='same', activation='relu') (x)     
+            x = Conv2D (ngf*4, 3, strides=1, padding='same', activation='relu') (x)
-                   
+
-            x = Conv2DTranspose (ngf // 2, 3, strides=2, padding='same', activation='relu') (x)            
+            x = Conv2DTranspose (ngf, 3, strides=2, padding='same', activation='relu') (x)
-            x = Concatenate(axis=3)([ x, x0])            
+            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

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):

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,

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) )

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)
 

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_S   = sample[0][1][0:4] #first 4 samples
-        test_A_m = sample[0][2][0:4] #first 4 samples
+        test_S_m = sample[0][2][0:4] #first 4 samples
-        test_B   = sample[1][1][0:4]
+        test_D   = sample[1][1][0:4]
-        test_B_m = sample[1][2][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)):
+            for i in range(0, len(test_S)):
-                ar = S[i], SS[i], D[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i])
+                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 )), ]

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

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() ) )
 

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)

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()):
+            if (inst.dfl_dict is not None):
-                inst.dfl_dict['face_type'] = FaceType.toString (FaceType.FULL)
+                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
+        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

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()):
+            if inst.dfl_dict is not None:
-                inst.fcwp_dict['face_type'] = FaceType.toString (FaceType.FULL)
+                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,
+    def embed_and_set(self, filename,   face_type=None,
-                                landmarks=None,
+                                        landmarks=None,
-                                ie_polys=None,
+                                        ie_polys=None,
-                                source_filename=None,
+                                        source_filename=None,
-                                source_rect=None,
+                                        source_rect=None,
-                                source_landmarks=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()
@@ -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_face_type(self): return self.dfl_dict['face_type']
-    def get_landmarks(self): return np.array ( self.fcwp_dict['landmarks'] )
+    def get_landmarks(self): return np.array ( self.dfl_dict['landmarks'] )
-    def get_ie_polys(self): return IEPolys.load(self.fcwp_dict.get('ie_polys',None))
+    def get_ie_polys(self): return IEPolys.load(self.dfl_dict.get('ie_polys',None))
-    def get_source_filename(self): return self.fcwp_dict['source_filename']
+    def get_source_filename(self): return self.dfl_dict['source_filename']
-    def get_source_rect(self): return self.fcwp_dict['source_rect']
+    def get_source_rect(self): return self.dfl_dict['source_rect']
-    def get_source_landmarks(self): return np.array ( self.fcwp_dict['source_landmarks'] )
+    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 "<PNG length={length} chunks={}>".format(len(self.chunks), **self.__dict__)