merge with mask bounds fix

.vscode/launch.json

@@ -12,7 +12,7 @@
             "type": "python",
             "request": "launch",
             "program": "${env:DFL_ROOT}\\main.py",
-            "pythonPath": "${env:PYTHONEXECUTABLE}",
+            "python": "${env:PYTHONEXECUTABLE}",
             "cwd": "${env:WORKSPACE}",
             "console": "integratedTerminal",    
             "args": ["train", 

DFLIMG/DFLJPG.py

@@ -281,6 +281,13 @@ class DFLJPG(object):
     def has_xseg_mask(self):
         return self.dfl_dict.get('xseg_mask',None) is not None
 
+    def get_xseg_mask_compressed(self):
+        mask_buf = self.dfl_dict.get('xseg_mask',None)
+        if mask_buf is None:
+            return None
+
+        return mask_buf
+        
     def get_xseg_mask(self):
         mask_buf = self.dfl_dict.get('xseg_mask',None)
         if mask_buf is None:
@@ -301,7 +308,7 @@ class DFLJPG(object):
         mask_a = imagelib.normalize_channels(mask_a, 1)
         img_data = np.clip( mask_a*255, 0, 255 ).astype(np.uint8)
 
-        data_max_len = 4096
+        data_max_len = 8192
 
         ret, buf = cv2.imencode('.png', img_data)
 

README.md

@@ -29,15 +29,18 @@ More than 95% of deepfake videos are created with DeepFaceLab.
 
 DeepFaceLab is used by such popular youtube channels as
 
-|![](doc/youtube_icon.png) [Ctrl Shift Face](https://www.youtube.com/channel/UCKpH0CKltc73e4wh0_pgL3g)|![](doc/youtube_icon.png) [VFXChris Ume](https://www.youtube.com/channel/UCGf4OlX_aTt8DlrgiH3jN3g/videos)|
+
+
+|![](doc/youtube_icon.png) [Ctrl Shift Face](https://www.youtube.com/channel/UCKpH0CKltc73e4wh0_pgL3g)|![](doc/youtube_icon.png) [VFXChris Ume](https://www.youtube.com/channel/UCGf4OlX_aTt8DlrgiH3jN3g/videos)|![](doc/youtube_icon.png) [Sham00k](https://www.youtube.com/channel/UCZXbWcv7fSZFTAZV4beckyw/videos)|
+|---|---|---|
+
+
+|![](doc/youtube_icon.png) [Collider videos](https://www.youtube.com/watch?v=A91P2qtPT54&list=PLayt6616lBclvOprvrC8qKGCO-mAhPRux)|![](doc/youtube_icon.png) [iFake](https://www.youtube.com/channel/UCC0lK2Zo2BMXX-k1Ks0r7dg/videos)|![](doc/youtube_icon.png) [NextFace](https://www.youtube.com/channel/UCFh3gL0a8BS21g-DHvXZEeQ/videos)|
+|---|---|---|
+
+|![](doc/youtube_icon.png) [Futuring Machine](https://www.youtube.com/channel/UCC5BbFxqLQgfnWPhprmQLVg)|![](doc/youtube_icon.png) [RepresentUS](https://www.youtube.com/channel/UCRzgK52MmetD9aG8pDOID3g)|
 |---|---|
 
-|![](doc/youtube_icon.png) [Sham00k](https://www.youtube.com/channel/UCZXbWcv7fSZFTAZV4beckyw/videos)|![](doc/youtube_icon.png) [Collider videos](https://www.youtube.com/watch?v=A91P2qtPT54&list=PLayt6616lBclvOprvrC8qKGCO-mAhPRux)|![](doc/youtube_icon.png) [iFake](https://www.youtube.com/channel/UCC0lK2Zo2BMXX-k1Ks0r7dg/videos)|
-|---|---|---|
-
-|![](doc/youtube_icon.png) [NextFace](https://www.youtube.com/channel/UCFh3gL0a8BS21g-DHvXZEeQ/videos)|![](doc/youtube_icon.png) [Futuring Machine](https://www.youtube.com/channel/UCC5BbFxqLQgfnWPhprmQLVg)|![](doc/youtube_icon.png) [RepresentUS](https://www.youtube.com/channel/UCRzgK52MmetD9aG8pDOID3g)|
-|---|---|---|
-
 </td></tr>
 
 <tr><td colspan=2 align="center">
@@ -49,7 +52,7 @@ DeepFaceLab is used by such popular youtube channels as
 
 ## Replace the face
 
-<img src="doc/replace_the_face.png" align="center">
+<img src="doc/replace_the_face.jpg" align="center">
 
 </td></tr>
 
@@ -189,7 +192,7 @@ Unfortunately, there is no "make everything ok" button in DeepFaceLab. You shoul
 </td></tr>
 
 <tr><td align="right">
-<a href="https://tinyurl.com/y6npm2su">Windows (magnet link)</a>
+<a href="https://tinyurl.com/yb6gw8hu">Windows (magnet link)</a>
 </td><td align="center">Last release. Use torrent client to download.</td></tr>
 
 <tr><td align="right">
@@ -284,10 +287,6 @@ Unfortunately, there is no "make everything ok" button in DeepFaceLab. You shoul
 <a href="https://mrdeepfakes.com/forums/">mrdeepfakes</a>
 </td><td align="center">the biggest NSFW English community</td></tr>
 
-<tr><td align="right">
-<a href="https://www.reddit.com/r/GifFakes/new/">reddit r/GifFakes/</a>
-</td><td align="center">Post your deepfakes there !</td></tr>
-
 <tr><td align="right">
 <a href="https://www.reddit.com/r/SFWdeepfakes/new/">reddit r/SFWdeepfakes/</a>
 </td><td align="center">Post your deepfakes there !</td></tr>
@@ -296,6 +295,10 @@ Unfortunately, there is no "make everything ok" button in DeepFaceLab. You shoul
 QQ 951138799
 </td><td align="center">中文 Chinese QQ group for ML/AI experts</td></tr>
 
+<tr><td align="right">
+<a href="https://www.dfldata.xyz">dfldata.xyz</a>
+</td><td align="center">中文交流论坛，免费软件教程、模型、人脸数据</td></tr>
+
 <tr><td align="right">
 <a href="https://www.deepfaker.xyz/">deepfaker.xyz</a>
 </td><td align="center">中文学习站（非官方)</td></tr>
@@ -317,7 +320,7 @@ QQ 951138799
 </td></tr>
 
 <tr><td colspan=2 align="center">
-<a href="https://www.paypal.com/cgi-bin/webscr?cmd=_donations&business=lepersorium@gmail.com&lc=US&no_note=0&item_name=Support+DeepFaceLab&cn=&curency_code=USD&bn=PP-DonationsBF:btn_donateCC_LG.gif:NonHosted">Donate via Paypal</a>
+<a href="https://www.paypal.com/paypalme/DeepFaceLab">Donate via Paypal</a>
 </td></tr>
 
 <tr><td colspan=2 align="center">

core/imagelib/warp.py

@@ -41,6 +41,9 @@ def gen_warp_params (w, flip, rotation_range=[-10,10], scale_range=[-0.5, 0.5],
     params['mapx'] = mapx
     params['mapy'] = mapy
     params['rmat'] = random_transform_mat
+    u_mat = random_transform_mat.copy()
+    u_mat[:,2] /= w
+    params['umat'] = u_mat
     params['w'] = w
     params['rw'] = rw
     params['flip'] = p_flip

core/leras/archis/DeepFakeArchi.py

@@ -72,12 +72,23 @@ class DeepFakeArchi(nn.ArchiBase):
                     return x
 
             class Encoder(nn.ModelBase):
-                def on_build(self, in_ch, e_ch):
+                def __init__(self, in_ch, e_ch, **kwargs ):
-                    self.down1 = DownscaleBlock(in_ch, e_ch, n_downscales=4, kernel_size=5)
+                    self.in_ch = in_ch
+                    self.e_ch = e_ch
+                    super().__init__(**kwargs)
+                    
+                def on_build(self):                    
+                    self.down1 = DownscaleBlock(self.in_ch, self.e_ch, n_downscales=4, kernel_size=5)
 
                 def forward(self, inp):
                     return nn.flatten(self.down1(inp))
                     
+                def get_out_res(self, res):
+                    return res // (2**4)
+                    
+                def get_out_ch(self):
+                    return self.e_ch * 8
+
             lowest_dense_res = resolution // (32 if 'd' in opts else 16)
 
             class Inter(nn.ModelBase):
@@ -104,9 +115,8 @@ class DeepFakeArchi(nn.ArchiBase):
                     x = self.upscale1(x)
                     return x
 
-                @staticmethod
+                def get_out_res(self):
-                def get_code_res():
+                    return lowest_dense_res * 2
-                    return lowest_dense_res
 
                 def get_out_ch(self):
                     return self.ae_out_ch

core/leras/layers/Saveable.py

@@ -76,6 +76,7 @@ class Saveable():
         if self.name is None:
             raise Exception("name must be defined.")
 
+        try:
             tuples = []
             for w in weights:
                 w_name_split = w.name.split('/')
@@ -94,6 +95,8 @@ class Saveable():
                     tuples.append ( (w, w_val) )
 
             nn.batch_set_value(tuples)
+        except:
+            return False
             
         return True
 

core/leras/models/ModelBase.py

@@ -116,41 +116,32 @@ class ModelBase(nn.Saveable):
 
         return self.forward(*args, **kwargs)
 
-    def compute_output_shape(self, shapes):
+    # def compute_output_shape(self, shapes):
-        if not self.built:
+    #     if not self.built:
-            self.build()
+    #         self.build()
 
-        not_list = False
+    #     not_list = False
-        if not isinstance(shapes, list):
+    #     if not isinstance(shapes, list):
-            not_list = True
+    #         not_list = True
-            shapes = [shapes]
+    #         shapes = [shapes]
 
-        with tf.device('/CPU:0'):
+    #     with tf.device('/CPU:0'):
-            # CPU tensors will not impact any performance, only slightly RAM "leakage"
+    #         # CPU tensors will not impact any performance, only slightly RAM "leakage"
-            phs = []
+    #         phs = []
-            for dtype,sh in shapes:
+    #         for dtype,sh in shapes:
-                phs += [ tf.placeholder(dtype, sh) ]
+    #             phs += [ tf.placeholder(dtype, sh) ]
 
-            result = self.__call__(phs[0] if not_list else phs)
+    #         result = self.__call__(phs[0] if not_list else phs)
 
-            if not isinstance(result, list):
+    #         if not isinstance(result, list):
-                result = [result]
+    #             result = [result]
 
-            result_shapes = []
+    #         result_shapes = []
 
-            for t in result:
+    #         for t in result:
-                result_shapes += [ t.shape.as_list() ]
+    #             result_shapes += [ t.shape.as_list() ]
 
-            return result_shapes[0] if not_list else result_shapes
+    #         return result_shapes[0] if not_list else result_shapes
-
-    def compute_output_channels(self, shapes):
-        shape = self.compute_output_shape(shapes)
-        shape_len = len(shape)
-
-        if shape_len == 4:
-            if nn.data_format == "NCHW":
-                return shape[1]
-        return shape[-1]
 
     def build_for_run(self, shapes_list):
         if not isinstance(shapes_list, list):

core/leras/nn.py

@@ -70,19 +70,23 @@ class nn():
                         first_run = True
                     os.environ['CUDA_CACHE_PATH'] = str(compute_cache_path)
 
-            os.environ['CUDA_​CACHE_​MAXSIZE'] = '536870912' #512Mb (32mb default)
+            #nvcuda.dll ignores this param : os.environ['CUDA_​CACHE_​MAXSIZE'] = '536870912' #512Mb (32mb default)
             os.environ['TF_MIN_GPU_MULTIPROCESSOR_COUNT'] = '2'
-            os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # tf log errors only
+            os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3' # tf log errors only
 
             if first_run:
                 io.log_info("Caching GPU kernels...")
 
-            import tensorflow as tf
+            #import tensorflow as tf
-            nn.tf = tf
+            import tensorflow.compat.v1 as tf
             
             import logging
             # Disable tensorflow warnings
-            logging.getLogger('tensorflow').setLevel(logging.ERROR)
+            tf_logger = logging.getLogger('tensorflow')
+            tf_logger.setLevel(logging.ERROR)
+            
+            tf.disable_v2_behavior()
+            nn.tf = tf
 
             # Initialize framework
             import core.leras.ops

core/leras/ops/__init__.py

@@ -334,6 +334,16 @@ def depth_to_space(x, size):
         return x
     else:
         return tf.depth_to_space(x, size, data_format=nn.data_format)
+        b,c,h,w = x.shape.as_list()
+        oh, ow = h * size, w * size
+        oc = c // (size * size)
+
+        x = tf.reshape(x, (-1, size, size, oc, h, w, ) )
+        x = tf.transpose(x, (0, 3, 4, 1, 5, 2))
+        x = tf.reshape(x, (-1, oc, oh, ow))
+        return x
+        
+        
 nn.depth_to_space = depth_to_space
 
 def rgb_to_lab(srgb):

core/leras/optimizers/AdaBelief.py

@@ -0,0 +1,80 @@
+from tensorflow.python.ops import control_flow_ops, state_ops
+from core.leras import nn
+tf = nn.tf
+
+class AdaBelief(nn.OptimizerBase):
+    def __init__(self, lr=0.001, beta_1=0.9, beta_2=0.999, lr_dropout=1.0, lr_cos=0, epsilon=1e-7, clipnorm=0.0, name=None, **kwargs):
+        super().__init__(name=name)
+
+        if name is None:
+            raise ValueError('name must be defined.')
+
+        self.lr = lr
+        self.beta_1 = beta_1
+        self.beta_2 = beta_2
+        self.lr_dropout = lr_dropout
+        self.lr_cos = lr_cos
+        self.clipnorm = clipnorm
+        self.epsilon = epsilon
+
+        with tf.device('/CPU:0') :
+            with tf.variable_scope(self.name):
+                self.iterations = tf.Variable(0, dtype=tf.int64, name='iters')
+
+        self.ms_dict = {}
+        self.vs_dict = {}
+        self.lr_rnds_dict = {}
+
+    def get_weights(self):
+        return [self.iterations] + list(self.ms_dict.values()) + list(self.vs_dict.values())
+
+    def initialize_variables(self, trainable_weights, vars_on_cpu=True, lr_dropout_on_cpu=False):
+        # Initialize here all trainable variables used in training
+        e = tf.device('/CPU:0') if vars_on_cpu else None
+        if e: e.__enter__()
+        with tf.variable_scope(self.name):
+            ms = { v.name : tf.get_variable ( f'ms_{v.name}'.replace(':','_'), v.shape, dtype=v.dtype, initializer=tf.initializers.constant(0.0), trainable=False) for v in trainable_weights }
+            vs = { v.name : tf.get_variable ( f'vs_{v.name}'.replace(':','_'), v.shape, dtype=v.dtype, initializer=tf.initializers.constant(0.0), trainable=False) for v in trainable_weights }
+            self.ms_dict.update (ms)
+            self.vs_dict.update (vs)
+            
+            if self.lr_dropout != 1.0:
+                e = tf.device('/CPU:0') if lr_dropout_on_cpu else None
+                if e: e.__enter__()                    
+                lr_rnds = [ nn.random_binomial( v.shape, p=self.lr_dropout, dtype=v.dtype) for v in trainable_weights ]
+                if e: e.__exit__(None, None, None)                
+                self.lr_rnds_dict.update ( { v.name : rnd for v,rnd in zip(trainable_weights,lr_rnds) } )
+        if e: e.__exit__(None, None, None)
+
+    def get_update_op(self, grads_vars):
+        updates = []
+
+        if self.clipnorm > 0.0:
+            norm = tf.sqrt( sum([tf.reduce_sum(tf.square(g)) for g,v in grads_vars]))
+        updates += [ state_ops.assign_add( self.iterations, 1) ]
+        for i, (g,v) in enumerate(grads_vars):
+            if self.clipnorm > 0.0:
+                g = self.tf_clip_norm(g, self.clipnorm, norm)
+
+            ms = self.ms_dict[ v.name ]
+            vs = self.vs_dict[ v.name ]
+            
+            m_t = self.beta_1*ms + (1.0-self.beta_1) * g
+            v_t = self.beta_2*vs + (1.0-self.beta_2) * tf.square(g-m_t)
+
+            lr = tf.constant(self.lr, g.dtype)
+            if self.lr_cos != 0:
+                lr *= (tf.cos(  tf.cast(self.iterations, g.dtype) * (2*3.1415926535/ float(self.lr_cos) )  ) + 1.0) / 2.0
+
+            v_diff = - lr * m_t / (tf.sqrt(v_t) + self.epsilon)
+            if self.lr_dropout != 1.0:
+                lr_rnd = self.lr_rnds_dict[v.name]
+                v_diff *= lr_rnd
+            new_v = v + v_diff
+
+            updates.append (state_ops.assign(ms, m_t))
+            updates.append (state_ops.assign(vs, v_t))
+            updates.append (state_ops.assign(v, new_v))
+
+        return control_flow_ops.group ( *updates, name=self.name+'_updates')
+nn.AdaBelief = AdaBelief

core/leras/optimizers/RMSprop.py

@@ -3,27 +3,29 @@ from core.leras import nn
 tf = nn.tf
 
 class RMSprop(nn.OptimizerBase):
-    def __init__(self, lr=0.001, rho=0.9, lr_dropout=1.0, epsilon=1e-7, clipnorm=0.0, name=None):
+    def __init__(self, lr=0.001, rho=0.9, lr_dropout=1.0, epsilon=1e-7, clipnorm=0.0, name=None, **kwargs):
         super().__init__(name=name)
 
         if name is None:
             raise ValueError('name must be defined.')
 
         self.lr_dropout = lr_dropout
+        self.lr = lr
+        self.rho = rho
+        self.epsilon = epsilon
+        
         self.clipnorm = clipnorm
 
         with tf.device('/CPU:0') :
             with tf.variable_scope(self.name):
-                self.lr = tf.Variable (lr, name="lr")
+                
-                self.rho = tf.Variable (rho, name="rho")
-                self.epsilon = tf.Variable (epsilon, name="epsilon")
                 self.iterations = tf.Variable(0, dtype=tf.int64, name='iters')
 
         self.accumulators_dict = {}
         self.lr_rnds_dict = {}
 
     def get_weights(self):
-        return [self.lr, self.rho, self.epsilon, self.iterations] + list(self.accumulators_dict.values())
+        return [self.iterations] + list(self.accumulators_dict.values())
 
     def initialize_variables(self, trainable_weights, vars_on_cpu=True, lr_dropout_on_cpu=False):
         # Initialize here all trainable variables used in training
@@ -53,13 +55,11 @@ class RMSprop(nn.OptimizerBase):
 
             a = self.accumulators_dict[ v.name ]
 
-            rho = tf.cast(self.rho, a.dtype)
+            new_a = self.rho * a + (1. - self.rho) * tf.square(g)
-            new_a = rho * a + (1. - rho) * tf.square(g)
 
-            lr = tf.cast(self.lr, a.dtype)
+            lr = tf.constant(self.lr, g.dtype)
-            epsilon = tf.cast(self.epsilon, a.dtype)
 
-            v_diff = - lr * g / (tf.sqrt(new_a) + epsilon)
+            v_diff = - lr * g / (tf.sqrt(new_a) + self.epsilon)
             if self.lr_dropout != 1.0:
                 lr_rnd = self.lr_rnds_dict[v.name]
                 v_diff *= lr_rnd

core/leras/optimizers/__init__.py

@@ -1,2 +1,3 @@
 from .OptimizerBase import *
 from .RMSprop import *
+from .AdaBelief import *

core/mplib/MPSharedList.py

@@ -60,9 +60,11 @@ class MPSharedList():
                 break
             key -= self.obj_counts[i]
 
-        offset_start, offset_end = struct.unpack('<QQ', bytes(sh_b[ table_offset + key*8     : table_offset + (key+2)*8]) )
+        sh_b = memoryview(sh_b).cast('B')
 
-        return pickle.loads( bytes(sh_b[ data_offset + offset_start : data_offset + offset_end ]) )
+        offset_start, offset_end = struct.unpack('<QQ', sh_b[ table_offset + key*8     : table_offset + (key+2)*8].tobytes() )
+
+        return pickle.loads( sh_b[ data_offset + offset_start : data_offset + offset_end ].tobytes() )
 
     def __iter__(self):
         for i in range(self.__len__()):
@@ -84,7 +86,8 @@ class MPSharedList():
             data_size    = sum([len(x) for x in obj_pickled_ar])
 
             sh_b = multiprocessing.RawArray('B', table_size + data_size)
-            sh_b[0:8] = struct.pack('<Q', obj_count)  
+            #sh_b[0:8] = struct.pack('<Q', obj_count)
+            sh_b_view = memoryview(sh_b).cast('B')
 
             offset = 0
 
@@ -97,51 +100,12 @@ class MPSharedList():
                 offset += len(obj_pickled_ar[i])
             offsets.append(offset)
 
-            sh_b[table_offset:table_offset+table_size] = struct.pack( '<'+'Q'*len(offsets), *offsets )
+            sh_b_view[table_offset:table_offset+table_size] = struct.pack( '<'+'Q'*len(offsets), *offsets )
 
-            ArrayFillerSubprocessor(sh_b, [ (data_offset+offsets[i], obj_pickled_ar[i] ) for i in range(obj_count) ] ).run()
+            for i, obj_pickled in enumerate(obj_pickled_ar):
+                offset = data_offset+offsets[i]
+                sh_b_view[offset:offset+len(obj_pickled)] = obj_pickled_ar[i]
 
             return obj_count, table_offset, data_offset, sh_b
         return 0, 0, 0, None
 
-        
-
-class ArrayFillerSubprocessor(Subprocessor):
-    """
-    Much faster to fill shared memory via subprocesses rather than direct whole bytes fill.
-    """
-    #override
-    def __init__(self, sh_b, data_list ):
-        self.sh_b = sh_b        
-        self.data_list = data_list
-        super().__init__('ArrayFillerSubprocessor', ArrayFillerSubprocessor.Cli, 60)
-
-    #override
-    def process_info_generator(self):
-        for i in range(min(multiprocessing.cpu_count(), 8)):
-            yield 'CPU%d' % (i), {}, {'sh_b':self.sh_b}
-
-    #override
-    def get_data(self, host_dict):
-        if len(self.data_list) > 0:
-            return self.data_list.pop(0)
-
-        return None
-
-    #override
-    def on_data_return (self, host_dict, data):
-        self.data_list.insert(0, data)
-
-    #override
-    def on_result (self, host_dict, data, result):
-        pass
-
-    class Cli(Subprocessor.Cli):
-        #overridable optional
-        def on_initialize(self, client_dict):
-            self.sh_b = client_dict['sh_b']
-            
-        def process_data(self, data):
-            offset, b = data
-            self.sh_b[offset:offset+len(b)]=b
-            return 0

facelib/XSegNet.py

@@ -30,6 +30,9 @@ class XSegNet(object):
         nn.initialize(data_format=data_format)
         tf = nn.tf
         
+        model_name = f'{name}_{resolution}'
+        self.model_filename_list = []
+        
         with tf.device ('/CPU:0'):
             #Place holders on CPU
             self.input_t  = tf.placeholder (nn.floatx, nn.get4Dshape(resolution,resolution,3) )
@@ -39,18 +42,17 @@ class XSegNet(object):
         with tf.device ('/CPU:0' if place_model_on_cpu else '/GPU:0'):
             self.model = nn.XSeg(3, 32, 1, name=name)
             self.model_weights = self.model.get_weights()
-
-        model_name = f'{name}_{resolution}'
-        self.model_filename_list = [ [self.model, f'{model_name}.npy'] ]
-
             if training:
                 if optimizer is None:
                     raise ValueError("Optimizer should be provided for training mode.")                
-
                 self.opt = optimizer              
                 self.opt.initialize_variables (self.model_weights, vars_on_cpu=place_model_on_cpu)                    
                 self.model_filename_list += [ [self.opt, f'{model_name}_opt.npy' ] ]
-        else:
+                
+        
+        self.model_filename_list += [ [self.model, f'{model_name}.npy'] ]
+
+        if not training:
             with tf.device ('/CPU:0' if run_on_cpu else '/GPU:0'):
                 _, pred = self.model(self.input_t)
 

main.py

@@ -70,7 +70,7 @@ if __name__ == "__main__":
 
     p = subparsers.add_parser( "sort", help="Sort faces in a directory.")
     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('--by', dest="sort_by_method", default=None, choices=("blur", "face-yaw", "face-pitch", "face-source-rect-size", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "oneface", "final", "final-faster", "absdiff"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." )
+    p.add_argument('--by', dest="sort_by_method", default=None, choices=("blur", "motion-blur", "face-yaw", "face-pitch", "face-source-rect-size", "hist", "hist-dissim", "brightness", "hue", "black", "origname", "oneface", "final-by-blur", "final-by-size", "absdiff"), help="Method of sorting. 'origname' sort by original filename to recover original sequence." )
     p.set_defaults (func=process_sort)
 
     def process_util(arguments):

mainscripts/Sorter.py

@@ -23,6 +23,9 @@ from facelib import LandmarksProcessor
 
 class BlurEstimatorSubprocessor(Subprocessor):
     class Cli(Subprocessor.Cli):
+        def on_initialize(self, client_dict):
+            self.estimate_motion_blur = client_dict['estimate_motion_blur']
+        
         #override
         def process_data(self, data):
             filepath = Path( data[0] )
@@ -37,7 +40,13 @@ class BlurEstimatorSubprocessor(Subprocessor):
                 face_mask = LandmarksProcessor.get_image_hull_mask (image.shape, dflimg.get_landmarks())
                 image = (image*face_mask).astype(np.uint8)
                 
-                return [ str(filepath), estimate_sharpness(image) ]
+                
+                if self.estimate_motion_blur:
+                    value = cv2.Laplacian(image, cv2.CV_64F, ksize=11).var()    
+                else:
+                    value = estimate_sharpness(image)
+                
+                return [ str(filepath), value ]
 
 
         #override
@@ -46,8 +55,9 @@ class BlurEstimatorSubprocessor(Subprocessor):
             return data[0]
 
     #override
-    def __init__(self, input_data ):
+    def __init__(self, input_data, estimate_motion_blur=False ):
         self.input_data = input_data
+        self.estimate_motion_blur = estimate_motion_blur
         self.img_list = []
         self.trash_img_list = []
         super().__init__('BlurEstimator', BlurEstimatorSubprocessor.Cli, 60)
@@ -66,7 +76,7 @@ class BlurEstimatorSubprocessor(Subprocessor):
         io.log_info(f'Running on {cpu_count} CPUs')
 
         for i in range(cpu_count):
-            yield 'CPU%d' % (i), {}, {}
+            yield 'CPU%d' % (i), {}, {'estimate_motion_blur':self.estimate_motion_blur}
 
     #override
     def get_data(self, host_dict):
@@ -104,6 +114,17 @@ def sort_by_blur(input_path):
 
     return img_list, trash_img_list
     
+def sort_by_motion_blur(input_path):
+    io.log_info ("Sorting by motion blur...")
+
+    img_list = [ (filename,[]) for filename in pathex.get_image_paths(input_path) ]
+    img_list, trash_img_list = BlurEstimatorSubprocessor (img_list, estimate_motion_blur=True).run()
+
+    io.log_info ("Sorting...")
+    img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
+
+    return img_list, trash_img_list
+    
 def sort_by_face_yaw(input_path):
     io.log_info ("Sorting by face yaw...")
     img_list = []
@@ -876,6 +897,7 @@ def final_process(input_path, img_list, trash_img_list):
 
 sort_func_methods = {
     'blur':        ("blur", sort_by_blur),
+    'motion-blur': ("motion_blur", sort_by_motion_blur),
     'face-yaw':    ("face yaw direction", sort_by_face_yaw),
     'face-pitch':  ("face pitch direction", sort_by_face_pitch),
     'face-source-rect-size' : ("face rect size in source image", sort_by_face_source_rect_size),
@@ -903,7 +925,7 @@ def main (input_path, sort_by_method=None):
             io.log_info(f"[{i}] {desc}")
 
         io.log_info("")
-        id = io.input_int("", 4, valid_list=[*range(len(key_list))] )
+        id = io.input_int("", 5, valid_list=[*range(len(key_list))] )
 
         sort_by_method = key_list[id]
     else:

mainscripts/Trainer.py

@@ -1,4 +1,5 @@
-import sys
+import os
+import sys
 import traceback
 import queue
 import threading
@@ -120,6 +121,12 @@ def trainerThread (s2c, c2s, e,
                             io.log_info("Trying to do the first iteration. If an error occurs, reduce the model parameters.")
                             io.log_info("")
                             
+                            if sys.platform[0:3] == 'win':
+                                io.log_info("!!!")
+                                io.log_info("Windows 10 users IMPORTANT notice. You should set this setting in order to work correctly.")
+                                io.log_info("https://i.imgur.com/B7cmDCB.jpg")
+                                io.log_info("!!!")
+
                         iter, iter_time = model.train_one_iter()
 
                         loss_history = model.get_loss_history()

mainscripts/dev_misc.py

@@ -357,25 +357,114 @@ def extract_umd_csv(input_file_csv,
     io.log_info ('Faces detected:      %d' % (faces_detected) )
     io.log_info ('-------------------------')
 
-def dev_test1(input_dir):
+
+    
+def dev_test(input_dir):
+    # LaPa dataset
+    
+    image_size = 1024
+    face_type = FaceType.HEAD
+    
     input_path = Path(input_dir)
+    images_path = input_path / 'images'    
+    if not images_path.exists:
+        raise ValueError('LaPa dataset: images folder not found.')
+    labels_path = input_path / 'labels'    
+    if not labels_path.exists:
+        raise ValueError('LaPa dataset: labels folder not found.')
+    landmarks_path = input_path / 'landmarks'    
+    if not landmarks_path.exists:
+        raise ValueError('LaPa dataset: landmarks folder not found.')
     
-    dir_names = pathex.get_all_dir_names(input_path)
+    output_path = input_path / 'out'    
+    if output_path.exists():
+        output_images_paths = pathex.get_image_paths(output_path)
+        if len(output_images_paths) != 0:
+            io.input(f"\n WARNING !!! \n {output_path} contains files! \n They will be deleted. \n Press enter to continue.\n")
+            for filename in output_images_paths:
+                Path(filename).unlink()
+    output_path.mkdir(parents=True, exist_ok=True)
     
-    for dir_name in io.progress_bar_generator(dir_names, desc="Processing"):
+    data = []
     
-        img_paths = pathex.get_image_paths (input_path / dir_name)
+    img_paths = pathex.get_image_paths (images_path)
     for filename in img_paths:
         filepath = Path(filename)
 
-            dflimg = DFLIMG.x (filepath)
+        landmark_filepath = landmarks_path / (filepath.stem + '.txt')
-            if dflimg is None:
+        if not landmark_filepath.exists():
-                raise ValueError
+            raise ValueError(f'no landmarks for {filepath}')
         
-            #dflimg.x(filename, person_name=dir_name)
+        #img = cv2_imread(filepath)
+        
+        lm = landmark_filepath.read_text()
+        lm = lm.split('\n')
+        if int(lm[0]) != 106:
+            raise ValueError(f'wrong landmarks format in {landmark_filepath}')
+        
+        lmrks = []
+        for i in range(106):
+            x,y = lm[i+1].split(' ')
+            x,y = float(x), float(y)
+            lmrks.append ( (x,y) )
+            
+        lmrks = np.array(lmrks)
+        
+        l,t = np.min(lmrks, 0)
+        r,b = np.max(lmrks, 0)
+        
+        l,t,r,b = ( int(x) for x in (l,t,r,b) )
+        
+        #for x, y in lmrks:
+        #    x,y = int(x), int(y)
+        #    cv2.circle(img, (x, y), 1, (0,255,0) , 1, lineType=cv2.LINE_AA)   
+         
+        #imagelib.draw_rect(img, (l,t,r,b), (0,255,0) )
+         
+        
+        data += [ ExtractSubprocessor.Data(filepath=filepath, rects=[ (l,t,r,b) ]) ]
+
+        #cv2.imshow("", img) 
+        #cv2.waitKey(0)
+ 
+    if len(data) > 0:
+        device_config = nn.DeviceConfig.BestGPU()
+        
+        io.log_info ("Performing 2nd pass...")
+        data = ExtractSubprocessor (data, 'landmarks', image_size, 95, face_type,  device_config=device_config).run()
+        io.log_info ("Performing 3rd pass...")
+        data = ExtractSubprocessor (data, 'final', image_size, 95, face_type, final_output_path=output_path, device_config=device_config).run()
+
+
+        for filename in pathex.get_image_paths (output_path):
+            filepath = Path(filename)
+            
+            
+            dflimg = DFLJPG.load(filepath)
+            
+            src_filename = dflimg.get_source_filename()
+            image_to_face_mat = dflimg.get_image_to_face_mat()
+
+            label_filepath = labels_path / ( Path(src_filename).stem + '.png')        
+            if not label_filepath.exists():
+                raise ValueError(f'{label_filepath} does not exist')
+            
+            mask = cv2_imread(label_filepath)        
+            #mask[mask == 10] = 0 # remove hair
+            mask[mask > 0] = 1
+            mask = cv2.warpAffine(mask, image_to_face_mat, (image_size, image_size), cv2.INTER_LINEAR)
+            mask = cv2.blur(mask, (3,3) )
+            
+            #cv2.imshow("", (mask*255).astype(np.uint8) ) 
+            #cv2.waitKey(0)
+            
+            dflimg.set_xseg_mask(mask)
+            dflimg.save()
+        
+    
+    import code
+    code.interact(local=dict(globals(), **locals()))
                     
-            #import code
-            #code.interact(local=dict(globals(), **locals()))
 
 def dev_resave_pngs(input_dir):
     input_path = Path(input_dir)

models/Model_Quick96/Model.py

@@ -56,10 +56,10 @@ class QModel(ModelBase):
         # Initializing model classes
         with tf.device (models_opt_device):
             self.encoder = model_archi.Encoder(in_ch=input_ch, e_ch=e_dims, name='encoder')
-            encoder_out_ch = self.encoder.compute_output_channels ( (nn.floatx, bgr_shape))
+            encoder_out_ch = self.encoder.get_out_ch()*self.encoder.get_out_res(resolution)**2
 
             self.inter = model_archi.Inter (in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims, name='inter')
-            inter_out_ch = self.inter.compute_output_channels ( (nn.floatx, (None,encoder_out_ch)))
+            inter_out_ch = self.inter.get_out_ch()
 
             self.decoder_src = model_archi.Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, name='decoder_src')
             self.decoder_dst = model_archi.Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, name='decoder_dst')

models/Model_SAEHD/Model.py

@@ -34,7 +34,7 @@ class SAEHDModel(ModelBase):
         default_face_type          = self.options['face_type']          = self.load_or_def_option('face_type', 'f')
         default_models_opt_on_gpu  = self.options['models_opt_on_gpu']  = self.load_or_def_option('models_opt_on_gpu', True)
 
-        archi = self.load_or_def_option('archi', 'df')
+        archi = self.load_or_def_option('archi', 'liae-ud')
         archi = {'dfuhd':'df-u','liaeuhd':'liae-u'}.get(archi, archi) #backward comp
         default_archi              = self.options['archi'] = archi
 
@@ -47,6 +47,8 @@ class SAEHDModel(ModelBase):
         default_mouth_prio         = self.options['mouth_prio']         = self.load_or_def_option('mouth_prio', False)
         default_uniform_yaw        = self.options['uniform_yaw']        = self.load_or_def_option('uniform_yaw', False)
 
+        default_adabelief          = self.options['adabelief']          = self.load_or_def_option('adabelief', True)
+
         lr_dropout = self.load_or_def_option('lr_dropout', 'n')
         lr_dropout = {True:'y', False:'n'}.get(lr_dropout, lr_dropout) #backward comp
         default_lr_dropout         = self.options['lr_dropout'] = lr_dropout
@@ -140,11 +142,13 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
         if self.is_first_run() or ask_override:
             self.options['models_opt_on_gpu'] = io.input_bool ("Place models and optimizer on GPU", default_models_opt_on_gpu, help_message="When you train on one GPU, by default model and optimizer weights are placed on GPU to accelerate the process. You can place they on CPU to free up extra VRAM, thus set bigger dimensions.")
 
+            self.options['adabelief'] = io.input_bool ("Use AdaBelief optimizer?", default_adabelief, help_message="Use AdaBelief optimizer. It requires more VRAM, but the accuracy and the generalization of the model is higher.")
+
             self.options['lr_dropout']  = io.input_str (f"Use learning rate dropout", default_lr_dropout, ['n','y','cpu'], help_message="When the face is trained enough, you can enable this option to get extra sharpness and reduce subpixel shake for less amount of iterations. Enabled it before `disable random warp` and before GAN. \nn - disabled.\ny - enabled\ncpu - enabled on CPU. This allows not to use extra VRAM, sacrificing 20% time of iteration.")
 
             self.options['random_warp'] = io.input_bool ("Enable random warp of samples", default_random_warp, help_message="Random warp is required to generalize facial expressions of both faces. When the face is trained enough, you can disable it to get extra sharpness and reduce subpixel shake for less amount of iterations.")
 
-            self.options['gan_power'] = np.clip ( io.input_number ("GAN power", default_gan_power, add_info="0.0 .. 10.0", help_message="Train the network in Generative Adversarial manner. Forces the neural network to learn small details of the face. Enable it only when the face is trained enough and don't disable. Typical value is 0.1"), 0.0, 10.0 )
+            self.options['gan_power'] = np.clip ( io.input_number ("GAN power", default_gan_power, add_info="0.0 .. 10.0", help_message="Train the network in Generative Adversarial manner. Forces the neural network to learn small details of the face. Enable it only when the face is trained enough and don't disable. Typical fine value is 0.05"), 0.0, 10.0 )
 
             if (self.options['gan_power'] != 0):
                 self.options['gan_old'] = io.input_bool ("Use old GAN version", default_gan_old, help_message="Use older version of GAN." )
@@ -200,6 +204,8 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
         if self.pretrain_just_disabled:
             self.set_iter(0)
 
+        adabelief = self.options['adabelief']
+
         self.gan_power = gan_power = 0.0 if self.pretrain else self.options['gan_power']
         self.gan_old = gan_old = self.options['gan_old']
 
@@ -235,8 +241,10 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
             self.target_src = tf.placeholder (nn.floatx, bgr_shape)
             self.target_dst = tf.placeholder (nn.floatx, bgr_shape)
 
-            self.target_srcm_all = tf.placeholder (nn.floatx, mask_shape)
+            self.target_srcm    = tf.placeholder (nn.floatx, mask_shape)
-            self.target_dstm_all = tf.placeholder (nn.floatx, mask_shape)
+            self.target_srcm_em = tf.placeholder (nn.floatx, mask_shape)
+            self.target_dstm    = tf.placeholder (nn.floatx, mask_shape)
+            self.target_dstm_em = tf.placeholder (nn.floatx, mask_shape)
 
         # Initializing model classes
         model_archi = nn.DeepFakeArchi(resolution, opts=archi_opts)
@@ -244,10 +252,10 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
         with tf.device (models_opt_device):
             if 'df' in archi_type:
                 self.encoder = model_archi.Encoder(in_ch=input_ch, e_ch=e_dims, name='encoder')
-                encoder_out_ch = self.encoder.compute_output_channels ( (nn.floatx, bgr_shape))
+                encoder_out_ch = self.encoder.get_out_ch()*self.encoder.get_out_res(resolution)**2
 
                 self.inter = model_archi.Inter (in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims, name='inter')
-                inter_out_ch = self.inter.compute_output_channels ( (nn.floatx, (None,encoder_out_ch)))
+                inter_out_ch = self.inter.get_out_ch()
 
                 self.decoder_src = model_archi.Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, name='decoder_src')
                 self.decoder_dst = model_archi.Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, name='decoder_dst')
@@ -259,19 +267,18 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
 
                 if self.is_training:
                     if self.options['true_face_power'] != 0:
-                        self.code_discriminator = nn.CodeDiscriminator(ae_dims, code_res=model_archi.Inter.get_code_res()*2, name='dis' )
+                        self.code_discriminator = nn.CodeDiscriminator(ae_dims, code_res=self.inter.get_out_res(), name='dis' )
                         self.model_filename_list += [ [self.code_discriminator, 'code_discriminator.npy'] ]
 
             elif 'liae' in archi_type:
                 self.encoder = model_archi.Encoder(in_ch=input_ch, e_ch=e_dims, name='encoder')
-                encoder_out_ch = self.encoder.compute_output_channels ( (nn.floatx, bgr_shape))
+                encoder_out_ch = self.encoder.get_out_ch()*self.encoder.get_out_res(resolution)**2
 
                 self.inter_AB = model_archi.Inter(in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims*2, name='inter_AB')
                 self.inter_B  = model_archi.Inter(in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims*2, name='inter_B')
 
-                inter_AB_out_ch = self.inter_AB.compute_output_channels ( (nn.floatx, (None,encoder_out_ch)))
+                inter_out_ch = self.inter_AB.get_out_ch()
-                inter_B_out_ch = self.inter_B.compute_output_channels ( (nn.floatx, (None,encoder_out_ch)))
+                inters_out_ch = inter_out_ch*2
-                inters_out_ch = inter_AB_out_ch+inter_B_out_ch
                 self.decoder = model_archi.Decoder(in_ch=inters_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, name='decoder')
 
                 self.model_filename_list += [ [self.encoder,  'encoder.npy'],
@@ -293,6 +300,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
                 # Initialize optimizers
                 lr=5e-5
                 lr_dropout = 0.3 if self.options['lr_dropout'] in ['y','cpu'] and not self.pretrain else 1.0
+                OptimizerClass = nn.AdaBelief if adabelief else nn.RMSprop
                 clipnorm = 1.0 if self.options['clipgrad'] else 0.0
 
                 if 'df' in archi_type:
@@ -300,17 +308,19 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
                 elif 'liae' in archi_type:
                     self.src_dst_trainable_weights = self.encoder.get_weights() + self.inter_AB.get_weights() + self.inter_B.get_weights() + self.decoder.get_weights()
 
-                self.src_dst_opt = nn.RMSprop(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='src_dst_opt')
+
+
+                self.src_dst_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='src_dst_opt')
                 self.src_dst_opt.initialize_variables (self.src_dst_trainable_weights, vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')
                 self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ]
 
                 if self.options['true_face_power'] != 0:
-                    self.D_code_opt = nn.RMSprop(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='D_code_opt')
+                    self.D_code_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='D_code_opt')
                     self.D_code_opt.initialize_variables ( self.code_discriminator.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')
                     self.model_filename_list += [ (self.D_code_opt, 'D_code_opt.npy') ]
 
                 if gan_power != 0:
-                    self.D_src_dst_opt = nn.RMSprop(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='D_src_dst_opt')
+                    self.D_src_dst_opt = OptimizerClass(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='D_src_dst_opt')
 
                     if gan_old:
                         self.D_src_dst_opt.initialize_variables ( self.D_src.get_weights()+self.D_src_x2.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')
@@ -349,8 +359,10 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
                         gpu_warped_dst      = self.warped_dst [batch_slice,:,:,:]
                         gpu_target_src      = self.target_src [batch_slice,:,:,:]
                         gpu_target_dst      = self.target_dst [batch_slice,:,:,:]
-                        gpu_target_srcm_all = self.target_srcm_all[batch_slice,:,:,:]
+                        gpu_target_srcm_all = self.target_srcm[batch_slice,:,:,:]
-                        gpu_target_dstm_all = self.target_dstm_all[batch_slice,:,:,:]
+                        gpu_target_srcm_em  = self.target_srcm_em[batch_slice,:,:,:]
+                        gpu_target_dstm_all = self.target_dstm[batch_slice,:,:,:]
+                        gpu_target_dstm_em  = self.target_dstm_em[batch_slice,:,:,:]
 
                     # process model tensors
                     if 'df' in archi_type:
@@ -537,7 +549,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
 
 
             # Average losses and gradients, and create optimizer update ops
-            with tf.device (models_opt_device):
+            with tf.device(f'/CPU:0'):
                 pred_src_src  = nn.concat(gpu_pred_src_src_list, 0)
                 pred_dst_dst  = nn.concat(gpu_pred_dst_dst_list, 0)
                 pred_src_dst  = nn.concat(gpu_pred_src_dst_list, 0)
@@ -545,6 +557,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
                 pred_dst_dstm = nn.concat(gpu_pred_dst_dstm_list, 0)
                 pred_src_dstm = nn.concat(gpu_pred_src_dstm_list, 0)
 
+            with tf.device (models_opt_device):
                 src_loss = tf.concat(gpu_src_losses, 0)
                 dst_loss = tf.concat(gpu_dst_losses, 0)
                 src_dst_loss_gv_op = self.src_dst_opt.get_update_op (nn.average_gv_list (gpu_G_loss_gvs))
@@ -557,15 +570,17 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
 
 
             # Initializing training and view functions
-            def src_dst_train(warped_src, target_src, target_srcm_all, \
+            def src_dst_train(warped_src, target_src, target_srcm, target_srcm_em,  \
-                              warped_dst, target_dst, target_dstm_all):
+                              warped_dst, target_dst, target_dstm, target_dstm_em, ):
                 s, d, _ = nn.tf_sess.run ( [ src_loss, dst_loss, src_dst_loss_gv_op],
                                             feed_dict={self.warped_src :warped_src,
                                                        self.target_src :target_src,
-                                                       self.target_srcm_all:target_srcm_all,
+                                                       self.target_srcm:target_srcm,
+                                                       self.target_srcm_em:target_srcm_em,
                                                        self.warped_dst :warped_dst,
                                                        self.target_dst :target_dst,
-                                                       self.target_dstm_all:target_dstm_all,
+                                                       self.target_dstm:target_dstm,
+                                                       self.target_dstm_em:target_dstm_em,
                                                        })
                 return s, d
             self.src_dst_train = src_dst_train
@@ -576,14 +591,16 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
                 self.D_train = D_train
 
             if gan_power != 0:
-                def D_src_dst_train(warped_src, target_src, target_srcm_all, \
+                def D_src_dst_train(warped_src, target_src, target_srcm, target_srcm_em,  \
-                                    warped_dst, target_dst, target_dstm_all):
+                                    warped_dst, target_dst, target_dstm, target_dstm_em, ):
                     nn.tf_sess.run ([src_D_src_dst_loss_gv_op], feed_dict={self.warped_src :warped_src,
                                                                            self.target_src :target_src,
-                                                                           self.target_srcm_all:target_srcm_all,
+                                                                           self.target_srcm:target_srcm,
+                                                                           self.target_srcm_em:target_srcm_em,
                                                                            self.warped_dst :warped_dst,
                                                                            self.target_dst :target_dst,
-                                                                           self.target_dstm_all:target_dstm_all})
+                                                                           self.target_dstm:target_dstm,
+                                                                           self.target_dstm_em:target_dstm_em})
                 self.D_src_dst_train = D_src_dst_train
 
 
@@ -657,6 +674,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
                         sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
                         output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': ct_mode,                                           'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                 {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False                      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': ct_mode,                                           'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
+                                                {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False                      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G,   'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                 {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False                      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G,   'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE_EYES, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                               ],
                         uniform_yaw_distribution=self.options['uniform_yaw'] or self.pretrain,
@@ -666,6 +684,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
                         sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
                         output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': fs_aug,                                                                 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                 {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False                      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': fs_aug,                                                                 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
+                                                {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False                      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G,   'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                                 {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False                      , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G,   'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE_EYES, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
                                               ],
                         uniform_yaw_distribution=self.options['uniform_yaw'] or self.pretrain,
@@ -699,26 +718,28 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
 
         bs = self.get_batch_size()
 
-        ( (warped_src, target_src, target_srcm_all), \
+        ( (warped_src, target_src, target_srcm, target_srcm_em), \
-          (warped_dst, target_dst, target_dstm_all) ) = self.generate_next_samples()
+          (warped_dst, target_dst, target_dstm, target_dstm_em) ) = self.generate_next_samples()
 
-        src_loss, dst_loss = self.src_dst_train (warped_src, target_src, target_srcm_all, warped_dst, target_dst, target_dstm_all)
+        src_loss, dst_loss = self.src_dst_train (warped_src, target_src, target_srcm, target_srcm_em, warped_dst, target_dst, target_dstm, target_dstm_em)
 
         for i in range(bs):
-            self.last_src_samples_loss.append (  (target_src[i], target_srcm_all[i], src_loss[i] )  )
+            self.last_src_samples_loss.append (  (target_src[i], target_srcm[i], target_srcm_em[i], src_loss[i] )  )
-            self.last_dst_samples_loss.append (  (target_dst[i], target_dstm_all[i], dst_loss[i] )  )
+            self.last_dst_samples_loss.append (  (target_dst[i], target_dstm[i], target_dstm_em[i], dst_loss[i] )  )
 
         if len(self.last_src_samples_loss) >= bs*16:
-            src_samples_loss = sorted(self.last_src_samples_loss, key=operator.itemgetter(2), reverse=True)
+            src_samples_loss = sorted(self.last_src_samples_loss, key=operator.itemgetter(3), reverse=True)
-            dst_samples_loss = sorted(self.last_dst_samples_loss, key=operator.itemgetter(2), reverse=True)
+            dst_samples_loss = sorted(self.last_dst_samples_loss, key=operator.itemgetter(3), reverse=True)
 
             target_src        = np.stack( [ x[0] for x in src_samples_loss[:bs] ] )
-            target_srcm_all = np.stack( [ x[1] for x in src_samples_loss[:bs] ] )
+            target_srcm       = np.stack( [ x[1] for x in src_samples_loss[:bs] ] )
+            target_srcm_em = np.stack( [ x[2] for x in src_samples_loss[:bs] ] )
 
             target_dst        = np.stack( [ x[0] for x in dst_samples_loss[:bs] ] )
-            target_dstm_all = np.stack( [ x[1] for x in dst_samples_loss[:bs] ] )
+            target_dstm       = np.stack( [ x[1] for x in dst_samples_loss[:bs] ] )
+            target_dstm_em = np.stack( [ x[2] for x in dst_samples_loss[:bs] ] )
 
-            src_loss, dst_loss = self.src_dst_train (target_src, target_src, target_srcm_all, target_dst, target_dst, target_dstm_all)
+            src_loss, dst_loss = self.src_dst_train (target_src, target_src, target_srcm, target_srcm_em, target_dst, target_dst, target_dstm, target_dstm_em)
             self.last_src_samples_loss = []
             self.last_dst_samples_loss = []
 
@@ -726,22 +747,19 @@ Examples: df, liae, df-d, df-ud, liae-ud, ...
             self.D_train (warped_src, warped_dst)
 
         if self.gan_power != 0:
-            self.D_src_dst_train (warped_src, target_src, target_srcm_all, warped_dst, target_dst, target_dstm_all)
+            self.D_src_dst_train (warped_src, target_src, target_srcm, target_srcm_em, warped_dst, target_dst, target_dstm, target_dstm_em)
 
         return ( ('src_loss', np.mean(src_loss) ), ('dst_loss', np.mean(dst_loss) ), )
 
     #override
     def onGetPreview(self, samples):
-        ( (warped_src, target_src, target_srcm_all,),
+        ( (warped_src, target_src, target_srcm, target_srcm_em),
-          (warped_dst, target_dst, target_dstm_all,) ) = samples
+          (warped_dst, target_dst, target_dstm, target_dstm_em) ) = samples
 
         S, D, SS, DD, DDM, SD, SDM = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([target_src,target_dst] + self.AE_view (target_src, target_dst) ) ]
         DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ]
 
-        target_srcm_all, target_dstm_all = [ nn.to_data_format(x,"NHWC", self.model_data_format) for x in ([target_srcm_all, target_dstm_all] )]
+        target_srcm, target_dstm = [ nn.to_data_format(x,"NHWC", self.model_data_format) for x in ([target_srcm, target_dstm] )]
-
-        target_srcm = np.clip(target_srcm_all, 0, 1)
-        target_dstm = np.clip(target_dstm_all, 0, 1)
 
         n_samples = min(4, self.get_batch_size(), 800 // self.resolution )
 

models/Model_XSeg/Model.py

@@ -81,8 +81,9 @@ class XSegModel(ModelBase):
             gpu_loss_gvs = []
 
             for gpu_id in range(gpu_count):
-                with tf.device( f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0' ):
 
+
+                with tf.device( f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0' ):
                     with tf.device(f'/CPU:0'):
                         # slice on CPU, otherwise all batch data will be transfered to GPU first
                         batch_slice = slice( gpu_id*bs_per_gpu, (gpu_id+1)*bs_per_gpu )
@@ -100,10 +101,10 @@ class XSegModel(ModelBase):
 
 
             # Average losses and gradients, and create optimizer update ops
+            #with tf.device(f'/CPU:0'): # Temporary fix. Unknown bug with training freeze starts from 2.4.0, but 2.3.1 was ok
             with tf.device (models_opt_device):
-                pred = nn.concat(gpu_pred_list, 0)                
+                pred = tf.concat(gpu_pred_list, 0)
-                loss = tf.reduce_mean(gpu_losses)
+                loss = tf.concat(gpu_losses, 0)
-                
                 loss_gv_op = self.model.opt.get_update_op (nn.average_gv_list (gpu_loss_gvs))
 
 
@@ -157,12 +158,9 @@ class XSegModel(ModelBase):
 
     #override
     def onTrainOneIter(self):
-        
-        
         image_np, mask_np = self.generate_next_samples()[0]
         loss = self.train (image_np, mask_np)
-
+        return ( ('loss', np.mean(loss) ), )
-        return ( ('loss', loss ), )
 
     #override
     def onGetPreview(self, samples):

project.code-workspace

@@ -1,50 +0,0 @@
-{
-	"folders": [
-		{
-			"path": "."
-		}
-	],
-	"settings": {
-        "workbench.colorTheme": "Visual Studio Light",
-        "diffEditor.ignoreTrimWhitespace": true,
-        "workbench.sideBar.location": "right",
-        "breadcrumbs.enabled": false,
-        "editor.renderWhitespace": "none",
-        "editor.minimap.enabled": false,
-        "workbench.activityBar.visible": true,
-        "window.menuBarVisibility": "default",
-        "editor.fastScrollSensitivity": 10,
-        "editor.mouseWheelScrollSensitivity": 2,
-        "window.zoomLevel": 0,
-        "extensions.ignoreRecommendations": true,
-        
-        "python.linting.pylintEnabled": false,
-        "python.linting.enabled": false,
-        "python.linting.pylamaEnabled": false,
-        "python.linting.pydocstyleEnabled": false,
-        "python.pythonPath": "${env:PYTHON_EXECUTABLE}",
-		"workbench.editor.tabCloseButton": "off",
-		"workbench.editor.tabSizing": "shrink",
-		"workbench.editor.highlightModifiedTabs": true,
-		"editor.mouseWheelScrollSensitivity": 3,
-		"editor.folding": false,
-		"editor.glyphMargin": false,
-		"files.exclude": {
-			"**/__pycache__": true,
-			"**/.github": true,
-			"**/.vscode": true,
-			"**/*.dat": true,
-			"**/*.h5": true,
-            "**/*.npy": true
-		},
-		"editor.quickSuggestions": {
-			"other": false,
-			"comments": false,
-			"strings": false
-			},
-			"editor.trimAutoWhitespace": false,
-			"python.linting.pylintArgs": [
-				"--disable=import-error"
-			]
-	}
-}

requirements-colab.txt

@@ -1,9 +1,9 @@
 tqdm
-numpy==1.17.0
+numpy==1.19.3
 h5py==2.9.0
 opencv-python==4.1.0.25
 ffmpeg-python==0.1.17
 scikit-image==0.14.2
 scipy==1.4.1
 colorama
-tensorflow-gpu==1.13.2
+tensorflow-gpu==2.3.1

requirements-cuda.txt

@@ -1,11 +1,10 @@
 tqdm
-numpy==1.17.0
+numpy==1.19.3
 h5py==2.9.0
 opencv-python==4.1.0.25
 ffmpeg-python==0.1.17
 scikit-image==0.14.2
 scipy==1.4.1
 colorama
-labelme==4.2.9
+tensorflow-gpu==2.4.0
-tensorflow-gpu==1.13.2
 pyqt5

samplelib/Sample.py

@@ -79,6 +79,9 @@ class Sample(object):
 
         self._filename_offset_size = None
 
+    def has_xseg_mask(self):
+        return self.xseg_mask is not None or self.xseg_mask_compressed is not None
+        
     def get_xseg_mask(self):
         if self.xseg_mask_compressed is not None:
             xseg_mask = cv2.imdecode(self.xseg_mask_compressed, cv2.IMREAD_UNCHANGED)

samplelib/SampleGeneratorFaceXSeg.py

@@ -26,11 +26,14 @@ class SampleGeneratorFaceXSeg(SampleGeneratorBase):
         samples = sum([ SampleLoader.load (SampleType.FACE, path) for path in paths ]  )
         seg_sample_idxs = SegmentedSampleFilterSubprocessor(samples).run()
 
-        seg_samples_len = len(seg_sample_idxs)
+        if len(seg_sample_idxs) == 0:
-        if seg_samples_len == 0:
+            seg_sample_idxs = SegmentedSampleFilterSubprocessor(samples, count_xseg_mask=True).run()
+            if len(seg_sample_idxs) == 0:
                 raise Exception(f"No segmented faces found.")
             else:
-            io.log_info(f"Using {seg_samples_len} segmented samples.")
+                io.log_info(f"Using {len(seg_sample_idxs)} xseg labeled samples.")                
+        else:
+            io.log_info(f"Using {len(seg_sample_idxs)} segmented samples.")
 
         if self.debug:
             self.generators_count = 1
@@ -80,8 +83,16 @@ class SampleGeneratorFaceXSeg(SampleGeneratorBase):
         def gen_img_mask(sample):
             img = sample.load_bgr()
             h,w,c = img.shape
+            
+            if sample.seg_ie_polys.has_polys():
                 mask = np.zeros ((h,w,1), dtype=np.float32)
                 sample.seg_ie_polys.overlay_mask(mask)
+            elif sample.has_xseg_mask():                
+                mask = sample.get_xseg_mask()
+                mask[mask < 0.5] = 0.0
+                mask[mask >= 0.5] = 1.0
+            else:
+                raise Exception(f'no mask in sample {sample.filename}')
 
             if face_type == sample.face_type:
                 if w != resolution:
@@ -158,9 +169,10 @@ class SampleGeneratorFaceXSeg(SampleGeneratorBase):
 
 class SegmentedSampleFilterSubprocessor(Subprocessor):
     #override
-    def __init__(self, samples ):
+    def __init__(self, samples, count_xseg_mask=False ):
         self.samples = samples
         self.samples_len = len(self.samples)
+        self.count_xseg_mask = count_xseg_mask
 
         self.idxs = [*range(self.samples_len)]
         self.result = []
@@ -169,7 +181,7 @@ class SegmentedSampleFilterSubprocessor(Subprocessor):
     #override
     def process_info_generator(self):
         for i in range(multiprocessing.cpu_count()):
-            yield 'CPU%d' % (i), {}, {'samples':self.samples}
+            yield 'CPU%d' % (i), {}, {'samples':self.samples, 'count_xseg_mask':self.count_xseg_mask}
 
     #override
     def on_clients_initialized(self):
@@ -203,6 +215,10 @@ class SegmentedSampleFilterSubprocessor(Subprocessor):
         #overridable optional
         def on_initialize(self, client_dict):
             self.samples = client_dict['samples']
+            self.count_xseg_mask = client_dict['count_xseg_mask']
 
         def process_data(self, idx):
+            if self.count_xseg_mask:
+                return idx, self.samples[idx].has_xseg_mask()
+            else:
                 return idx, self.samples[idx].seg_ie_polys.get_pts_count() != 0

samplelib/SampleLoader.py

@@ -81,7 +81,7 @@ class SampleLoader:
               shape,
               landmarks,
               seg_ie_polys,
-              xseg_mask,
+              xseg_mask_compressed,
               eyebrows_expand_mod,
               source_filename ) = data
               
@@ -91,7 +91,7 @@ class SampleLoader:
                                         shape=shape,
                                         landmarks=landmarks,
                                         seg_ie_polys=seg_ie_polys,
-                                        xseg_mask=xseg_mask,
+                                        xseg_mask_compressed=xseg_mask_compressed,
                                         eyebrows_expand_mod=eyebrows_expand_mod,
                                         source_filename=source_filename,
                                     ))
@@ -163,7 +163,7 @@ class FaceSamplesLoaderSubprocessor(Subprocessor):
                         dflimg.get_shape(),
                         dflimg.get_landmarks(),
                         dflimg.get_seg_ie_polys(),
-                        dflimg.get_xseg_mask(),
+                        dflimg.get_xseg_mask_compressed(),
                         dflimg.get_eyebrows_expand_mod(),
                         dflimg.get_source_filename() )
 

samplelib/SampleProcessor.py

@@ -29,9 +29,9 @@ class SampleProcessor(object):
 
     class FaceMaskType(IntEnum):
         NONE          = 0
-        FULL_FACE      = 1  #mask all hull as grayscale
+        FULL_FACE      = 1  # mask all hull as grayscale
-        EYES     = 2  #mask eyes hull as grayscale
+        EYES           = 2  # mask eyes hull as grayscale
-        FULL_FACE_EYES = 3  #combo all + eyes as grayscale
+        FULL_FACE_EYES     = 3  # eyes and mouse
 
     class Options(object):
         def __init__(self, random_flip = True, rotation_range=[-10,10], scale_range=[-0.05, 0.05], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05] ):
@@ -149,9 +149,9 @@ class SampleProcessor(object):
                             # sets both eyes and mouth mask parts
                             img = get_full_face_mask()                            
                             eye_mask = get_eyes_mask()
-                            img = np.where(eye_mask > 1, eye_mask, img)
+                            img = np.where(eye_mask > 1 and img > 0, eye_mask, img)
                             mouth_mask = get_mouth_mask()
-                            img = np.where(mouth_mask > 2, mouth_mask, img)
+                            img = np.where(mouth_mask > 2 and img > 0, mouth_mask, img)
                         else:
                             img = np.zeros ( sample_bgr.shape[0:2]+(1,), dtype=np.float32)