update == 04.20.2019 == (#242)

* superb improved fanseg

* _

* _

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

* -

* _

* _

* update to 'partial' func

* _

* trained FANSeg_256_full_face.h5,
new experimental models: AVATAR, RecycleGAN

* _

* _

* _

* fix for TCC mode cards(tesla), was conflict with plaidML initialization.

* _

* update manuals

* _

converters/Converter.py

@@ -5,8 +5,9 @@ You can implement your own Converter, check example ConverterMasked.py
 
 class Converter(object):
     TYPE_FACE = 0                   #calls convert_face
-    TYPE_IMAGE = 1                  #calls convert_image without landmarks
+    TYPE_FACE_AVATAR = 1            #calls convert_face with avatar_operator_face
-    TYPE_IMAGE_WITH_LANDMARKS = 2   #calls convert_image with landmarks
+    TYPE_IMAGE = 2                  #calls convert_image without landmarks
+    TYPE_IMAGE_WITH_LANDMARKS = 3   #calls convert_image with landmarks
 
     #overridable
     def __init__(self, predictor_func, type):
@@ -23,13 +24,13 @@ class Converter(object):
         pass
 
     #overridable
-    def cli_convert_face (self, img_bgr, img_face_landmarks, debug):
+    def cli_convert_face (self, img_bgr, img_face_landmarks, debug, avaperator_face_bgr=None, **kwargs):
         #return float32 image
         #if debug , return tuple ( images of any size and channels, ...)
         return image
         
     #overridable
-    def convert_image (self, img_bgr, img_landmarks, debug):
+    def cli_convert_image (self, img_bgr, img_landmarks, debug):
         #img_landmarks not None, if input image is png with embedded data
         #return float32 image
         #if debug , return tuple ( images of any size and channels, ...)

converters/ConverterAvatar.py

@@ -0,0 +1,70 @@
+import time
+
+import cv2
+import numpy as np
+
+from facelib import FaceType, LandmarksProcessor
+from joblib import SubprocessFunctionCaller
+from utils.pickle_utils import AntiPickler
+
+from .Converter import Converter
+
+class ConverterAvatar(Converter):
+
+    #override
+    def __init__(self,  predictor_func,
+                        predictor_input_size=0):
+
+        super().__init__(predictor_func, Converter.TYPE_FACE_AVATAR)
+
+        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 ) )
+        time.sleep(2)
+
+        predictor_func_host, predictor_func = SubprocessFunctionCaller.make_pair(predictor_func)
+        self.predictor_func_host = AntiPickler(predictor_func_host)
+        self.predictor_func = predictor_func
+
+    #overridable
+    def on_host_tick(self):
+        self.predictor_func_host.obj.process_messages()
+        
+    #override
+    def cli_convert_face (self, img_bgr, img_face_landmarks, debug, avaperator_face_bgr=None, **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)
+        
+        output_size = self.predictor_input_size        
+        face_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, output_size, face_type=FaceType.FULL)
+
+        dst_face_mask_a_0 = cv2.warpAffine( img_face_mask_a, face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC )
+
+        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]
+
+        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)
+        
+        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

converters/ConverterImage.py

@@ -1,40 +1,50 @@
-from .Converter import Converter
+import time
-from facelib import LandmarksProcessor
-from facelib import FaceType
 
 import cv2
 import numpy as np
 
-'''
+from facelib import FaceType, LandmarksProcessor
-predictor_func:
+from joblib import SubprocessFunctionCaller
-    input:  [predictor_input_size, predictor_input_size, BGR]
+from utils.pickle_utils import AntiPickler
-    output: [predictor_input_size, predictor_input_size, BGR]
+
-'''
+from .Converter import Converter
 
 class ConverterImage(Converter):
 
     #override
     def __init__(self,  predictor_func,
-                        predictor_input_size=0,
+                        predictor_input_size=0):
-                        output_size=0):
 
         super().__init__(predictor_func, Converter.TYPE_IMAGE)
 
         self.predictor_input_size = predictor_input_size
-        self.output_size = output_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 ) )
+        time.sleep(2)
+
+        predictor_func_host, predictor_func = SubprocessFunctionCaller.make_pair(predictor_func)
+        self.predictor_func_host = AntiPickler(predictor_func_host)
+        self.predictor_func = predictor_func
+
+    #overridable
+    def on_host_tick(self):
+        self.predictor_func_host.obj.process_messages()
         
     #override
-    def dummy_predict(self):
+    def cli_convert_image (self, img_bgr, img_landmarks, debug):
-        self.predictor_func ( np.zeros ( (self.predictor_input_size, self.predictor_input_size,3), dtype=np.float32) )
-
-    #override
-    def convert_image (self, img_bgr, img_landmarks, debug):
         img_size = img_bgr.shape[1], img_bgr.shape[0]
 
         predictor_input_bgr = cv2.resize ( img_bgr, (self.predictor_input_size, self.predictor_input_size), cv2.INTER_LANCZOS4 )
-        predicted_bgr = self.predictor_func ( predictor_input_bgr )
         
-        output = cv2.resize ( predicted_bgr, (self.output_size, self.output_size), cv2.INTER_LANCZOS4 )
+        if debug:
+            debugs = [predictor_input_bgr]
+            
+        output = self.predictor_func ( predictor_input_bgr )
+
         if debug:
             return (predictor_input_bgr,output,)
-        return output
+        if debug:
+            debugs += [out_img.copy()]
+
+        return debugs if debug else output

converters/ConverterMasked.py

@@ -30,7 +30,8 @@ class ConverterMasked(Converter):
                         base_blur_mask_modifier = 0,
                         default_erode_mask_modifier = 0,
                         default_blur_mask_modifier = 0,
-                        clip_hborder_mask_per = 0):
+                        clip_hborder_mask_per = 0,
+                        force_mask_mode=-1):
 
         super().__init__(predictor_func, Converter.TYPE_FACE)
 
@@ -76,6 +77,9 @@ class ConverterMasked(Converter):
             if self.mode == 'hist-match' or self.mode == 'hist-match-bw' or self.mode == 'seamless-hist-match':
                 self.hist_match_threshold = np.clip ( io.input_int("Hist match threshold [0..255] (skip:255) :  ", 255), 0, 255)
 
+        if force_mask_mode != -1:
+            self.mask_mode = force_mask_mode
+        else:
             if face_type == FaceType.FULL:
                 self.mask_mode = np.clip ( io.input_int ("Mask mode: (1) learned, (2) dst, (3) FAN-prd, (4) FAN-dst , (5) FAN-prd*FAN-dst (6) learned*FAN-prd*FAN-dst (?) help. Default - %d : " % (1) , 1, help_message="If you learned mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images. 'FAN-prd' - using super smooth mask by pretrained FAN-model from predicted face. 'FAN-dst' - using super smooth mask by pretrained FAN-model from dst face. 'FAN-prd*FAN-dst' or 'learned*FAN-prd*FAN-dst' - using multiplied masks."), 1, 6 )
             else:
@@ -118,10 +122,10 @@ class ConverterMasked(Converter):
     #overridable
     def on_cli_initialize(self):
         if (self.mask_mode >= 3 and self.mask_mode <= 6) and self.fan_seg == None:
-            self.fan_seg = FANSegmentator(256, FaceType.toString(FaceType.FULL) )
+            self.fan_seg = FANSegmentator(256, FaceType.toString( self.face_type ) )
 
     #override
-    def cli_convert_face (self, img_bgr, img_face_landmarks, debug):
+    def cli_convert_face (self, img_bgr, img_face_landmarks, debug, **kwargs):
         if debug:
             debugs = [img_bgr.copy()]
 
@@ -171,13 +175,13 @@ class ConverterMasked(Converter):
 
             if self.mask_mode == 3 or self.mask_mode == 5 or self.mask_mode == 6: 
                 prd_face_bgr_256 = cv2.resize (prd_face_bgr, (256,256) )
