mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-20 13:33:24 -07:00
merge with mask bounds fix
This commit is contained in:
sinofis
commit
40a416e415
31 changed files with 534 additions and 354 deletions
2
.vscode/launch.json
vendored
2
.vscode/launch.json
vendored
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
31
README.md
31
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
|
||||
|
||||
| [Ctrl Shift Face](https://www.youtube.com/channel/UCKpH0CKltc73e4wh0_pgL3g)| [VFXChris Ume](https://www.youtube.com/channel/UCGf4OlX_aTt8DlrgiH3jN3g/videos)|
|
||||
|
||||
|
||||
| [Ctrl Shift Face](https://www.youtube.com/channel/UCKpH0CKltc73e4wh0_pgL3g)| [VFXChris Ume](https://www.youtube.com/channel/UCGf4OlX_aTt8DlrgiH3jN3g/videos)| [Sham00k](https://www.youtube.com/channel/UCZXbWcv7fSZFTAZV4beckyw/videos)|
|
||||
|---|---|---|
|
||||
|
||||
|
||||
| [Collider videos](https://www.youtube.com/watch?v=A91P2qtPT54&list=PLayt6616lBclvOprvrC8qKGCO-mAhPRux)| [iFake](https://www.youtube.com/channel/UCC0lK2Zo2BMXX-k1Ks0r7dg/videos)| [NextFace](https://www.youtube.com/channel/UCFh3gL0a8BS21g-DHvXZEeQ/videos)|
|
||||
|---|---|---|
|
||||
|
||||
| [Futuring Machine](https://www.youtube.com/channel/UCC5BbFxqLQgfnWPhprmQLVg)| [RepresentUS](https://www.youtube.com/channel/UCRzgK52MmetD9aG8pDOID3g)|
|
||||
|---|---|
|
||||
|
||||
| [Sham00k](https://www.youtube.com/channel/UCZXbWcv7fSZFTAZV4beckyw/videos)| [Collider videos](https://www.youtube.com/watch?v=A91P2qtPT54&list=PLayt6616lBclvOprvrC8qKGCO-mAhPRux)| [iFake](https://www.youtube.com/channel/UCC0lK2Zo2BMXX-k1Ks0r7dg/videos)|
|
||||
|---|---|---|
|
||||
|
||||
| [NextFace](https://www.youtube.com/channel/UCFh3gL0a8BS21g-DHvXZEeQ/videos)| [Futuring Machine](https://www.youtube.com/channel/UCC5BbFxqLQgfnWPhprmQLVg)| [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">
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
80
core/leras/optimizers/AdaBelief.py
Normal file
80
core/leras/optimizers/AdaBelief.py
Normal file
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -1,2 +1,3 @@
|
|||
from .OptimizerBase import *
|
||||
from .RMSprop import *
|
||||
from .RMSprop import *
|
||||
from .AdaBelief import *
|
||||
|
|
@ -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('<QQ', bytes(sh_b[ table_offset + key*8 : table_offset + (key+2)*8]) )
|
||||
sh_b = memoryview(sh_b).cast('B')
|
||||
|
||||
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() )
|
||||
|
||||
return pickle.loads( bytes(sh_b[ data_offset + offset_start : data_offset + offset_end ]) )
|
||||
|
||||
def __iter__(self):
|
||||
for i in range(self.__len__()):
|
||||
yield self.__getitem__(i)
|
||||
|
||||
|
||||
@staticmethod
|
||||
def bake_data(obj_list):
|
||||
if not isinstance(obj_list, list):
|
||||
raise ValueError("MPSharedList: obj_list should be list type.")
|
||||
|
||||
|
||||
obj_count = len(obj_list)
|
||||
|
||||
if obj_count != 0:
|
||||
|
||||
if obj_count != 0:
|
||||
obj_pickled_ar = [pickle.dumps(o, 4) for o in obj_list]
|
||||
|
||||
|
||||
table_offset = 0
|
||||
table_size = (obj_count+1)*8
|
||||
data_offset = table_offset + table_size
|
||||
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)
|
||||
|
||||
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_view = memoryview(sh_b).cast('B')
|
||||
|
||||
offset = 0
|
||||
|
||||
sh_b_table = bytes()
|
||||
offsets = []
|
||||
|
||||
|
||||
offset = 0
|
||||
for i in range(obj_count):
|
||||
offsets.append(offset)
|
||||
offset += len(obj_pickled_ar[i])
|
||||
offsets.append(offset)
|
||||
|
||||
sh_b[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()
|
||||
|
||||
sh_b_view[table_offset:table_offset+table_size] = struct.pack( '<'+'Q'*len(offsets), *offsets )
|
||||
|
||||
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
|
||||
|
|
|
|||
BIN
doc/replace_the_face.jpg
Normal file
BIN
doc/replace_the_face.jpg
Normal file
Binary file not shown.
|
After Width: | Height: | Size: 378 KiB |
Binary file not shown.
|
Before Width: | Height: | Size: 1.2 MiB |
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
2
main.py
2
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):
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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')
|
||||
|
|
|
|||
|
|
@ -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 )
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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"
|
||||
]
|
||||
}
|
||||
}
|
||||
|
|
@ -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
|
||||
|
|
@ -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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
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
|
||||
|
|
@ -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() )
|
||||
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue