|
diff --git a/core/imagelib/warp.py b/core/imagelib/warp.py
index 9252532..ac4f324 100644
--- a/core/imagelib/warp.py
+++ b/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
diff --git a/core/leras/archis/DeepFakeArchi.py b/core/leras/archis/DeepFakeArchi.py
index ffd9201..9818ab8 100644
--- a/core/leras/archis/DeepFakeArchi.py
+++ b/core/leras/archis/DeepFakeArchi.py
@@ -72,11 +72,22 @@ class DeepFakeArchi(nn.ArchiBase):
return x
class Encoder(nn.ModelBase):
- def on_build(self, in_ch, e_ch):
- self.down1 = DownscaleBlock(in_ch, e_ch, n_downscales=4, kernel_size=5)
+ def __init__(self, in_ch, e_ch, **kwargs ):
+ 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)
@@ -104,9 +115,8 @@ class DeepFakeArchi(nn.ArchiBase):
x = self.upscale1(x)
return x
- @staticmethod
- def get_code_res():
- return lowest_dense_res
+ def get_out_res(self):
+ return lowest_dense_res * 2
def get_out_ch(self):
return self.ae_out_ch
diff --git a/core/leras/layers/Saveable.py b/core/leras/layers/Saveable.py
index 13eff3b..0cc0e94 100644
--- a/core/leras/layers/Saveable.py
+++ b/core/leras/layers/Saveable.py
@@ -76,25 +76,28 @@ class Saveable():
if self.name is None:
raise Exception("name must be defined.")
- tuples = []
- for w in weights:
- w_name_split = w.name.split('/')
- if self.name != w_name_split[0]:
- raise Exception("weight first name != Saveable.name")
+ try:
+ tuples = []
+ for w in weights:
+ w_name_split = w.name.split('/')
+ if self.name != w_name_split[0]:
+ raise Exception("weight first name != Saveable.name")
- sub_w_name = "/".join(w_name_split[1:])
+ sub_w_name = "/".join(w_name_split[1:])
- w_val = d.get(sub_w_name, None)
+ w_val = d.get(sub_w_name, None)
- if w_val is None:
- #io.log_err(f"Weight {w.name} was not loaded from file {filename}")
- tuples.append ( (w, w.initializer) )
- else:
- w_val = np.reshape( w_val, w.shape.as_list() )
- tuples.append ( (w, w_val) )
-
- nn.batch_set_value(tuples)
+ if w_val is None:
+ #io.log_err(f"Weight {w.name} was not loaded from file {filename}")
+ tuples.append ( (w, w.initializer) )
+ else:
+ w_val = np.reshape( w_val, w.shape.as_list() )
+ tuples.append ( (w, w_val) )
+ nn.batch_set_value(tuples)
+ except:
+ return False
+
return True
def init_weights(self):
diff --git a/core/leras/models/ModelBase.py b/core/leras/models/ModelBase.py
index 77ac284..cc558a4 100644
--- a/core/leras/models/ModelBase.py
+++ b/core/leras/models/ModelBase.py
@@ -116,41 +116,32 @@ class ModelBase(nn.Saveable):
return self.forward(*args, **kwargs)
- def compute_output_shape(self, shapes):
- if not self.built:
- self.build()
+ # def compute_output_shape(self, shapes):
+ # if not self.built:
+ # self.build()
- not_list = False
- if not isinstance(shapes, list):
- not_list = True
- shapes = [shapes]
+ # not_list = False
+ # if not isinstance(shapes, list):
+ # not_list = True
+ # shapes = [shapes]
- with tf.device('/CPU:0'):
- # CPU tensors will not impact any performance, only slightly RAM "leakage"
- phs = []
- for dtype,sh in shapes:
- phs += [ tf.placeholder(dtype, sh) ]
+ # with tf.device('/CPU:0'):
+ # # CPU tensors will not impact any performance, only slightly RAM "leakage"
+ # phs = []
+ # for dtype,sh in shapes:
+ # 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):
- result = [result]
+ # if not isinstance(result, list):
+ # result = [result]
- result_shapes = []
+ # result_shapes = []
- for t in result:
- result_shapes += [ t.shape.as_list() ]
+ # for t in result:
+ # result_shapes += [ t.shape.as_list() ]
- 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]
+ # return result_shapes[0] if not_list else result_shapes
def build_for_run(self, shapes_list):
if not isinstance(shapes_list, list):
diff --git a/core/leras/nn.py b/core/leras/nn.py
index 6504698..7d62393 100644
--- a/core/leras/nn.py
+++ b/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
- nn.tf = tf
+ #import tensorflow as 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
diff --git a/core/leras/ops/__init__.py b/core/leras/ops/__init__.py
index ddf1c06..648407c 100644
--- a/core/leras/ops/__init__.py
+++ b/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):
diff --git a/core/leras/optimizers/AdaBelief.py b/core/leras/optimizers/AdaBelief.py
new file mode 100644
index 0000000..aed4308
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/core/leras/optimizers/RMSprop.py b/core/leras/optimizers/RMSprop.py
index 37cebe9..345b2a7 100644
--- a/core/leras/optimizers/RMSprop.py
+++ b/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 = rho * a + (1. - rho) * tf.square(g)
+ new_a = self.rho * a + (1. - self.rho) * tf.square(g)
- lr = tf.cast(self.lr, a.dtype)
- epsilon = tf.cast(self.epsilon, a.dtype)
+ lr = tf.constant(self.lr, g.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
diff --git a/core/leras/optimizers/__init__.py b/core/leras/optimizers/__init__.py
index aec36af..4f8a7e4 100644
--- a/core/leras/optimizers/__init__.py
+++ b/core/leras/optimizers/__init__.py
@@ -1,2 +1,3 @@
from .OptimizerBase import *
-from .RMSprop import *
\ No newline at end of file
+from .RMSprop import *
+from .AdaBelief import *
\ No newline at end of file
diff --git a/core/mplib/MPSharedList.py b/core/mplib/MPSharedList.py
index 658f9fc..874c56a 100644
--- a/core/mplib/MPSharedList.py
+++ b/core/mplib/MPSharedList.py
@@ -7,52 +7,52 @@ class MPSharedList():
"""
Provides read-only pickled list of constant objects via shared memory aka 'multiprocessing.Array'
Thus no 4GB limit for subprocesses.
-
+
supports list concat via + or sum()
"""
-
- def __init__(self, obj_list):
+
+ def __init__(self, obj_list):
if obj_list is None:
self.obj_counts = None
self.table_offsets = None
self.data_offsets = None
self.sh_bs = None
- else:
+ else:
obj_count, table_offset, data_offset, sh_b = MPSharedList.bake_data(obj_list)
self.obj_counts = [obj_count]
self.table_offsets = [table_offset]
self.data_offsets = [data_offset]
self.sh_bs = [sh_b]
-
+
def __add__(self, o):
if isinstance(o, MPSharedList):
m = MPSharedList(None)
- m.obj_counts = self.obj_counts + o.obj_counts
+ m.obj_counts = self.obj_counts + o.obj_counts
m.table_offsets = self.table_offsets + o.table_offsets
- m.data_offsets = self.data_offsets + o.data_offsets
- m.sh_bs = self.sh_bs + o.sh_bs
+ m.data_offsets = self.data_offsets + o.data_offsets
+ m.sh_bs = self.sh_bs + o.sh_bs
return m
elif isinstance(o, int):
return self
else:
raise ValueError(f"MPSharedList object of class {o.__class__} is not supported for __add__ operator.")
-
+
def __radd__(self, o):
return self+o
-
+
def __len__(self):
return sum(self.obj_counts)
-
- def __getitem__(self, key):
+
+ def __getitem__(self, key):
obj_count = sum(self.obj_counts)
if key < 0:
key = obj_count+key
if key < 0 or key >= obj_count:
raise ValueError("out of range")
-
+
for i in range(len(self.obj_counts)):
-
+
if key < self.obj_counts[i]:
table_offset = self.table_offsets[i]
data_offset = self.data_offsets[i]
@@ -60,88 +60,52 @@ class MPSharedList():
break
key -= self.obj_counts[i]
- offset_start, offset_end = struct.unpack(' 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
diff --git a/doc/replace_the_face.jpg b/doc/replace_the_face.jpg
new file mode 100644
index 0000000..55501d0
Binary files /dev/null and b/doc/replace_the_face.jpg differ
diff --git a/doc/replace_the_face.png b/doc/replace_the_face.png
deleted file mode 100644
index 7d612dc..0000000
Binary files a/doc/replace_the_face.png and /dev/null differ
diff --git a/facelib/LandmarksProcessor.py b/facelib/LandmarksProcessor.py
index 23430e5..d4ef910 100644
--- a/facelib/LandmarksProcessor.py
+++ b/facelib/LandmarksProcessor.py
@@ -454,7 +454,7 @@ def get_image_mouth_mask (image_shape, image_landmarks):
hull_mask = hull_mask[...,None]
return hull_mask
-
+
def alpha_to_color (img_alpha, color):
if len(img_alpha.shape) == 2:
img_alpha = img_alpha[...,None]
diff --git a/facelib/XSegNet.py b/facelib/XSegNet.py
index f260e21..761ab94 100644
--- a/facelib/XSegNet.py
+++ b/facelib/XSegNet.py
@@ -29,7 +29,10 @@ 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()
+ 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' ] ]
+
+
+ self.model_filename_list += [ [self.model, f'{model_name}.npy'] ]
- 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:
+ if not training:
with tf.device ('/CPU:0' if run_on_cpu else '/GPU:0'):
_, pred = self.model(self.input_t)
diff --git a/main.py b/main.py
index fb64f96..2ba085f 100644
--- a/main.py
+++ b/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):
diff --git a/mainscripts/Sorter.py b/mainscripts/Sorter.py
index 9b09e43..39eec5e 100644
--- a/mainscripts/Sorter.py
+++ b/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] )
@@ -36,8 +39,14 @@ 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):
@@ -103,7 +113,18 @@ def sort_by_blur(input_path):
img_list = sorted(img_list, key=operator.itemgetter(1), reverse=True)
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:
diff --git a/mainscripts/Trainer.py b/mainscripts/Trainer.py
index 07d1307..7d73e2f 100644
--- a/mainscripts/Trainer.py
+++ b/mainscripts/Trainer.py
@@ -1,4 +1,5 @@
-import sys
+import os
+import sys
import traceback
import queue
import threading
@@ -119,6 +120,12 @@ def trainerThread (s2c, c2s, e,
io.log_info("")
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()
diff --git a/mainscripts/dev_misc.py b/mainscripts/dev_misc.py
index 9f9c1fa..a71bcdf 100644
--- a/mainscripts/dev_misc.py
+++ b/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.')
+
+ 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)
+
+ data = []
+
+ img_paths = pathex.get_image_paths (images_path)
+ for filename in img_paths:
+ filepath = Path(filename)
- dir_names = pathex.get_all_dir_names(input_path)
+ landmark_filepath = landmarks_path / (filepath.stem + '.txt')
+ if not landmark_filepath.exists():
+ raise ValueError(f'no landmarks for {filepath}')
+
+ #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) ]) ]
- for dir_name in io.progress_bar_generator(dir_names, desc="Processing"):
+ #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()
- img_paths = pathex.get_image_paths (input_path / dir_name)
- for filename in img_paths:
+
+ 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()
- dflimg = DFLIMG.x (filepath)
- if dflimg is None:
- raise ValueError
-
- #dflimg.x(filename, person_name=dir_name)
-
- #import code
- #code.interact(local=dict(globals(), **locals()))
+ 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()))
+
def dev_resave_pngs(input_dir):
input_path = Path(input_dir)
diff --git a/models/Model_Quick96/Model.py b/models/Model_Quick96/Model.py
index c85354f..3c39e46 100644
--- a/models/Model_Quick96/Model.py
+++ b/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')
diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py
index d34e7dd..c5a44ce 100644
--- a/models/Model_SAEHD/Model.py
+++ b/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_dstm_all = tf.placeholder (nn.floatx, mask_shape)
+ self.target_srcm = 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_B_out_ch = self.inter_B.compute_output_channels ( (nn.floatx, (None,encoder_out_ch)))
- inters_out_ch = inter_AB_out_ch+inter_B_out_ch
+ inter_out_ch = self.inter_AB.get_out_ch()
+ inters_out_ch = inter_out_ch*2
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_dstm_all = self.target_dstm_all[batch_slice,:,:,:]
+ gpu_target_srcm_all = self.target_srcm[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, \
- warped_dst, target_dst, target_dstm_all):
+ def src_dst_train(warped_src, target_src, target_srcm, target_srcm_em, \
+ 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, \
- warped_dst, target_dst, target_dstm_all):
+ def D_src_dst_train(warped_src, target_src, target_srcm, target_srcm_em, \
+ 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_dst, target_dst, target_dstm_all) ) = self.generate_next_samples()
+ ( (warped_src, target_src, target_srcm, target_srcm_em), \
+ (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_dst_samples_loss.append ( (target_dst[i], target_dstm_all[i], dst_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[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)
- dst_samples_loss = sorted(self.last_dst_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(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_src = np.stack( [ x[0] 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_dst = np.stack( [ x[0] 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_dst, target_dst, target_dstm_all,) ) = samples
+ ( (warped_src, target_src, target_srcm, target_srcm_em),
+ (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 = np.clip(target_srcm_all, 0, 1)
- target_dstm = np.clip(target_dstm_all, 0, 1)
+ target_srcm, target_dstm = [ nn.to_data_format(x,"NHWC", self.model_data_format) for x in ([target_srcm, target_dstm] )]
n_samples = min(4, self.get_batch_size(), 800 // self.resolution )
diff --git a/models/Model_XSeg/Model.py b/models/Model_XSeg/Model.py
index 03a6fb5..4ab23fc 100644
--- a/models/Model_XSeg/Model.py
+++ b/models/Model_XSeg/Model.py
@@ -15,23 +15,23 @@ class XSegModel(ModelBase):
def __init__(self, *args, **kwargs):
super().__init__(*args, force_model_class_name='XSeg', **kwargs)
-
+
#override
- def on_initialize_options(self):
+ def on_initialize_options(self):
ask_override = self.ask_override()
- if not self.is_first_run() and ask_override:
+ if not self.is_first_run() and ask_override:
if io.input_bool(f"Restart training?", False, help_message="Reset model weights and start training from scratch."):
self.set_iter(0)
default_face_type = self.options['face_type'] = self.load_or_def_option('face_type', 'wf')
-
+
if self.is_first_run():
self.options['face_type'] = io.input_str ("Face type", default_face_type, ['h','mf','f','wf','head'], help_message="Half / mid face / full face / whole face / head. Choose the same as your deepfake model.").lower()
if self.is_first_run() or ask_override:
- self.ask_batch_size(4, range=[2,16])
-
+ self.ask_batch_size(4, range=[2,16])
+
#override
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
@@ -44,29 +44,29 @@ class XSegModel(ModelBase):
self.resolution = resolution = 256
-
+
self.face_type = {'h' : FaceType.HALF,
'mf' : FaceType.MID_FULL,
'f' : FaceType.FULL,
'wf' : FaceType.WHOLE_FACE,
'head' : FaceType.HEAD}[ self.options['face_type'] ]
-
+
place_model_on_cpu = len(devices) == 0
models_opt_device = '/CPU:0' if place_model_on_cpu else '/GPU:0'
bgr_shape = nn.get4Dshape(resolution,resolution,3)
mask_shape = nn.get4Dshape(resolution,resolution,1)
-
+
# Initializing model classes
- self.model = XSegNet(name='XSeg',
- resolution=resolution,
+ self.model = XSegNet(name='XSeg',
+ resolution=resolution,
load_weights=not self.is_first_run(),
weights_file_root=self.get_model_root_path(),
training=True,
place_model_on_cpu=place_model_on_cpu,
optimizer=nn.RMSprop(lr=0.0001, lr_dropout=0.3, name='opt'),
data_format=nn.data_format)
-
+
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices) )
@@ -79,20 +79,21 @@ class XSegModel(ModelBase):
gpu_losses = []
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' ):
+ 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'/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 )
gpu_input_t = self.model.input_t [batch_slice,:,:,:]
- gpu_target_t = self.model.target_t [batch_slice,:,:,:]
-
+ gpu_target_t = self.model.target_t [batch_slice,:,:,:]
+
# process model tensors
- gpu_pred_logits_t, gpu_pred_t = self.model.flow(gpu_input_t)
+ gpu_pred_logits_t, gpu_pred_t = self.model.flow(gpu_input_t)
gpu_pred_list.append(gpu_pred_t)
-
+
gpu_loss = tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=gpu_target_t, logits=gpu_pred_logits_t), axis=[1,2,3])
gpu_losses += [gpu_loss]
@@ -100,13 +101,13 @@ 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)
- loss = tf.reduce_mean(gpu_losses)
-
+ pred = tf.concat(gpu_pred_list, 0)
+ loss = tf.concat(gpu_losses, 0)
loss_gv_op = self.model.opt.get_update_op (nn.average_gv_list (gpu_loss_gvs))
-
-
+
+
# Initializing training and view functions
def train(input_np, target_np):
l, _ = nn.tf_sess.run ( [loss, loss_gv_op], feed_dict={self.model.input_t :input_np, self.model.target_t :target_np })
@@ -122,29 +123,29 @@ class XSegModel(ModelBase):
src_dst_generators_count = cpu_count // 2
src_generators_count = cpu_count // 2
dst_generators_count = cpu_count // 2
-
-
+
+
srcdst_generator = SampleGeneratorFaceXSeg([self.training_data_src_path, self.training_data_dst_path],
debug=self.is_debug(),
batch_size=self.get_batch_size(),
resolution=resolution,
face_type=self.face_type,
- generators_count=src_dst_generators_count,
+ generators_count=src_dst_generators_count,
data_format=nn.data_format)
-
+
src_generator = SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=False),
output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE, 'warp':False, 'transform':False, 'channel_type' : SampleProcessor.ChannelType.BGR, 'border_replicate':False, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
],
generators_count=src_generators_count,
- raise_on_no_data=False )
+ raise_on_no_data=False )
dst_generator = SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=False),
output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE, 'warp':False, 'transform':False, 'channel_type' : SampleProcessor.ChannelType.BGR, 'border_replicate':False, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
],
generators_count=dst_generators_count,
raise_on_no_data=False )
-
+
self.set_training_data_generators ([srcdst_generator, src_generator, dst_generator])
#override
@@ -154,21 +155,18 @@ class XSegModel(ModelBase):
#override
def onSave(self):
self.model.save_weights()
-
- #override
- def onTrainOneIter(self):
-
-
- image_np, mask_np = self.generate_next_samples()[0]
- loss = self.train (image_np, mask_np)
- return ( ('loss', loss ), )
+ #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) ), )
#override
def onGetPreview(self, samples):
n_samples = min(4, self.get_batch_size(), 800 // self.resolution )
- srcdst_samples, src_samples, dst_samples = samples
+ srcdst_samples, src_samples, dst_samples = samples
image_np, mask_np = srcdst_samples
I, M, IM, = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([image_np,mask_np] + self.view (image_np) ) ]
@@ -176,41 +174,41 @@ class XSegModel(ModelBase):
green_bg = np.tile( np.array([0,1,0], dtype=np.float32)[None,None,...], (self.resolution,self.resolution,1) )
- result = []
+ result = []
st = []
for i in range(n_samples):
ar = I[i]*M[i]+0.5*I[i]*(1-M[i])+0.5*green_bg*(1-M[i]), IM[i], I[i]*IM[i]+0.5*I[i]*(1-IM[i]) + 0.5*green_bg*(1-IM[i])
st.append ( np.concatenate ( ar, axis=1) )
result += [ ('XSeg training faces', np.concatenate (st, axis=0 )), ]
-
+
if len(src_samples) != 0:
src_np, = src_samples
-
+
D, DM, = [ np.clip(nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([src_np] + self.view (src_np) ) ]
DM, = [ np.repeat (x, (3,), -1) for x in [DM] ]
-
+
st = []
for i in range(n_samples):
ar = D[i], DM[i], D[i]*DM[i] + 0.5*D[i]*(1-DM[i]) + 0.5*green_bg*(1-DM[i])
st.append ( np.concatenate ( ar, axis=1) )
-
+
result += [ ('XSeg src faces', np.concatenate (st, axis=0 )), ]
-
+
if len(dst_samples) != 0:
dst_np, = dst_samples
-
+
D, DM, = [ np.clip(nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([dst_np] + self.view (dst_np) ) ]
DM, = [ np.repeat (x, (3,), -1) for x in [DM] ]
-
+
st = []
for i in range(n_samples):
ar = D[i], DM[i], D[i]*DM[i] + 0.5*D[i]*(1-DM[i]) + 0.5*green_bg*(1-DM[i])
st.append ( np.concatenate ( ar, axis=1) )
-
+
result += [ ('XSeg dst faces', np.concatenate (st, axis=0 )), ]
-
+
return result
Model = XSegModel
\ No newline at end of file
diff --git a/project.code-workspace b/project.code-workspace
deleted file mode 100644
index 07fae2f..0000000
--- a/project.code-workspace
+++ /dev/null
@@ -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"
- ]
- }
-}
\ No newline at end of file
diff --git a/requirements-colab.txt b/requirements-colab.txt
index 128a518..814ed07 100644
--- a/requirements-colab.txt
+++ b/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
\ No newline at end of file
+tensorflow-gpu==2.3.1
\ No newline at end of file
diff --git a/requirements-cuda.txt b/requirements-cuda.txt
index 8df3478..5fb89ee 100644
--- a/requirements-cuda.txt
+++ b/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==1.13.2
+tensorflow-gpu==2.4.0
pyqt5
\ No newline at end of file
diff --git a/samplelib/Sample.py b/samplelib/Sample.py
index 7ced210..a379275 100644
--- a/samplelib/Sample.py
+++ b/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)
diff --git a/samplelib/SampleGeneratorFaceXSeg.py b/samplelib/SampleGeneratorFaceXSeg.py
index 7999d16..9c7cf6e 100644
--- a/samplelib/SampleGeneratorFaceXSeg.py
+++ b/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 seg_samples_len == 0:
- raise Exception(f"No segmented faces found.")
+ if len(seg_sample_idxs) == 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 {len(seg_sample_idxs)} xseg labeled samples.")
else:
- io.log_info(f"Using {seg_samples_len} segmented samples.")
+ 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
- mask = np.zeros ((h,w,1), dtype=np.float32)
- sample.seg_ie_polys.overlay_mask(mask)
+
+ 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):
- return idx, self.samples[idx].seg_ie_polys.get_pts_count() != 0
\ No newline at end of file
+ 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
\ No newline at end of file
diff --git a/samplelib/SampleLoader.py b/samplelib/SampleLoader.py
index 46d0962..edb3775 100644
--- a/samplelib/SampleLoader.py
+++ b/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() )
diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py
index 2cfc8b3..e2894a2 100644
--- a/samplelib/SampleProcessor.py
+++ b/samplelib/SampleProcessor.py
@@ -29,9 +29,9 @@ class SampleProcessor(object):
class FaceMaskType(IntEnum):
NONE = 0
- FULL_FACE = 1 #mask all hull as grayscale
- EYES = 2 #mask eyes hull as grayscale
- FULL_FACE_EYES = 3 #combo all + eyes as grayscale
+ FULL_FACE = 1 # mask all hull as grayscale
+ EYES = 2 # mask eyes hull 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)
@@ -170,7 +170,7 @@ class SampleProcessor(object):
img = cv2.resize( img, (resolution, resolution), interpolation=cv2.INTER_LINEAR )
img = imagelib.warp_by_params (params_per_resolution[resolution], img, warp, transform, can_flip=True, border_replicate=border_replicate, cv2_inter=cv2.INTER_LINEAR)
-
+
if len(img.shape) == 2:
img = img[...,None]
|
+