-                prd_face_bgr_256_mask = self.fan_seg.extract_from_bgr( prd_face_bgr_256[np.newaxis,...] ) [0]
+                prd_face_bgr_256_mask = self.fan_seg.extract( prd_face_bgr_256 )
                 FAN_prd_face_mask_a_0 = cv2.resize (prd_face_bgr_256_mask, (output_size,output_size), cv2.INTER_CUBIC)
 
             if self.mask_mode == 4 or self.mask_mode == 5 or self.mask_mode == 6: 
                 face_256_mat     = LandmarksProcessor.get_transform_mat (img_face_landmarks, 256, face_type=FaceType.FULL)
                 dst_face_256_bgr = cv2.warpAffine(img_bgr, face_256_mat, (256, 256), flags=cv2.INTER_LANCZOS4 )
-                dst_face_256_mask = self.fan_seg.extract_from_bgr( dst_face_256_bgr[np.newaxis,...] ) [0]
+                dst_face_256_mask = self.fan_seg.extract( dst_face_256_bgr )
                 FAN_dst_face_mask_a_0 = cv2.resize (dst_face_256_mask, (output_size,output_size), cv2.INTER_CUBIC)
 
             if self.mask_mode == 3:   #FAN-prd

converters/__init__.py

@@ -1,3 +1,4 @@
 from .Converter import Converter
 from .ConverterMasked import ConverterMasked
 from .ConverterImage import ConverterImage
+from .ConverterAvatar import ConverterAvatar

facelib/FANSegmentator.py

@@ -1,15 +1,30 @@
-import numpy as np
 import os
-import cv2
+import pickle
+from functools import partial
 from pathlib import Path
-from nnlib import nnlib
+
+import cv2
+import numpy as np
+
 from interact import interact as io
+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
+"""
 
 class FANSegmentator(object):
+    VERSION = 1
     def __init__ (self, resolution, face_type_str, load_weights=True, weights_file_root=None, training=False):
         exec( nnlib.import_all(), locals(), globals() )
 
-        self.model = FANSegmentator.BuildModel(resolution, ngf=32)
+        self.model = FANSegmentator.BuildModel(resolution, ngf=64)
 
         if weights_file_root:
             weights_file_root = Path(weights_file_root)
@@ -22,13 +37,20 @@ class FANSegmentator(object):
             self.model.load_weights (str(self.weights_path))
         else:
             if training:
-                conv_weights_list = []
+                try:
-                for layer in self.model.layers:
+                    with open( Path(__file__).parent / 'vgg11_enc_weights.npy', 'rb' ) as f:
-                    if type(layer) == keras.layers.Conv2D:
+                        d = pickle.loads (f.read())
-                        conv_weights_list += [layer.weights[0]]  # Conv2D kernel_weights
+
-                CAInitializerMP(conv_weights_list)
+                    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))
+            #self.model.compile(loss='mse', optimizer=Adam(tf_cpu_mode=2))
+            self.model.compile(loss='binary_crossentropy', optimizer=Adam(tf_cpu_mode=2), metrics=['accuracy'])
             
     def __enter__(self):
         return self
@@ -42,66 +64,76 @@ class FANSegmentator(object):
     def train_on_batch(self, inp, outp):
         return self.model.train_on_batch(inp, outp)
 
-    def extract_from_bgr (self, input_image):
+    def extract (self, input_image, is_input_tanh=False):
-        return np.clip ( (self.model.predict(input_image) + 1) / 2.0, 0, 1.0 )
+        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 )
+        result[result < 0.1] = 0 #get rid of noise
+
+        if input_shape_len == 3:
+            result = result[0]
+
+        return result
 
     @staticmethod
-    def BuildModel ( resolution, ngf=64):
+    def BuildModel ( resolution, ngf=64, norm='', act='lrelu'):
         exec( nnlib.import_all(), locals(), globals() )
         inp = Input ( (resolution,resolution,3) )
         x = inp
-        x = FANSegmentator.EncFlow(ngf=ngf)(x)
+        x = FANSegmentator.Flow(ngf=ngf, norm=norm, act=act)(x)
-        x = FANSegmentator.DecFlow(ngf=ngf)(x)
         model = Model(inp,x)
         return model
 
     @staticmethod
-    def EncFlow(ngf=64, num_downs=4):
+    def Flow(ngf=64, num_downs=4, norm='', act='lrelu'):
         exec( nnlib.import_all(), locals(), globals() )
 
-        use_bias = True
-        def XNormalization(x):
-            return InstanceNormalization (axis=3, gamma_initializer=RandomNormal(1., 0.02))(x)
-
-        def downscale (dim):
-            def func(x):
-                return LeakyReLU(0.1)(XNormalization(Conv2D(dim, kernel_size=5, strides=2, padding='same', kernel_initializer=RandomNormal(0, 0.02))(x)))
-            return func
-
         def func(input):
             x = input
 
-            result = []
+            x0 = x = Conv2D(ngf, kernel_size=3, strides=1, padding='same', activation='relu', name='features.0')(x)
-            for i in range(num_downs):
+            x = MaxPooling2D()(x)
-               x = downscale ( min(ngf*(2**i), ngf*8) )(x)
+
-               result += [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 = 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 result
-        return func
-
-    @staticmethod
-    def DecFlow(output_nc=1, ngf=64, activation='tanh'):
-        exec (nnlib.import_all(), locals(), globals())
-
-        use_bias = True
-        def XNormalization(x):
-            return InstanceNormalization (axis=3, gamma_initializer=RandomNormal(1., 0.02))(x)
-
-        def Conv2D (filters, kernel_size, strides=(1, 1), padding='valid', data_format=None, dilation_rate=(1, 1), activation=None, use_bias=use_bias, kernel_initializer=RandomNormal(0, 0.02), bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, kernel_constraint=None, bias_constraint=None):
-            return keras.layers.Conv2D( filters=filters, kernel_size=kernel_size, strides=strides, padding=padding, data_format=data_format, dilation_rate=dilation_rate, activation=activation, use_bias=use_bias, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, kernel_regularizer=kernel_regularizer, bias_regularizer=bias_regularizer, activity_regularizer=activity_regularizer, kernel_constraint=kernel_constraint, bias_constraint=bias_constraint )
-
-        def upscale (dim):
-            def func(x):
-                return SubpixelUpscaler()( LeakyReLU(0.1)(XNormalization(Conv2D(dim, kernel_size=3, strides=1, padding='same', kernel_initializer=RandomNormal(0, 0.02))(x))))
-            return func
-
-        def func(input):
-            input_len = len(input)
-            x = input[input_len-1]
-            for i in range(input_len-1, -1, -1):
-                x = upscale( min(ngf* (2**i) *4, ngf*8 *4 ) )(x)
-                if i != 0:
-                    x = Concatenate(axis=3)([ input[i-1] , x])
-
-            return Conv2D(output_nc, 3, 1, 'same', activation=activation)(x)
         return func

joblib/SubprocessFunctionCaller.py

@@ -8,9 +8,9 @@ class SubprocessFunctionCaller(object):
             self.c2s = c2s
             self.lock = lock
 
-        def __call__(self, value):
+        def __call__(self, *args, **kwargs):
             self.lock.acquire()
-            self.c2s.put (value)
+            self.c2s.put ( {'args':args, 'kwargs':kwargs} )
             while True:
                 if not self.s2c.empty():
                     obj = self.s2c.get()
@@ -27,7 +27,7 @@ class SubprocessFunctionCaller(object):
         def process_messages(self):
             while not self.c2s.empty():
                 obj = self.c2s.get()
-                result = self.func (obj)
+                result = self.func ( *obj['args'], **obj['kwargs'] )
                 self.s2c.put (result)
 
     @staticmethod

main.py

@@ -49,6 +49,23 @@ 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
@@ -72,11 +89,16 @@ 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):
@@ -112,6 +134,7 @@ 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,
@@ -125,7 +148,8 @@ if __name__ == "__main__":
     p = subparsers.add_parser( "convert", help="Converter")
     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. Not used in AVATAR model.")
+    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.")

mainscripts/Converter.py

@@ -1,19 +1,23 @@
 import sys
+import multiprocessing
+import operator
 import os
+import shutil
+import time
 import traceback
 from pathlib import Path
-from utils import Path_utils
+
 import cv2
-from utils.DFLPNG import DFLPNG
-from utils.DFLJPG import DFLJPG
-from utils.cv2_utils import *
-import shutil
 import numpy as np
-import time
+
-import multiprocessing
 from converters import Converter
-from joblib import Subprocessor, SubprocessFunctionCaller
 from interact import interact as io
+from joblib import SubprocessFunctionCaller, Subprocessor
+from utils import Path_utils
+from utils.cv2_utils import *
+from utils.DFLJPG import DFLJPG
+from utils.DFLPNG import DFLPNG
+from imagelib import normalize_channels
 
 class ConvertSubprocessor(Subprocessor):
     class Cli(Subprocessor.Cli):
@@ -26,6 +30,7 @@ class ConvertSubprocessor(Subprocessor):
             self.converter   = client_dict['converter']
             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
@@ -45,13 +50,15 @@ class ConvertSubprocessor(Subprocessor):
 
         #override
         def process_data(self, data):
-            filename_path = Path(data)
+            idx, filename = data
+            filename_path = Path(filename)
             files_processed = 1
             faces_processed = 0
 
             output_filename_path = self.output_path / (filename_path.stem + '.png')
 
-            if self.converter.type == Converter.TYPE_FACE and filename_path.stem not in self.alignments.keys():
+            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) )
 
@@ -62,19 +69,18 @@ class ConvertSubprocessor(Subprocessor):
                         cv2_imwrite ( str(output_filename_path), image )
             else:
                 image = (cv2_imread(str(filename_path)) / 255.0).astype(np.float32)
-                h,w,c = image.shape
+                image = normalize_channels (image, 3)
-                if c > 3:
-                    image = image[...,0:3]
                 
                 if self.converter.type == Converter.TYPE_IMAGE:
-                    image = self.converter.convert_image(image, None, self.debug)
+                    image = self.converter.cli_convert_image(image, None, self.debug)
+
                     if self.debug:
-                        raise NotImplementedError
+                        return (1, image)
-                        #for img in image:
+                        
-                        #    io.show_image ('Debug convert', img )
-                        #    cv2.waitKey(0)
                     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':
@@ -96,7 +102,13 @@ class ConvertSubprocessor(Subprocessor):
                     else:
                         self.log_err ("%s is not a dfl image file" % (filename_path.name) )
 
-                elif self.converter.type == Converter.TYPE_FACE:
+                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:
@@ -108,9 +120,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)
+                                debug_images += self.converter.cli_convert_face(image, image_landmarks, self.debug, avaperator_face_bgr=ava_face)
                             else:
-                                image = self.converter.cli_convert_face(image, image_landmarks, self.debug)
+                                image = self.converter.cli_convert_face(image, image_landmarks, self.debug, avaperator_face_bgr=ava_face)
 
                         except Exception as e:
                             e_str = traceback.format_exc()
@@ -133,16 +145,19 @@ class ConvertSubprocessor(Subprocessor):
         #overridable
         def get_data_name (self, data):
             #return string identificator of your data
-            return data
+            idx, filename = data
+            return filename
 
     #override
-    def __init__(self, converter, input_path_image_paths, output_path, alignments, debug = False):
+    def __init__(self, converter, input_path_image_paths, output_path, alignments, avatar_image_paths=None, 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_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
@@ -158,6 +173,7 @@ class ConvertSubprocessor(Subprocessor):
                                       'converter' : self.converter,
                                       '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
                                       }
@@ -167,7 +183,7 @@ class ConvertSubprocessor(Subprocessor):
         if self.debug:
             io.named_window ("Debug convert")
 
-        io.progress_bar ("Converting", len (self.input_data) )
+        io.progress_bar ("Converting", len (self.input_data_idxs) )
 
     #overridable optional
     def on_clients_finalized(self):
@@ -178,13 +194,15 @@ class ConvertSubprocessor(Subprocessor):
 
     #override
     def get_data(self, host_dict):
-        if len (self.input_data) > 0:
+        if len (self.input_data_idxs) > 0:
-            return self.input_data.pop(0)
+            idx = self.input_data_idxs.pop(0)
+            return (idx, self.input_data[idx])
         return None
 
     #override
     def on_data_return (self, host_dict, data):
-        self.input_data.insert(0, data)
+        idx, filename = data
+        self.input_data_idxs.insert(0, idx)
 
     #override
     def on_result (self, host_dict, data, result):
@@ -209,6 +227,7 @@ def main (args, device_args):
     io.log_info ("Running converter.\r\n")
 
     aligned_dir = args.get('aligned_dir', None)
+    avaperator_aligned_dir = args.get('avaperator_aligned_dir', None)
     
     try:
         input_path = Path(args['input_dir'])
@@ -233,9 +252,10 @@ def main (args, device_args):
         model = models.import_model( args['model_name'] )(model_path, device_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:
+        if converter.type == Converter.TYPE_FACE or converter.type == Converter.TYPE_FACE_AVATAR:
             if aligned_dir is None:
                 io.log_err('Aligned directory not found. Please ensure it exists.')
                 return
@@ -268,11 +288,44 @@ def main (args, device_args):
 
                 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
+
+            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':
+                    dflimg = DFLJPG.load ( str(filepath) )
+                else:
+                    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)) ]
+                    
+            if len(input_path_image_paths) < len(avatar_image_paths):
+                io.log_err("Input faces count must be >= avatar operator faces count.")
+                return
+                
         files_processed, faces_processed = ConvertSubprocessor (
                     converter              = converter,
-                    input_path_image_paths = Path_utils.get_image_paths(input_path),
+                    input_path_image_paths = input_path_image_paths,                    
                     output_path            = output_path,
                     alignments             = alignments,
+                    avatar_image_paths     = avatar_image_paths,
                     debug                  = args.get('debug',False)
                     ).run()
 

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
@@ -80,6 +82,11 @@ class ExtractSubprocessor(Subprocessor):
                 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) )
 
@@ -254,6 +263,13 @@ class ExtractSubprocessor(Subprocessor):
 
                 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):
@@ -261,7 +277,7 @@ class ExtractSubprocessor(Subprocessor):
             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:
@@ -658,6 +674,33 @@ class DeletedFilesSearcherSubprocessor(Subprocessor):
         return self.result
 
 
+#currently unused
+def extract_fanseg(input_dir, device_args={} ):
+    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,

mainscripts/MaskEditorTool.py

@@ -397,12 +397,16 @@ 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
 
+                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) )

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_FANSEG/Model.py

@@ -17,13 +17,22 @@ class Model(ModelBase):
                             ask_random_flip=False,
                             ask_src_scale_mod=False)
     
+    #override
+    def onInitializeOptions(self, is_first_run, ask_override):            
+        default_face_type = 'f'
+        if is_first_run:
+            self.options['face_type'] = io.input_str ("Half or Full face? (h/f, ?:help skip:f) : ", default_face_type, ['h','f'], help_message="Half face has better resolution, but covers less area of cheeks.").lower()
+        else:
+            self.options['face_type'] = self.options.get('face_type', default_face_type)
+     
     #override
     def onInitialize(self):
         exec(nnlib.import_all(), locals(), globals())
         self.set_vram_batch_requirements( {1.5:4} )
 
         self.resolution = 256
-        self.face_type = FaceType.FULL
+        self.face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
+
         
         self.fan_seg = FANSegmentator(self.resolution, 
                                       FaceType.toString(self.face_type), 
@@ -33,18 +42,18 @@ class Model(ModelBase):
 
         if self.is_training_mode:
             f = SampleProcessor.TypeFlags
-            f_type = f.FACE_TYPE_FULL
+            face_type = f.FACE_TYPE_FULL if self.options['face_type'] == 'f' else f.FACE_TYPE_HALF
             
             self.set_training_data_generators ([    
                     SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, 
-                            sample_process_options=SampleProcessor.Options(random_flip=True, motion_blur = [25, 1], normalize_tanh = True ), 
+                            sample_process_options=SampleProcessor.Options(random_flip=True, motion_blur = [25, 1] ), 
-                            output_sample_types=[ [f.TRANSFORMED | f_type | f.MODE_BGR_SHUFFLE | f.OPT_APPLY_MOTION_BLUR, self.resolution],
+                            output_sample_types=[ [f.WARPED_TRANSFORMED | face_type | f.MODE_BGR_SHUFFLE | f.OPT_APPLY_MOTION_BLUR, self.resolution],
-                                                  [f.TRANSFORMED | f_type | f.MODE_M | f.FACE_MASK_FULL, self.resolution]
+                                                  [f.WARPED_TRANSFORMED | face_type | f.MODE_M | f.FACE_MASK_FULL, self.resolution]
                                                 ]),
                                                 
                     SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, 
-                            sample_process_options=SampleProcessor.Options(random_flip=True, normalize_tanh = True ), 
+                            sample_process_options=SampleProcessor.Options(random_flip=True ), 
-                            output_sample_types=[ [f.TRANSFORMED | f_type | f.MODE_BGR_SHUFFLE, self.resolution]
+                            output_sample_types=[ [f.TRANSFORMED | face_type | f.MODE_BGR_SHUFFLE, self.resolution]
                                                 ])
                                                ])
                 
@@ -56,19 +65,17 @@ class Model(ModelBase):
     def onTrainOneIter(self, generators_samples, generators_list):
         target_src, target_src_mask = generators_samples[0]
 
-        loss = self.fan_seg.train_on_batch( [target_src], [target_src_mask] )
+        loss,acc = self.fan_seg.train_on_batch( [target_src], [target_src_mask] )
 
-        return ( ('loss', loss), )
+        return ( ('loss', loss), ('acc',acc))
         
     #override
     def onGetPreview(self, sample):
         test_A   = sample[0][0][0:4] #first 4 samples
         test_B   = sample[1][0][0:4] #first 4 samples
         
-        mAA = self.fan_seg.extract_from_bgr([test_A])
+        mAA = self.fan_seg.extract(test_A)
-        mBB = self.fan_seg.extract_from_bgr([test_B])
+        mBB = self.fan_seg.extract(test_B)
-        
-        test_A, test_B, = [ np.clip( (x + 1.0)/2.0, 0.0, 1.0)  for x in [test_A, test_B] ]
 
         mAA = np.repeat ( mAA, (3,), -1)
         mBB = np.repeat ( mBB, (3,), -1)
@@ -89,6 +96,6 @@ class Model(ModelBase):
                 test_B[i,:,:,0:3]*mBB[i],
                 ), axis=1) )
                 
-        return [ ('FANSegmentator', np.concatenate ( st, axis=0 ) ),
+        return [ ('training data', np.concatenate ( st, axis=0 ) ),
-                 ('never seen', np.concatenate ( st2, axis=0 ) ),
+                 ('evaluating data', np.concatenate ( st2, axis=0 ) ),
                  ]

models/Model_RecycleGAN/Model.py

@@ -0,0 +1,477 @@
+from functools import partial
+
+import cv2
+import numpy as np
+
+from facelib import FaceType
+from interact import interact as io
+from mathlib import get_power_of_two
+from models import ModelBase
+from nnlib import nnlib
+from samplelib import *
+
+class RecycleGANModel(ModelBase):
+    
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs, 
+                            ask_sort_by_yaw=False,
+                            ask_random_flip=False,
+                            ask_src_scale_mod=False)
+                            
+    #override
+    def onInitializeOptions(self, is_first_run, ask_override):
+        if is_first_run:
+            self.options['resolution'] = io.input_int("Resolution ( 128,256 ?:help skip:128) : ", 128, [128,256], help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16.")
+        else:
+            self.options['resolution'] = self.options.get('resolution', 128)
+    
+    #override
+    def onInitialize(self, batch_size=-1, **in_options):
+        exec(nnlib.code_import_all, locals(), globals())
+        self.set_vram_batch_requirements({6:16})
+        
+        resolution = self.options['resolution']
+        bgr_shape = (resolution, resolution, 3)
+        ngf = 64
+        npf = 32
+        ndf = 64
+        lambda_A = 10
+        lambda_B = 10
+
+        use_batch_norm = True #created_batch_size > 1
+        self.GA = modelify(RecycleGANModel.ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=True))(Input(bgr_shape))
+        self.GB = modelify(RecycleGANModel.ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=True))(Input(bgr_shape))
+
+        #self.GA = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape))
+        #self.GB = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape))
+        
+        self.PA = modelify(RecycleGANModel.UNetTemporalPredictor(bgr_shape[2], use_batch_norm, ngf=npf))([Input(bgr_shape), Input(bgr_shape)])
+        self.PB = modelify(RecycleGANModel.UNetTemporalPredictor(bgr_shape[2], use_batch_norm, ngf=npf))([Input(bgr_shape), Input(bgr_shape)])
+
+        self.DA = modelify(RecycleGANModel.PatchDiscriminator(ndf=ndf) ) (Input(bgr_shape))
+        self.DB = modelify(RecycleGANModel.PatchDiscriminator(ndf=ndf) ) (Input(bgr_shape))
+
+        if not self.is_first_run():
+            weights_to_load = [
+                (self.GA, 'GA.h5'),
+                (self.DA, 'DA.h5'),
+                (self.PA, 'PA.h5'),
+                (self.GB, 'GB.h5'),
+                (self.DB, 'DB.h5'),
+                (self.PB, 'PB.h5'),
+            ]
+            self.load_weights_safe(weights_to_load)
+            
+        real_A0 = Input(bgr_shape, name="real_A0")
+        real_A1 = Input(bgr_shape, name="real_A1")
+        real_A2 = Input(bgr_shape, name="real_A2")
+        
+        real_B0 = Input(bgr_shape, name="real_B0")
+        real_B1 = Input(bgr_shape, name="real_B1")
+        real_B2 = Input(bgr_shape, name="real_B2")
+        
+        DA_ones =  K.ones_like ( K.shape(self.DA.outputs[0]) )
+        DA_zeros = K.zeros_like ( K.shape(self.DA.outputs[0] ))
+        DB_ones = K.ones_like ( K.shape(self.DB.outputs[0] ))
+        DB_zeros = K.zeros_like ( K.shape(self.DB.outputs[0] ))
+        
+        def DLoss(labels,logits):
+            return K.mean(K.binary_crossentropy(labels,logits))
+
+        def CycleLoss (t1,t2):
+            return K.mean(K.abs(t1 - t2))
+        
+        def RecurrentLOSS(t1,t2):
+            return K.mean(K.abs(t1 - t2))
+            
+        def RecycleLOSS(t1,t2):
+            return K.mean(K.abs(t1 - t2))
+            
+        fake_B0 = self.GA(real_A0)
+        fake_B1 = self.GA(real_A1)
+        
+        fake_A0 = self.GB(real_B0)      
+        fake_A1 = self.GB(real_B1)
+        
+        real_A0_d = self.DA(real_A0)
+        real_A0_d_ones = K.ones_like(real_A0_d)
+        real_A1_d = self.DA(real_A1)
+        real_A1_d_ones = K.ones_like(real_A1_d)
+        
+        fake_A0_d = self.DA(fake_A0)
+        fake_A0_d_ones = K.ones_like(fake_A0_d)
+        fake_A0_d_zeros = K.zeros_like(fake_A0_d)
+        
+        fake_A1_d = self.DA(fake_A1)
+        fake_A1_d_ones = K.ones_like(fake_A1_d)
+        fake_A1_d_zeros = K.zeros_like(fake_A1_d)
+        
+        real_B0_d = self.DB(real_B0)
+        real_B0_d_ones = K.ones_like(real_B0_d)
+        
+        real_B1_d = self.DB(real_B1)
+        real_B1_d_ones = K.ones_like(real_B1_d)
+        
+        fake_B0_d = self.DB(fake_B0)
+        fake_B0_d_ones = K.ones_like(fake_B0_d)
+        fake_B0_d_zeros = K.zeros_like(fake_B0_d)
+        
+        fake_B1_d = self.DB(fake_B1)
+        fake_B1_d_ones = K.ones_like(fake_B1_d)
+        fake_B1_d_zeros = K.zeros_like(fake_B1_d)
+
+        pred_A2 = self.PA ( [real_A0, real_A1])
+        pred_B2 = self.PB ( [real_B0, real_B1])
+        rec_A2 = self.GB ( self.PB ( [fake_B0, fake_B1]) )
+        rec_B2 = self.GA ( self.PA ( [fake_A0, fake_A1]))
+
+                    
+        loss_GA = DLoss(fake_B0_d_ones, fake_B0_d ) + \
+                  DLoss(fake_B1_d_ones, fake_B1_d ) + \
+                  lambda_A * (RecurrentLOSS(pred_A2, real_A2) + \
+                              RecycleLOSS(rec_B2, real_B2) )
+                              
+                              
+        weights_GA = self.GA.trainable_weights + self.PA.trainable_weights
+        
+        loss_GB = DLoss(fake_A0_d_ones, fake_A0_d ) + \
+                  DLoss(fake_A1_d_ones, fake_A1_d ) + \
+                  lambda_B * (RecurrentLOSS(pred_B2, real_B2) + \
+                              RecycleLOSS(rec_A2, real_A2) )
+                              
+        weights_GB = self.GB.trainable_weights + self.PB.trainable_weights
+        
+        def opt():
+            return Adam(lr=2e-4, beta_1=0.5, beta_2=0.999, tf_cpu_mode=2)#, clipnorm=1)
+        
+        self.GA_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_GA],
+                                    opt().get_updates(loss_GA, weights_GA) )
+                                    
+        self.GB_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_GB],
+                                    opt().get_updates(loss_GB, weights_GB) )
+                                    
+        ###########        
+        
+        loss_D_A0 = ( DLoss(real_A0_d_ones, real_A0_d ) + \
+                      DLoss(fake_A0_d_zeros, fake_A0_d ) ) * 0.5
+        
+        loss_D_A1 = ( DLoss(real_A1_d_ones, real_A1_d ) + \
+                      DLoss(fake_A1_d_zeros, fake_A1_d ) ) * 0.5
+                      
+        loss_D_A = loss_D_A0 + loss_D_A1
+        
+        self.DA_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_D_A],
+                                    opt().get_updates(loss_D_A, self.DA.trainable_weights) )
+        
+        ############
+        
+        loss_D_B0 = ( DLoss(real_B0_d_ones, real_B0_d ) + \
+                      DLoss(fake_B0_d_zeros, fake_B0_d ) ) * 0.5
+        
+        loss_D_B1 = ( DLoss(real_B1_d_ones, real_B1_d ) + \
+                      DLoss(fake_B1_d_zeros, fake_B1_d ) ) * 0.5
+                      
+        loss_D_B = loss_D_B0 + loss_D_B1
+        
+        self.DB_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_D_B],
+                                    opt().get_updates(loss_D_B, self.DB.trainable_weights) )
+        
+        ############
+        
+
+        self.G_view = K.function([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[fake_A0, fake_A1, pred_A2, rec_A2, fake_B0, fake_B1, pred_B2, rec_B2 ])
+        
+        
+        
+        if self.is_training_mode:
+            f = SampleProcessor.TypeFlags
+            self.set_training_data_generators ([            
+                    SampleGeneratorImageTemporal(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, 
+                        temporal_image_count=3,
+                        sample_process_options=SampleProcessor.Options(random_flip = False, normalize_tanh = True), 
+                        output_sample_types=[ [f.SOURCE | f.MODE_BGR, resolution] ] ),
+                        
+                    SampleGeneratorImageTemporal(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, 
+                        temporal_image_count=3,
+                        sample_process_options=SampleProcessor.Options(random_flip = False, normalize_tanh = True), 
+                        output_sample_types=[ [f.SOURCE | f.MODE_BGR, resolution] ] ),
+                   ])
+        else:
+            self.G_convert = K.function([real_B0],[fake_A0])
+            
+    #override
+    def onSave(self):
+        self.save_weights_safe( [[self.GA,   'GA.h5'],
+                                 [self.GB,   'GB.h5'],
+                                 [self.DA,   'DA.h5'],
+                                 [self.DB,   'DB.h5'],
+                                 [self.PA,   'PA.h5'],
+                                 [self.PB,   'PB.h5'] ])
+        
+    #override
+    def onTrainOneIter(self, generators_samples, generators_list):
+        source_src_0, source_src_1, source_src_2, = generators_samples[0]
+        source_dst_0, source_dst_1, source_dst_2, = generators_samples[1]        
+     
+        feed = [source_src_0, source_src_1, source_src_2, source_dst_0, source_dst_1, source_dst_2]
+
+        loss_GA, = self.GA_train ( feed )
+        loss_GB, = self.GB_train ( feed )
+        loss_DA, = self.DA_train( feed )
+        loss_DB, = self.DB_train( feed )
+        
+        return ( ('GA', loss_GA), ('GB', loss_GB), ('DA', loss_DA),  ('DB', loss_DB)  )
+
+    #override
+    def onGetPreview(self, sample):
+        test_A0   = sample[0][0]
+        test_A1   = sample[0][1]
+        test_A2   = sample[0][2]
+        
+        test_B0   = sample[1][0]
+        test_B1   = sample[1][1]
+        test_B2   = sample[1][2]
+        
+        G_view_result = self.G_view([test_A0, test_A1, test_A2, test_B0, test_B1, test_B2])        
+
+        fake_A0, fake_A1, pred_A2, rec_A2, fake_B0, fake_B1, pred_B2, rec_B2 = [ x[0] / 2 + 0.5 for x in G_view_result]        
+        test_A0, test_A1, test_A2, test_B0, test_B1, test_B2 = [ x[0] / 2 + 0.5 for x in [test_A0, test_A1, test_A2, test_B0, test_B1, test_B2] ]
+
+        r = np.concatenate ((np.concatenate ( (test_A0, test_A1, test_A2, pred_A2, fake_B0, fake_B1, rec_A2), axis=1),
+                             np.concatenate ( (test_B0, test_B1, test_B2, pred_B2, fake_A0, fake_A1, rec_B2), axis=1)
+                             ), axis=0)                            
+                
+        return [ ('RecycleGAN', r ) ]
+    
+    def predictor_func (self, face):
+        x = self.G_convert ( [ face[np.newaxis,...]*2-1 ]  )[0]
+        return np.clip ( x[0] / 2 + 0.5 , 0, 1)
+        
+    #override
+    def get_converter(self, **in_options):
+        from converters import ConverterImage                   
+        return ConverterImage(self.predictor_func, 
+                              predictor_input_size=self.options['resolution'], 
+                              **in_options)
+    
+    @staticmethod
+    def ResNet(output_nc, use_batch_norm, ngf=64, n_blocks=6, use_dropout=False):
+        exec (nnlib.import_all(), locals(), globals())
+
+        if not use_batch_norm:
+            use_bias = True
+            def XNormalization(x):
+                return InstanceNormalization (axis=-1)(x)
+        else:
+            use_bias = False
+            def XNormalization(x):
+                return BatchNormalization (axis=-1)(x)
+                
+        XConv2D = partial(Conv2D, padding='same', use_bias=use_bias)
+        XConv2DTranspose = partial(Conv2DTranspose, padding='same', use_bias=use_bias)
+
+        def func(input):
+
+
+            def ResnetBlock(dim, use_dropout=False):
+                def func(input):
+                    x = input
+
+                    x = XConv2D(dim, 3, strides=1)(x)
+                    x = XNormalization(x)
+                    x = ReLU()(x)
+
+                    if use_dropout:
+                        x = Dropout(0.5)(x)
+
+                    x = XConv2D(dim, 3, strides=1)(x)
+                    x = XNormalization(x)
+                    x = ReLU()(x)
+                    return Add()([x,input])
+                return func
+
+            x = input
+
+            x = ReLU()(XNormalization(XConv2D(ngf, 7, strides=1)(x)))
+
+            x = ReLU()(XNormalization(XConv2D(ngf*2, 3, strides=2)(x)))
+            x = ReLU()(XNormalization(XConv2D(ngf*4, 3, strides=2)(x)))
+
+            for i in range(n_blocks):
+                x = ResnetBlock(ngf*4, use_dropout=use_dropout)(x)
+
+            x = ReLU()(XNormalization(XConv2DTranspose(ngf*2, 3, strides=2)(x)))
+            x = ReLU()(XNormalization(XConv2DTranspose(ngf  , 3, strides=2)(x)))
+
+            x = XConv2D(output_nc, 7, strides=1, activation='tanh', use_bias=True)(x)
+
+            return x
+
+        return func
+        
+    @staticmethod
+    def UNet(output_nc, use_batch_norm, ngf=64, use_dropout=False):
+        exec (nnlib.import_all(), locals(), globals())
+
+        if not use_batch_norm:
+            use_bias = True
+            def XNormalizationL():
+                return InstanceNormalization (axis=-1)
+        else:
+            use_bias = False
+            def XNormalizationL():
+                return BatchNormalization (axis=-1)
+                
+        def XNormalization(x):
+            return XNormalizationL()(x)
+                
+        XConv2D = partial(Conv2D, padding='same', use_bias=use_bias)
+        XConv2DTranspose = partial(Conv2DTranspose, padding='same', use_bias=use_bias)
+      
+        def func(input):
+            
+            b,h,w,c = K.int_shape(input)
+            
+            n_downs = get_power_of_two(w) - 4
+            
+            Norm = XNormalizationL()
+            Norm2 = XNormalizationL()
+            Norm4 = XNormalizationL()
+            Norm8 = XNormalizationL()
+            
+            x = input
+            
+            x = e1 = XConv2D( ngf, 4, strides=2, use_bias=True ) (x)
+
+            x = e2 = Norm2( XConv2D( ngf*2, 4, strides=2  )( LeakyReLU(0.2)(x) ) )
+            x = e3 = Norm4( XConv2D( ngf*4, 4, strides=2  )( LeakyReLU(0.2)(x) ) )
+            
+            l = []
+            for i in range(n_downs):
+                x = Norm8( XConv2D( ngf*8, 4, strides=2  )( LeakyReLU(0.2)(x) ) )
+                l += [x]
+            
+            x = XConv2D( ngf*8, 4, strides=2, use_bias=True  )( LeakyReLU(0.2)(x) )
+            
+            for i in range(n_downs):
+                x = Norm8( XConv2DTranspose( ngf*8, 4, strides=2  )( ReLU()(x) ) )
+                if i <= n_downs-2:
+                    x = Dropout(0.5)(x)                
+                x = Concatenate(axis=-1)([x, l[-i-1] ])
+  
+            x = Norm4( XConv2DTranspose( ngf*4, 4, strides=2  )( ReLU()(x) ) )
+            x = Concatenate(axis=-1)([x, e3])
+
+            x = Norm2( XConv2DTranspose( ngf*2, 4, strides=2  )( ReLU()(x) ) )
+            x = Concatenate(axis=-1)([x, e2])  
+            
+            x = Norm( XConv2DTranspose( ngf, 4, strides=2  )( ReLU()(x) ) )
+            x = Concatenate(axis=-1)([x, e1])   
+            
+            x = XConv2DTranspose(output_nc, 4, strides=2, activation='tanh', use_bias=True)( ReLU()(x) )
+
+            return x
+        return func
+    nnlib.UNet = UNet
+
+    @staticmethod
+    def UNetTemporalPredictor(output_nc, use_batch_norm, ngf=64, use_dropout=False):
+        exec (nnlib.import_all(), locals(), globals())
+        def func(inputs):
+            past_2_image_tensor, past_1_image_tensor = inputs
+
+            x = Concatenate(axis=-1)([ past_2_image_tensor, past_1_image_tensor ])
+            x = UNet(3, use_batch_norm, ngf=ngf, use_dropout=use_dropout) (x)
+
+            return x
+
+        return func
+        
+    @staticmethod
+    def PatchDiscriminator(ndf=64):
+        exec (nnlib.import_all(), locals(), globals())
+
+        #use_bias = True
+        #def XNormalization(x):
+        #    return InstanceNormalization (axis=-1)(x)
+        use_bias = False
+        def XNormalization(x):
+            return BatchNormalization (axis=-1)(x)
+                
+        XConv2D = partial(Conv2D, use_bias=use_bias)
+ 
+        def func(input):
+            b,h,w,c = K.int_shape(input)
+
+            x = input
+
+            x = ZeroPadding2D((1,1))(x)
+            x = XConv2D( ndf, 4, strides=2, padding='valid', use_bias=True)(x)
+            x = LeakyReLU(0.2)(x)
+            
+            x = ZeroPadding2D((1,1))(x)
+            x = XConv2D( ndf*2, 4, strides=2, padding='valid')(x)
+            x = XNormalization(x)
+            x = LeakyReLU(0.2)(x)
+            
+            x = ZeroPadding2D((1,1))(x)           
+            x = XConv2D( ndf*4, 4, strides=2, padding='valid')(x)
+            x = XNormalization(x)
+            x = LeakyReLU(0.2)(x)
+            
+            x = ZeroPadding2D((1,1))(x)
+            x = XConv2D( ndf*8, 4, strides=2, padding='valid')(x)
+            x = XNormalization(x)
+            x = LeakyReLU(0.2)(x)
+
+            x = ZeroPadding2D((1,1))(x)
+            x = XConv2D( ndf*8, 4, strides=2, padding='valid')(x)
+            x = XNormalization(x)
+            x = LeakyReLU(0.2)(x)
+
+            x = ZeroPadding2D((1,1))(x)
+            return XConv2D( 1, 4, strides=1, padding='valid', use_bias=True, activation='sigmoid')(x)#
+        return func
+        
+    @staticmethod
+    def NLayerDiscriminator(ndf=64, n_layers=3):
+        exec (nnlib.import_all(), locals(), globals())
+
+        #use_bias = True
+        #def XNormalization(x):
+        #    return InstanceNormalization (axis=-1)(x)
+        use_bias = False
+        def XNormalization(x):
+            return BatchNormalization (axis=-1)(x)
+                
+        XConv2D = partial(Conv2D, use_bias=use_bias)
+ 
+        def func(input):
+            b,h,w,c = K.int_shape(input)
+
+            x = input
+            
+            f = ndf
+
+            x = ZeroPadding2D((1,1))(x)
+            x = XConv2D( f, 4, strides=2, padding='valid', use_bias=True)(x)
+            f = min( ndf*8, f*2 )
+            x = LeakyReLU(0.2)(x)
+            
+            for i in range(n_layers):
+                x = ZeroPadding2D((1,1))(x)
+                x = XConv2D( f, 4, strides=2, padding='valid')(x)               
+                f = min( ndf*8, f*2 )
+                x = XNormalization(x)
+                x = LeakyReLU(0.2)(x)
+            
+            x = ZeroPadding2D((1,1))(x)
+            x = XConv2D( f, 4, strides=1, padding='valid')(x)
+            x = XNormalization(x)
+            x = LeakyReLU(0.2)(x)
+
+            x = ZeroPadding2D((1,1))(x)
+            return XConv2D( 1, 4, strides=1, padding='valid', use_bias=True, activation='sigmoid')(x)#
+        return func
+        
+Model = RecycleGANModel

models/Model_RecycleGAN/__init__.py

@@ -0,0 +1 @@
+from .Model import Model

models/Model_SAE/Model.py

@@ -1,3 +1,4 @@
+from functools import partial
 import numpy as np
 
 from nnlib import nnlib
@@ -385,20 +386,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 +407,8 @@ 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 )), ]
@@ -485,57 +486,29 @@ class SAEModel(ModelBase):
                 return x
         SAEModel.ResidualBlock = ResidualBlock
 
-        def ResidualBlock_pre (**base_kwargs):
-            def func(*args, **kwargs):
-                kwargs.update(base_kwargs)
-                return ResidualBlock(*args, **kwargs)
-            return func
-        SAEModel.ResidualBlock_pre = ResidualBlock_pre
-
         def downscale (dim, padding='zero', norm='', act='', **kwargs):
             def func(x):
                 return Norm(norm)( Act(act) (Conv2D(dim, kernel_size=5, strides=2, padding=padding)(x)) )
             return func
         SAEModel.downscale = downscale
 
-        def downscale_pre (**base_kwargs):
-            def func(*args, **kwargs):
-                kwargs.update(base_kwargs)
-                return downscale(*args, **kwargs)
-            return func
-        SAEModel.downscale_pre = downscale_pre
-
         def upscale (dim, padding='zero', norm='', act='', **kwargs):
             def func(x):
                 return SubpixelUpscaler()(Norm(norm)(Act(act)(Conv2D(dim * 4, kernel_size=3, strides=1, padding=padding)(x))))
             return func
         SAEModel.upscale = upscale
 
-        def upscale_pre (**base_kwargs):
-            def func(*args, **kwargs):
-                kwargs.update(base_kwargs)
-                return upscale(*args, **kwargs)
-            return func
-        SAEModel.upscale_pre = upscale_pre
-
         def to_bgr (output_nc, padding='zero', **kwargs):
             def func(x):
                 return Conv2D(output_nc, kernel_size=5, padding=padding, activation='sigmoid')(x)
             return func
         SAEModel.to_bgr = to_bgr
 
-        def to_bgr_pre (**base_kwargs):
-            def func(*args, **kwargs):
-                kwargs.update(base_kwargs)
-                return to_bgr(*args, **kwargs)
-            return func
-        SAEModel.to_bgr_pre = to_bgr_pre
-
     @staticmethod
     def LIAEEncFlow(resolution, ch_dims, **kwargs):
         exec (nnlib.import_all(), locals(), globals())
-        upscale = SAEModel.upscale_pre(**kwargs)
+        upscale = partial(SAEModel.upscale, **kwargs)
-        downscale = SAEModel.downscale_pre(**kwargs)
+        downscale = partial(SAEModel.downscale, **kwargs)
 
         def func(input):
             dims = K.int_shape(input)[-1]*ch_dims
@@ -553,7 +526,7 @@ class SAEModel(ModelBase):
     @staticmethod
     def LIAEInterFlow(resolution, ae_dims=256, **kwargs):
         exec (nnlib.import_all(), locals(), globals())
-        upscale = SAEModel.upscale_pre(**kwargs)
+        upscale = partial(SAEModel.upscale, **kwargs)
         lowest_dense_res=resolution // 16
 
         def func(input):
@@ -568,10 +541,10 @@ class SAEModel(ModelBase):
     @staticmethod
     def LIAEDecFlow(output_nc,ch_dims, multiscale_count=1, add_residual_blocks=False, padding='zero', norm='', **kwargs):
         exec (nnlib.import_all(), locals(), globals())
-        upscale = SAEModel.upscale_pre(**kwargs)
+        upscale = partial(SAEModel.upscale, **kwargs)
-        to_bgr = SAEModel.to_bgr_pre(**kwargs)
+        to_bgr = partial(SAEModel.to_bgr, **kwargs)
         dims = output_nc * ch_dims
-        ResidualBlock = SAEModel.ResidualBlock_pre(**kwargs)
+        ResidualBlock = partial(SAEModel.ResidualBlock, **kwargs)
 
         def func(input):
             x = input[0]
@@ -609,8 +582,8 @@ class SAEModel(ModelBase):
     @staticmethod
     def DFEncFlow(resolution, ae_dims, ch_dims, padding='zero', **kwargs):
         exec (nnlib.import_all(), locals(), globals())
-        upscale = SAEModel.upscale_pre(padding=padding)
+        upscale = partial(SAEModel.upscale, padding=padding)
-        downscale = SAEModel.downscale_pre(padding=padding)
+        downscale = partial(SAEModel.downscale, padding=padding)
         lowest_dense_res = resolution // 16
 
         def func(input):
@@ -634,10 +607,10 @@ class SAEModel(ModelBase):
     @staticmethod
     def DFDecFlow(output_nc, ch_dims, multiscale_count=1, add_residual_blocks=False, padding='zero', **kwargs):
         exec (nnlib.import_all(), locals(), globals())
-        upscale = SAEModel.upscale_pre(padding=padding)
+        upscale = partial(SAEModel.upscale, padding=padding)
-        to_bgr = SAEModel.to_bgr_pre(padding=padding)
+        to_bgr = partial(SAEModel.to_bgr, padding=padding)
         dims = output_nc * ch_dims
-        ResidualBlock = SAEModel.ResidualBlock_pre(padding=padding)
+        ResidualBlock = partial(SAEModel.ResidualBlock, padding=padding)
 
         def func(input):
             x = input[0]

nnlib/device.py

@@ -274,10 +274,14 @@ has_nvml_cap = False
 #- CUDA build of DFL
 has_nvidia_device = os.environ.get("DFL_FORCE_HAS_NVIDIA_DEVICE", "0") == "1"
 
-plaidML_devices = []
+plaidML_devices = None
-
+def get_plaidML_devices():
-# Using plaidML OpenCL backend to determine system devices and has_nvidia_device
+    global plaidML_devices
-try:
+    global has_nvidia_device
+    if plaidML_devices is None:
+        plaidML_devices = []
+        # Using plaidML OpenCL backend to determine system devices and has_nvidia_device
+        try:
             os.environ['PLAIDML_EXPERIMENTAL'] = 'false' #this enables work plaidML without run 'plaidml-setup'
             import plaidml
             ctx = plaidml.Context()
@@ -292,10 +296,12 @@ try:
                                     'description' : d.description.decode()
                                 }]
             ctx.shutdown()
-except:
+        except:
             pass
+    return plaidML_devices
  
-plaidML_devices_count = len(plaidML_devices)
+if not has_nvidia_device:
+    get_plaidML_devices()    
 
 #choosing backend
 
@@ -324,7 +330,7 @@ if device.backend is None and not force_tf_cpu:
 
 if force_plaidML or (device.backend is None and not has_nvidia_device):
     #tensorflow backend was failed without has_nvidia_device , or forcing plaidML, trying to use plaidML backend
-    if plaidML_devices_count == 0:
+    if len(get_plaidML_devices()) == 0:
         #print ("plaidML: No capable OpenCL devices found. Falling back to tensorflow backend.")
         device.backend = None
     else:

nnlib/nnlib.py

@@ -52,7 +52,7 @@ Input = KL.Input
 
 Dense = KL.Dense
 Conv2D = nnlib.Conv2D
-Conv2DTranspose = KL.Conv2DTranspose
+Conv2DTranspose = nnlib.Conv2DTranspose
 SeparableConv2D = KL.SeparableConv2D
 MaxPooling2D = KL.MaxPooling2D
 UpSampling2D = KL.UpSampling2D
@@ -696,6 +696,26 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
                 return self.func(x)
         nnlib.Conv2D = Conv2D
         
+        class Conv2DTranspose():
+            def __init__ (self, *args, **kwargs):
+                self.reflect_pad = False
+                padding = kwargs.get('padding','')
+                if padding == 'zero':
+                    kwargs['padding'] = 'same'
+                if padding == 'reflect':
+                    kernel_size = kwargs['kernel_size']
+                    if (kernel_size % 2) == 1:
+                        self.pad = (kernel_size // 2,)*2
+                        kwargs['padding'] = 'valid'
+                        self.reflect_pad = True
+                self.func = keras.layers.Conv2DTranspose (*args, **kwargs)
+
+            def __call__(self,x):
+                if self.reflect_pad:
+                    x = ReflectionPadding2D( self.pad ) (x)
+                return self.func(x)
+        nnlib.Conv2DTranspose = Conv2DTranspose
+
     @staticmethod
     def import_keras_contrib(device_config):
         if nnlib.keras_contrib is not None:

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
@@ -10,13 +16,12 @@ class SampleType(IntEnum):
     FACE = 1                      #aligned face unsorted
     FACE_YAW_SORTED = 2           #sorted by yaw
     FACE_YAW_SORTED_AS_TARGET = 3 #sorted by yaw and included only yaws which exist in TARGET also automatic mirrored
-    FACE_WITH_CLOSE_TO_SELF = 4
+    FACE_END = 3
-    FACE_END = 4
 
-    QTY = 5
+    QTY = 4
 
 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, source_filename=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
@@ -25,10 +30,12 @@ class Sample(object):
         self.ie_polys = ie_polys
         self.pitch = pitch
         self.yaw = yaw
+        self.source_filename = source_filename
         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, source_filename=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,
@@ -38,8 +45,10 @@ class Sample(object):
             ie_polys=ie_polys if ie_polys is not None else self.ie_polys,
             pitch=pitch if pitch is not None else self.pitch,
             yaw=yaw if yaw is not None else self.yaw,
+            source_filename=source_filename if source_filename is not None else self.source_filename,
             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 +56,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/SampleGeneratorFace.py

@@ -15,7 +15,7 @@ output_sample_types = [
                       ]
 '''
 class SampleGeneratorFace(SampleGeneratorBase):
-    def __init__ (self, samples_path, debug, batch_size, sort_by_yaw=False, sort_by_yaw_target_samples_path=None, with_close_to_self=False, sample_process_options=SampleProcessor.Options(), output_sample_types=[], add_sample_idx=False, add_pitch=False, add_yaw=False, generators_count=2, generators_random_seed=None, **kwargs):
+    def __init__ (self, samples_path, debug, batch_size, sort_by_yaw=False, sort_by_yaw_target_samples_path=None, sample_process_options=SampleProcessor.Options(), output_sample_types=[], add_sample_idx=False, add_pitch=False, add_yaw=False, generators_count=2, generators_random_seed=None, **kwargs):
         super().__init__(samples_path, debug, batch_size)
         self.sample_process_options = sample_process_options
         self.output_sample_types = output_sample_types
@@ -27,8 +27,6 @@ class SampleGeneratorFace(SampleGeneratorBase):
             self.sample_type = SampleType.FACE_YAW_SORTED_AS_TARGET
         elif sort_by_yaw:
             self.sample_type = SampleType.FACE_YAW_SORTED
-        elif with_close_to_self:
-            self.sample_type = SampleType.FACE_WITH_CLOSE_TO_SELF
         else:
             self.sample_type = SampleType.FACE
 
@@ -82,7 +80,7 @@ class SampleGeneratorFace(SampleGeneratorBase):
             if all ( [ samples[idx] == None for idx in samples_idxs] ):
                 raise ValueError('Not enough training data. Gather more faces!')
 
-        if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
+        if self.sample_type == SampleType.FACE:
             shuffle_idxs = []
         elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
             shuffle_idxs = []
@@ -102,12 +100,12 @@ class SampleGeneratorFace(SampleGeneratorBase):
 
                     if len(repeat_samples_idxs) > 0:
                         idx = repeat_samples_idxs.pop()
-                        if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
+                        if self.sample_type == SampleType.FACE:
                             sample = samples[idx]
                         elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
                             sample = samples[(idx >> 16) & 0xFFFF][idx & 0xFFFF]
                     else:
-                        if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
+                        if self.sample_type == SampleType.FACE:
                             if len(shuffle_idxs) == 0:
                                 shuffle_idxs = samples_idxs.copy()
                                 np.random.shuffle(shuffle_idxs)

samplelib/SampleLoader.py

@@ -1,19 +1,19 @@
+import operator
 import traceback
 from enum import IntEnum
-import cv2
-import numpy as np
 from pathlib import Path
 
-from utils import Path_utils
+import cv2
-from utils.DFLPNG import DFLPNG
+import numpy as np
-from utils.DFLJPG import DFLJPG
 
-from .Sample import Sample
+from facelib import FaceType, LandmarksProcessor
-from .Sample import SampleType
-
-from facelib import FaceType
-from facelib import LandmarksProcessor
 from interact import interact as io
+from utils import Path_utils
+from utils.DFLJPG import DFLJPG
+from utils.DFLPNG import DFLPNG
+
+from .Sample import Sample, SampleType
+
 
 class SampleLoader:
     cache = dict()
@@ -35,6 +35,10 @@ class SampleLoader:
             if  datas[sample_type] is None:
                 datas[sample_type] = SampleLoader.upgradeToFaceSamples( [ Sample(filename=filename) for filename in Path_utils.get_image_paths(samples_path) ] )
 
+        # elif          sample_type == SampleType.FACE_TEMPORAL_SORTED:
+        #     if  datas[sample_type] is None:
+        #         datas[sample_type] = SampleLoader.upgradeToFaceTemporalSortedSamples( SampleLoader.load(SampleType.FACE, samples_path) )
+
         elif          sample_type == SampleType.FACE_YAW_SORTED:
             if  datas[sample_type] is None:
                 datas[sample_type] = SampleLoader.upgradeToFaceYawSortedSamples( SampleLoader.load(SampleType.FACE, samples_path) )
@@ -44,10 +48,6 @@ class SampleLoader:
                 if target_samples_path is None:
                     raise Exception('target_samples_path is None for FACE_YAW_SORTED_AS_TARGET')
                 datas[sample_type] = SampleLoader.upgradeToFaceYawSortedAsTargetSamples( SampleLoader.load(SampleType.FACE_YAW_SORTED, samples_path), SampleLoader.load(SampleType.FACE_YAW_SORTED, target_samples_path) )
-        elif          sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
-            if  datas[sample_type] is None:
-                datas[sample_type] = SampleLoader.upgradeToFaceCloseToSelfSamples( SampleLoader.load(SampleType.FACE, samples_path) )
-
 
         return datas[sample_type]
 
@@ -77,44 +77,20 @@ class SampleLoader:
                                                    landmarks=dflimg.get_landmarks(),
                                                    ie_polys=dflimg.get_ie_polys(),
                                                    pitch=pitch,
-                                                   yaw=yaw) )
+                                                   yaw=yaw,
+                                                   source_filename=dflimg.get_source_filename(),
+                                                   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() ) )
 
         return sample_list
 
-    @staticmethod
+    # @staticmethod
-    def upgradeToFaceCloseToSelfSamples (samples):
+    # def upgradeToFaceTemporalSortedSamples( samples ):
-        yaw_samples = SampleLoader.upgradeToFaceYawSortedSamples(samples)
+    #     new_s = [ (s, s.source_filename) for s in samples]
-        yaw_samples_len = len(yaw_samples)
+    #     new_s = sorted(new_s, key=operator.itemgetter(1))
 
-        sample_list = []
+    #     return [ s[0] for s in new_s]
-        for i in io.progress_bar_generator( range(yaw_samples_len), "Sorting"):
-            if yaw_samples[i] is not None:
-                for s in yaw_samples[i]:
-                    s_t = []
-
-                    for n in range(2000):
-                        yaw_idx = np.clip ( i-10 +np.random.randint(20), 0, yaw_samples_len-1 )
-                        if yaw_samples[yaw_idx] is None:
-                            continue
-
-                        yaw_idx_samples_len = len(yaw_samples[yaw_idx])
-
-                        yaw_idx_sample = yaw_samples[yaw_idx][ np.random.randint(yaw_idx_samples_len) ]
-                        if s.filename == yaw_idx_sample.filename:
-                            continue
-
-                        s_t.append ( yaw_idx_sample )
-                        if len(s_t) >= 50:
-                            break
-
-                    if len(s_t) == 0:
-                        s_t = [s]
-
-                    sample_list.append( s.copy_and_set(close_target_list = s_t) )
-
-        return sample_list
 
     @staticmethod
     def upgradeToFaceYawSortedSamples( samples ):

samplelib/SampleProcessor.py

@@ -172,7 +172,14 @@ class SampleProcessor(object):
                                 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)
@@ -202,7 +209,7 @@ class SampleProcessor(object):
                 img_mask = img[...,3:4]
 
                 if f & SPTF.MODE_BGR != 0:
-                    img = img
+                    img = img_bgr
                 elif f & SPTF.MODE_BGR_SHUFFLE != 0:
                     img_bgr = np.take (img_bgr, np.random.permutation(img_bgr.shape[-1]), axis=-1)
                     img = np.concatenate ( (img_bgr,img_mask) , -1 )

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,9 +141,16 @@ 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):
+                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,13 +215,17 @@ 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""
@@ -257,3 +286,8 @@ class DFLJPG(object):
         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

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:
@@ -298,7 +321,8 @@ 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 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__)