def on_initialize(self, client_dict):
self.type = client_dict['type']
self.image_size = client_dict['image_size']
self.face_type = client_dict['face_type']
self.max_faces_from_image = client_dict['max_faces_from_image']
self.device_idx = client_dict['device_idx']
self.cpu_only = client_dict['device_type'] == 'CPU'
self.final_output_path = client_dict['final_output_path']
self.output_debug_path = client_dict['output_debug_path']
#transfer and set stdin in order to work code.interact in debug subprocess
stdin_fd = client_dict['stdin_fd']
if stdin_fd is not None and DEBUG:
sys.stdin = os.fdopen(stdin_fd)
if self.cpu_only:
device_config = nn.DeviceConfig.CPU()
place_model_on_cpu = True
else:
device_config = nn.DeviceConfig.GPUIndexes ([self.device_idx])
place_model_on_cpu = device_config.devices[0].total_mem_gb < 4
if self.type == 'all' or 'rects' in self.type or 'landmarks' in self.type:
nn.initialize (device_config)
self.log_info (f"Running on {client_dict['device_name'] }")
if self.type == 'all' or self.type == 'rects-s3fd' or 'landmarks' in self.type:
self.rects_extractor = facelib.S3FDExtractor(place_model_on_cpu=place_model_on_cpu)
if self.type == 'all' or 'landmarks' in self.type:
self.landmarks_extractor = facelib.FANExtractor(place_model_on_cpu=place_model_on_cpu)
self.cached_image = (None, None)
def on_initialize(self, client_dict):
device_idx = client_dict['device_idx']
cpu_only = client_dict['device_type'] == 'CPU'
self.output_dirpath = client_dict['output_dirpath']
nn_initialize_mp_lock = client_dict['nn_initialize_mp_lock']
if cpu_only:
device_config = nn.DeviceConfig.CPU()
device_vram = 99
else:
device_config = nn.DeviceConfig.GPUIndexes([device_idx])
device_vram = device_config.devices[0].total_mem_gb
nn.initialize(device_config)
intro_str = 'Running on %s.' % (client_dict['device_name'])
self.log_info(intro_str)
from facelib import FaceEnhancer
self.fe = FaceEnhancer(place_model_on_cpu=(device_vram <= 2))
def apply_xseg(input_path, model_path):
if not input_path.exists():
raise ValueError(f'{input_path} not found. Please ensure it exists.')
if not model_path.exists():
raise ValueError(f'{model_path} not found. Please ensure it exists.')
io.log_info(f'Applying trained XSeg model to {input_path.name}/ folder.')
device_config = nn.DeviceConfig.ask_choose_device(choose_only_one=True)
nn.initialize(device_config)
xseg = XSegNet(name='XSeg',
load_weights=True,
weights_file_root=model_path,
data_format=nn.data_format,
raise_on_no_model_files=True)
res = xseg.get_resolution()
images_paths = pathex.get_image_paths(input_path, return_Path_class=True)
for filepath in io.progress_bar_generator(images_paths, "Processing"):
dflimg = DFLIMG.load(filepath)
if dflimg is None or not dflimg.has_data():
io.log_info(f'{filepath} is not a DFLIMG')
continue
img = cv2_imread(filepath).astype(np.float32) / 255.0
h, w, c = img.shape
if w != res:
img = cv2.resize(img, (res, res), interpolation=cv2.INTER_CUBIC)
if len(img.shape) == 2:
img = img[..., None]
mask = xseg.extract(img)
mask[mask < 0.5] = 0
mask[mask >= 0.5] = 1
dflimg.set_xseg_mask(mask)
dflimg.save()
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
self.model_data_format = "NCHW" if len(
device_config.devices) != 0 and not self.is_debug() else "NHWC"
nn.initialize(data_format=self.model_data_format)
tf = nn.tf
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
self.resolution = resolution = 256
self.face_type = FaceType.WHOLE_FACE
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=f'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))
bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
self.set_batch_size(gpu_count * bs_per_gpu)
# Compute losses per GPU
gpu_pred_list = []
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'):
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, :, :, :]
# process model tensors
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]
gpu_loss_gvs += [
nn.gradients(gpu_loss, self.model.get_weights())
]
# Average losses and gradients, and create optimizer update ops
with tf.device(models_opt_device):
pred = nn.concat(gpu_pred_list, 0)
loss = tf.reduce_mean(gpu_losses)
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
})
return l
self.train = train
def view(input_np):
return nn.tf_sess.run([pred],
feed_dict={self.model.input_t: input_np})
self.view = view
# initializing sample generators
cpu_count = min(multiprocessing.cpu_count(), 8)
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,
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)
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])
def on_initialize(self):
nn.initialize()
tf = nn.tf
nn.set_floatx(tf.float32)
conv_kernel_initializer = nn.initializers.ca
class Downscale(nn.ModelBase):
def __init__(self, in_ch, out_ch, kernel_size=5, dilations=1, subpixel=True, use_activator=True, *kwargs ):
self.in_ch = in_ch
self.out_ch = out_ch
self.kernel_size = kernel_size
self.dilations = dilations
self.subpixel = subpixel
self.use_activator = use_activator
super().__init__(*kwargs)
def on_build(self, *args, **kwargs ):
self.conv1 = nn.Conv2D( self.in_ch,
self.out_ch // (4 if self.subpixel else 1),
kernel_size=self.kernel_size,
strides=1 if self.subpixel else 2,
padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer )
def forward(self, x):
x = self.conv1(x)
if self.subpixel:
x = tf.nn.space_to_depth(x, 2)
if self.use_activator:
x = tf.nn.leaky_relu(x, 0.1)
return x
def get_out_ch(self):
return (self.out_ch // 4) * 4
class DownscaleBlock(nn.ModelBase):
def on_build(self, in_ch, ch, n_downscales, kernel_size, dilations=1, subpixel=True):
self.downs = []
last_ch = in_ch
for i in range(n_downscales):
cur_ch = ch*( min(2**i, 8) )
self.downs.append ( Downscale(last_ch, cur_ch, kernel_size=kernel_size, dilations=dilations, subpixel=subpixel) )
last_ch = self.downs[-1].get_out_ch()
def forward(self, inp):
x = inp
for down in self.downs:
x = down(x)
return x
class Upscale(nn.ModelBase):
def on_build(self, in_ch, out_ch, kernel_size=3 ):
self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
def forward(self, x):
x = self.conv1(x)
x = tf.nn.leaky_relu(x, 0.1)
x = tf.nn.depth_to_space(x, 2)
return x
class ResidualBlock(nn.ModelBase):
def on_build(self, ch, kernel_size=3 ):
self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
def forward(self, inp):
x = self.conv1(inp)
x = tf.nn.leaky_relu(x, 0.2)
x = self.conv2(x)
x = tf.nn.leaky_relu(inp + x, 0.2)
return x
class UpdownResidualBlock(nn.ModelBase):
def on_build(self, ch, inner_ch, kernel_size=3 ):
self.up = Upscale (ch, inner_ch, kernel_size=kernel_size)
self.res = ResidualBlock (inner_ch, kernel_size=kernel_size)
self.down = Downscale (inner_ch, ch, kernel_size=kernel_size, use_activator=False)
def forward(self, inp):
x = self.up(inp)
x = upx = self.res(x)
x = self.down(x)
x = x + inp
x = tf.nn.leaky_relu(x, 0.2)
return x, upx
class Encoder(nn.ModelBase):
def on_build(self, in_ch, e_ch, is_hd):
self.is_hd=is_hd
if self.is_hd:
self.down1 = DownscaleBlock(in_ch, e_ch*2, n_downscales=4, kernel_size=3, dilations=1)
self.down2 = DownscaleBlock(in_ch, e_ch*2, n_downscales=4, kernel_size=5, dilations=1)
self.down3 = DownscaleBlock(in_ch, e_ch//2, n_downscales=4, kernel_size=5, dilations=2)
self.down4 = DownscaleBlock(in_ch, e_ch//2, n_downscales=4, kernel_size=7, dilations=2)
else:
self.down1 = DownscaleBlock(in_ch, e_ch, n_downscales=4, kernel_size=5, dilations=1, subpixel=False)
def forward(self, inp):
if self.is_hd:
x = tf.concat([ nn.tf_flatten(self.down1(inp)),
nn.tf_flatten(self.down2(inp)),
nn.tf_flatten(self.down3(inp)),
nn.tf_flatten(self.down4(inp)) ], -1 )
else:
x = nn.tf_flatten(self.down1(inp))
return x
class Inter(nn.ModelBase):
def __init__(self, in_ch, lowest_dense_res, ae_ch, ae_out_ch, **kwargs):
self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch = in_ch, lowest_dense_res, ae_ch, ae_out_ch
super().__init__(**kwargs)
def on_build(self):
in_ch, lowest_dense_res, ae_ch, ae_out_ch = self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch
self.dense1 = nn.Dense( in_ch, ae_ch, kernel_initializer=tf.initializers.orthogonal )
self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch, kernel_initializer=tf.initializers.orthogonal )
self.upscale1 = Upscale(ae_out_ch, ae_out_ch)
def forward(self, inp):
x = self.dense1(inp)
x = self.dense2(x)
x = tf.reshape (x, (-1, lowest_dense_res, lowest_dense_res, self.ae_out_ch))
x = self.upscale1(x)
return x
def get_out_ch(self):
return self.ae_out_ch
class Decoder(nn.ModelBase):
def on_build(self, in_ch, d_ch, d_mask_ch, is_hd ):
self.is_hd = is_hd
self.upscale0 = Upscale(in_ch, d_ch*8, kernel_size=3)
self.upscale1 = Upscale(d_ch*8, d_ch*4, kernel_size=3)
self.upscale2 = Upscale(d_ch*4, d_ch*2, kernel_size=3)
if is_hd:
self.res0 = UpdownResidualBlock(in_ch, d_ch*8, kernel_size=3)
self.res1 = UpdownResidualBlock(d_ch*8, d_ch*4, kernel_size=3)
self.res2 = UpdownResidualBlock(d_ch*4, d_ch*2, kernel_size=3)
self.res3 = UpdownResidualBlock(d_ch*2, d_ch, kernel_size=3)
else:
self.res0 = ResidualBlock(d_ch*8, kernel_size=3)
self.res1 = ResidualBlock(d_ch*4, kernel_size=3)
self.res2 = ResidualBlock(d_ch*2, kernel_size=3)
self.out_conv = nn.Conv2D( d_ch*2, 3, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3)
self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3)
self.upscalem2 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3)
self.out_convm = nn.Conv2D( d_mask_ch*2, 1, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
def get_weights_ex(self, include_mask):
# Call internal get_weights in order to initialize inner logic
self.get_weights()
weights = self.upscale0.get_weights() + self.upscale1.get_weights() + self.upscale2.get_weights() \
+ self.res0.get_weights() + self.res1.get_weights() + self.res2.get_weights() + self.out_conv.get_weights()
if include_mask:
weights += self.upscalem0.get_weights() + self.upscalem1.get_weights() + self.upscalem2.get_weights() \
+ self.out_convm.get_weights()
return weights
def forward(self, inp):
z = inp
if self.is_hd:
x, upx = self.res0(z)
x = self.upscale0(x)
x = tf.nn.leaky_relu(x + upx, 0.2)
x, upx = self.res1(x)
x = self.upscale1(x)
x = tf.nn.leaky_relu(x + upx, 0.2)
x, upx = self.res2(x)
x = self.upscale2(x)
x = tf.nn.leaky_relu(x + upx, 0.2)
x, upx = self.res3(x)
else:
x = self.upscale0(z)
x = self.res0(x)
x = self.upscale1(x)
x = self.res1(x)
x = self.upscale2(x)
x = self.res2(x)
m = self.upscalem0(z)
m = self.upscalem1(m)
m = self.upscalem2(m)
return tf.nn.sigmoid(self.out_conv(x)), \
tf.nn.sigmoid(self.out_convm(m))
class CodeDiscriminator(nn.ModelBase):
def on_build(self, in_ch, code_res, ch=256):
n_downscales = 2 + code_res // 8
self.convs = []
prev_ch = in_ch
for i in range(n_downscales):
cur_ch = ch * min( (2**i), 8 )
self.convs.append ( nn.Conv2D( prev_ch, cur_ch, kernel_size=4 if i == 0 else 3, strides=2, padding='SAME', kernel_initializer=conv_kernel_initializer) )
prev_ch = cur_ch
self.out_conv = nn.Conv2D( prev_ch, 1, kernel_size=1, padding='VALID', kernel_initializer=conv_kernel_initializer)
def forward(self, x):
for conv in self.convs:
x = tf.nn.leaky_relu( conv(x), 0.1 )
return self.out_conv(x)
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
resolution = self.options['resolution']
learn_mask = self.options['learn_mask']
archi = self.options['archi']
ae_dims = self.options['ae_dims']
e_dims = self.options['e_dims']
d_dims = self.options['d_dims']
d_mask_dims = self.options['d_mask_dims']
self.pretrain = self.options['pretrain']
masked_training = True
models_opt_on_gpu = False if len(devices) != 1 else self.options['models_opt_on_gpu']
models_opt_device = '/GPU:0' if models_opt_on_gpu and self.is_training else '/CPU:0'
optimizer_vars_on_cpu = models_opt_device=='/CPU:0'
input_nc = 3
output_nc = 3
bgr_shape = (resolution, resolution, output_nc)
mask_shape = (resolution, resolution, 1)
lowest_dense_res = resolution // 16
self.model_filename_list = []
with tf.device ('/CPU:0'):
#Place holders on CPU
self.warped_src = tf.placeholder (tf.float32, (None,)+bgr_shape)
self.warped_dst = tf.placeholder (tf.float32, (None,)+bgr_shape)
self.target_src = tf.placeholder (tf.float32, (None,)+bgr_shape)
self.target_dst = tf.placeholder (tf.float32, (None,)+bgr_shape)
self.target_srcm = tf.placeholder (tf.float32, (None,)+mask_shape)
self.target_dstm = tf.placeholder (tf.float32, (None,)+mask_shape)
self.target_dst_0 = tf.placeholder (tf.float32, (None,)+bgr_shape)
self.target_dst_1 = tf.placeholder (tf.float32, (None,)+bgr_shape)
self.target_dst_2 = tf.placeholder (tf.float32, (None,)+bgr_shape)
# Initializing model classes
with tf.device (models_opt_device):
if 'df' in archi:
self.encoder = Encoder(in_ch=input_nc, e_ch=e_dims, is_hd='hd' in archi, name='encoder')
encoder_out_ch = self.encoder.compute_output_shape ( (tf.float32, (None,resolution,resolution,input_nc)))[-1]
self.inter = Inter (in_ch=encoder_out_ch, lowest_dense_res=lowest_dense_res, ae_ch=ae_dims, ae_out_ch=ae_dims, name='inter')
inter_out_ch = self.inter.compute_output_shape ( (tf.float32, (None,encoder_out_ch)))[-1]
self.decoder_src = Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd='hd' in archi, name='decoder_src')
self.decoder_dst = Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd='hd' in archi, name='decoder_dst')
self.model_filename_list += [ [self.encoder, 'encoder.npy' ],
[self.inter, 'inter.npy' ],
[self.decoder_src, 'decoder_src.npy'],
[self.decoder_dst, 'decoder_dst.npy'] ]
if self.is_training:
if self.options['true_face_training']:
self.dis = CodeDiscriminator(ae_dims, code_res=lowest_dense_res*2, name='dis' )
self.model_filename_list += [ [self.dis, 'dis.npy'] ]
elif 'liae' in archi:
self.encoder = Encoder(in_ch=input_nc, e_ch=e_dims, is_hd='hd' in archi, name='encoder')
encoder_out_ch = self.encoder.compute_output_shape ( (tf.float32, (None,resolution,resolution,input_nc)))[-1]
self.inter_AB = Inter(in_ch=encoder_out_ch, lowest_dense_res=lowest_dense_res, ae_ch=ae_dims, ae_out_ch=ae_dims*2, name='inter_AB')
self.inter_B = Inter(in_ch=encoder_out_ch, lowest_dense_res=lowest_dense_res, ae_ch=ae_dims, ae_out_ch=ae_dims*2, name='inter_B')
inter_AB_out_ch = self.inter_AB.compute_output_shape ( (tf.float32, (None,encoder_out_ch)))[-1]
inter_B_out_ch = self.inter_B.compute_output_shape ( (tf.float32, (None,encoder_out_ch)))[-1]
inters_out_ch = inter_AB_out_ch+inter_B_out_ch
self.decoder = Decoder(in_ch=inters_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd='hd' in archi, name='decoder')
self.model_filename_list += [ [self.encoder, 'encoder.npy'],
[self.inter_AB, 'inter_AB.npy'],
[self.inter_B , 'inter_B.npy'],
[self.decoder , 'decoder.npy'] ]
if self.is_training:
# Initialize optimizers
lr=5e-5
lr_dropout = 0.3 if self.options['lr_dropout'] else 1.0
clipnorm = 1.0 if self.options['clipgrad'] else 0.0
self.src_dst_opt = nn.TFRMSpropOptimizer(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='src_dst_opt')
self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ]
if 'df' in archi:
self.src_dst_all_trainable_weights = self.encoder.get_weights() + self.inter.get_weights() + self.decoder_src.get_weights() + self.decoder_dst.get_weights()
self.src_dst_trainable_weights = self.encoder.get_weights() + self.inter.get_weights() + self.decoder_src.get_weights_ex(learn_mask) + self.decoder_dst.get_weights_ex(learn_mask)
self.src_trainable_weights = self.encoder.get_weights() + self.inter.get_weights() + self.decoder_src.get_weights_ex(False)
elif 'liae' in archi:
self.src_dst_all_trainable_weights = self.encoder.get_weights() + self.inter_AB.get_weights() + self.inter_B.get_weights() + self.decoder.get_weights()
self.src_dst_trainable_weights = self.encoder.get_weights() + self.inter_AB.get_weights() + self.inter_B.get_weights() + self.decoder.get_weights_ex(learn_mask)
self.src_dst_opt.initialize_variables (self.src_dst_all_trainable_weights, vars_on_cpu=optimizer_vars_on_cpu)
if self.options['true_face_training']:
self.D_opt = nn.TFRMSpropOptimizer(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='D_opt')
self.D_opt.initialize_variables ( self.dis.get_weights(), vars_on_cpu=optimizer_vars_on_cpu)
self.model_filename_list += [ (self.D_opt, 'D_opt.npy') ]
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices) )
bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
self.set_batch_size( gpu_count*bs_per_gpu)
# Compute losses per GPU
gpu_pred_src_src_list = []
gpu_pred_dst_dst_list = []
gpu_pred_src_dst_list = []
gpu_pred_src_srcm_list = []
gpu_pred_dst_dstm_list = []
gpu_pred_src_dstm_list = []
gpu_src_losses = []
gpu_dst_losses = []
gpu_src_dst_loss_gvs = []
gpu_D_loss_gvs = []
gpu_var_loss_gvs = []
for gpu_id in range(gpu_count):
with tf.device( f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0' ):
batch_slice = slice( gpu_id*bs_per_gpu, (gpu_id+1)*bs_per_gpu )
with tf.device(f'/CPU:0'):
# slice on CPU, otherwise all batch data will be transfered to GPU first
gpu_warped_src = self.warped_src [batch_slice,:,:,:]
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 = self.target_srcm[batch_slice,:,:,:]
gpu_target_dstm = self.target_dstm[batch_slice,:,:,:]
gpu_target_dst_0 = self.target_dst_0[batch_slice,:,:,:]
gpu_target_dst_1 = self.target_dst_1[batch_slice,:,:,:]
gpu_target_dst_2 = self.target_dst_2[batch_slice,:,:,:]
# process model tensors
if 'df' in archi:
gpu_src_code = self.inter(self.encoder(gpu_warped_src))
gpu_dst_code = self.inter(self.encoder(gpu_warped_dst))
gpu_dst_0_code = self.inter(self.encoder(gpu_target_dst_0))
gpu_dst_1_code = self.inter(self.encoder(gpu_target_dst_1))
gpu_dst_2_code = self.inter(self.encoder(gpu_target_dst_2))
gpu_pred_src_src, gpu_pred_src_srcm = self.decoder_src(gpu_src_code)
gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code)
gpu_pred_src_dst_0, _ = self.decoder_src(gpu_dst_0_code)
gpu_pred_src_dst_1, _ = self.decoder_src(gpu_dst_1_code)
gpu_pred_src_dst_2, _ = self.decoder_src(gpu_dst_2_code)
elif 'liae' in archi:
gpu_src_code = self.encoder (gpu_warped_src)
gpu_src_inter_AB_code = self.inter_AB (gpu_src_code)
gpu_src_code = tf.concat([gpu_src_inter_AB_code,gpu_src_inter_AB_code],-1)
gpu_dst_code = self.encoder (gpu_warped_dst)
gpu_dst_inter_B_code = self.inter_B (gpu_dst_code)
gpu_dst_inter_AB_code = self.inter_AB (gpu_dst_code)
gpu_dst_code = tf.concat([gpu_dst_inter_B_code,gpu_dst_inter_AB_code],-1)
gpu_src_dst_code = tf.concat([gpu_dst_inter_AB_code,gpu_dst_inter_AB_code],-1)
gpu_pred_src_src, gpu_pred_src_srcm = self.decoder(gpu_src_code)
gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder(gpu_dst_code)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
gpu_pred_src_src_list.append(gpu_pred_src_src)
gpu_pred_dst_dst_list.append(gpu_pred_dst_dst)
gpu_pred_src_dst_list.append(gpu_pred_src_dst)
gpu_pred_src_srcm_list.append(gpu_pred_src_srcm)
gpu_pred_dst_dstm_list.append(gpu_pred_dst_dstm)
gpu_pred_src_dstm_list.append(gpu_pred_src_dstm)
gpu_target_srcm_blur = nn.tf_gaussian_blur(gpu_target_srcm, max(1, resolution // 32) )
gpu_target_dstm_blur = nn.tf_gaussian_blur(gpu_target_dstm, max(1, resolution // 32) )
gpu_target_dst_masked = gpu_target_dst*gpu_target_dstm_blur
gpu_target_dst_anti_masked = gpu_target_dst*(1.0 - gpu_target_dstm_blur)
gpu_target_srcmasked_opt = gpu_target_src*gpu_target_srcm_blur if masked_training else gpu_target_src
gpu_target_dst_masked_opt = gpu_target_dst_masked if masked_training else gpu_target_dst
gpu_pred_src_src_masked_opt = gpu_pred_src_src*gpu_target_srcm_blur if masked_training else gpu_pred_src_src
gpu_pred_dst_dst_masked_opt = gpu_pred_dst_dst*gpu_target_dstm_blur if masked_training else gpu_pred_dst_dst
gpu_psd_target_dst_masked = gpu_pred_src_dst*gpu_target_dstm_blur
gpu_psd_target_dst_anti_masked = gpu_pred_src_dst*(1.0 - gpu_target_dstm_blur)
gpu_src_loss = tf.reduce_mean ( 10*nn.tf_dssim(gpu_target_srcmasked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_srcmasked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3])
if learn_mask:
gpu_src_loss += tf.reduce_mean ( tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] )
face_style_power = self.options['face_style_power'] / 100.0
if face_style_power != 0 and not self.pretrain:
gpu_src_loss += nn.tf_style_loss(gpu_psd_target_dst_masked, gpu_target_dst_masked, gaussian_blur_radius=resolution//16, loss_weight=10000*face_style_power)
bg_style_power = self.options['bg_style_power'] / 100.0
if bg_style_power != 0 and not self.pretrain:
gpu_src_loss += tf.reduce_mean( (10*bg_style_power)*nn.tf_dssim(gpu_psd_target_dst_anti_masked, gpu_target_dst_anti_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
gpu_src_loss += tf.reduce_mean( (10*bg_style_power)*tf.square( gpu_psd_target_dst_anti_masked - gpu_target_dst_anti_masked), axis=[1,2,3] )
gpu_dst_loss = tf.reduce_mean ( 10*nn.tf_dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3])
if learn_mask:
gpu_dst_loss += tf.reduce_mean ( tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] )
gpu_src_losses += [gpu_src_loss]
gpu_dst_losses += [gpu_dst_loss]
gpu_src_dst_loss = gpu_src_loss + gpu_dst_loss
if self.options['true_face_training']:
def DLoss(labels,logits):
return tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits), axis=[1,2,3])
gpu_src_code_d = self.dis( gpu_src_code )
gpu_src_code_d_ones = tf.ones_like(gpu_src_code_d)
gpu_src_code_d_zeros = tf.zeros_like(gpu_src_code_d)
gpu_dst_code_d = self.dis( gpu_dst_code )
gpu_dst_code_d_ones = tf.ones_like(gpu_dst_code_d)
gpu_src_dst_loss += 0.01*DLoss(gpu_src_code_d_ones, gpu_src_code_d)
gpu_D_loss = (DLoss(gpu_src_code_d_ones , gpu_dst_code_d) + \
DLoss(gpu_src_code_d_zeros, gpu_src_code_d) ) * 0.5
gpu_D_loss_gvs += [ nn.tf_gradients (gpu_D_loss, self.dis.get_weights() ) ]
gpu_src_dst_loss_gvs += [ nn.tf_gradients ( gpu_src_dst_loss, self.src_dst_trainable_weights ) ]
gpu_var_loss = nn.tf_style_loss (gpu_pred_src_dst_1, gpu_pred_src_dst_0, gaussian_blur_radius=resolution//16, loss_weight=100)
gpu_var_loss += nn.tf_style_loss (gpu_pred_src_dst_1, gpu_pred_src_dst_2, gaussian_blur_radius=resolution//16, loss_weight=100)
gpu_var_loss_gvs += [ nn.tf_gradients ( gpu_var_loss, self.src_trainable_weights ) ]
# Average losses and gradients, and create optimizer update ops
with tf.device (models_opt_device):
if gpu_count == 1:
pred_src_src = gpu_pred_src_src_list[0]
pred_dst_dst = gpu_pred_dst_dst_list[0]
pred_src_dst = gpu_pred_src_dst_list[0]
pred_src_srcm = gpu_pred_src_srcm_list[0]
pred_dst_dstm = gpu_pred_dst_dstm_list[0]
pred_src_dstm = gpu_pred_src_dstm_list[0]
src_loss = gpu_src_losses[0]
dst_loss = gpu_dst_losses[0]
src_dst_loss_gv = gpu_src_dst_loss_gvs[0]
var_loss_gv = gpu_var_loss_gvs[0]
else:
pred_src_src = tf.concat(gpu_pred_src_src_list, 0)
pred_dst_dst = tf.concat(gpu_pred_dst_dst_list, 0)
pred_src_dst = tf.concat(gpu_pred_src_dst_list, 0)
pred_src_srcm = tf.concat(gpu_pred_src_srcm_list, 0)
pred_dst_dstm = tf.concat(gpu_pred_dst_dstm_list, 0)
pred_src_dstm = tf.concat(gpu_pred_src_dstm_list, 0)
src_loss = nn.tf_average_tensor_list(gpu_src_losses)
dst_loss = nn.tf_average_tensor_list(gpu_dst_losses)
src_dst_loss_gv = nn.tf_average_gv_list (gpu_src_dst_loss_gvs)
var_loss_gv = nn.tf_average_gv_list (gpu_var_loss_gvs)
if self.options['true_face_training']:
D_loss_gv = nn.tf_average_gv_list(gpu_D_loss_gvs)
src_dst_loss_gv_op = self.src_dst_opt.get_update_op (src_dst_loss_gv )
if self.options['true_face_training']:
D_loss_gv_op = self.D_opt.get_update_op (D_loss_gv )
var_loss_gv_op = self.src_dst_opt.get_update_op (var_loss_gv )
# Initializing training and view functions
def src_dst_train(warped_src, target_src, target_srcm, \
warped_dst, target_dst, target_dstm):
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:target_srcm,
self.warped_dst :warped_dst,
self.target_dst :target_dst,
self.target_dstm:target_dstm,
})
s = np.mean(s)
d = np.mean(d)
return s, d
self.src_dst_train = src_dst_train
def var_train(target_dst_0, target_dst_1, target_dst_2):
_ = nn.tf_sess.run ( [ var_loss_gv_op],
feed_dict={self.target_dst_0 :target_dst_0,
self.target_dst_1 :target_dst_1,
self.target_dst_2 :target_dst_2,
})
self.var_train = var_train
if self.options['true_face_training']:
def D_train(warped_src, warped_dst):
nn.tf_sess.run ([D_loss_gv_op], feed_dict={self.warped_src: warped_src, self.warped_dst: warped_dst})
self.D_train = D_train
if learn_mask:
def AE_view(warped_src, warped_dst):
return nn.tf_sess.run ( [pred_src_src, pred_dst_dst, pred_dst_dstm, pred_src_dst, pred_src_dstm],
feed_dict={self.warped_src:warped_src,
self.warped_dst:warped_dst})
else:
def AE_view(warped_src, warped_dst):
return nn.tf_sess.run ( [pred_src_src, pred_dst_dst, pred_src_dst],
feed_dict={self.warped_src:warped_src,
self.warped_dst:warped_dst})
self.AE_view = AE_view
else:
# Initializing merge function
with tf.device( f'/GPU:0' if len(devices) != 0 else f'/CPU:0'):
if 'df' in archi:
gpu_dst_code = self.inter(self.encoder(self.warped_dst))
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code)
_, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
elif 'liae' in archi:
gpu_dst_code = self.encoder (self.warped_dst)
gpu_dst_inter_B_code = self.inter_B (gpu_dst_code)
gpu_dst_inter_AB_code = self.inter_AB (gpu_dst_code)
gpu_dst_code = tf.concat([gpu_dst_inter_B_code,gpu_dst_inter_AB_code],-1)
gpu_src_dst_code = tf.concat([gpu_dst_inter_AB_code,gpu_dst_inter_AB_code],-1)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
_, gpu_pred_dst_dstm = self.decoder(gpu_dst_code)
if learn_mask:
def AE_merge( warped_dst):
return nn.tf_sess.run ( [gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm], feed_dict={self.warped_dst:warped_dst})
else:
def AE_merge( warped_dst):
return nn.tf_sess.run ( [gpu_pred_src_dst], feed_dict={self.warped_dst:warped_dst})
self.AE_merge = AE_merge
# Loading/initializing all models/optimizers weights
for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"):
do_init = self.is_first_run()
if self.pretrain_just_disabled:
if 'df' in archi:
if model == self.inter:
do_init = True
elif 'liae' in archi:
if model == self.inter_AB:
do_init = True
if not do_init:
do_init = not model.load_weights( self.get_strpath_storage_for_file(filename) )
if do_init:
model.init_weights()
# initializing sample generators
if self.is_training:
t = SampleProcessor.Types
if self.options['face_type'] == 'h':
face_type = t.FACE_TYPE_HALF
elif self.options['face_type'] == 'mf':
face_type = t.FACE_TYPE_MID_FULL
elif self.options['face_type'] == 'f':
face_type = t.FACE_TYPE_FULL
training_data_src_path = self.training_data_src_path if not self.pretrain else self.get_pretraining_data_path()
training_data_dst_path = self.training_data_dst_path if not self.pretrain else self.get_pretraining_data_path()
random_ct_samples_path=training_data_dst_path if self.options['ct_mode'] != 'none' and not self.pretrain else None
t_img_warped = t.IMG_WARPED_TRANSFORMED if self.options['random_warp'] else t.IMG_TRANSFORMED
cpu_count = multiprocessing.cpu_count()
src_generators_count = cpu_count // 2
if self.options['ct_mode'] != 'none':
src_generators_count = int(src_generators_count * 1.5)
dst_generators_count = cpu_count - src_generators_count
self.set_training_data_generators ([
SampleGeneratorFace(training_data_src_path, random_ct_samples_path=random_ct_samples_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
output_sample_types = [ {'types' : (t_img_warped, face_type, t.MODE_BGR), 'resolution':resolution, 'ct_mode': self.options['ct_mode'] },
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution': resolution, 'ct_mode': self.options['ct_mode'] },
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_FACE_MASK_ALL_HULL), 'resolution': resolution } ],
generators_count=src_generators_count ),
SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
output_sample_types = [ {'types' : (t_img_warped, face_type, t.MODE_BGR), 'resolution':resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution': resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_FACE_MASK_ALL_HULL), 'resolution': resolution} ],
generators_count=dst_generators_count ),
SampleGeneratorFaceTemporal(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 = [{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution': resolution}],
generators_count=dst_generators_count )
])
def __init__(self, place_model_on_cpu=False, run_on_cpu=False):
nn.initialize(data_format="NHWC")
tf = nn.tf
class FaceEnhancer(nn.ModelBase):
def __init__(self, name='FaceEnhancer'):
super().__init__(name=name)
def on_build(self):
self.conv1 = nn.Conv2D(3,
64,
kernel_size=3,
strides=1,
padding='SAME')
self.dense1 = nn.Dense(1, 64, use_bias=False)
self.dense2 = nn.Dense(1, 64, use_bias=False)
self.e0_conv0 = nn.Conv2D(64,
64,
kernel_size=3,
strides=1,
padding='SAME')
self.e0_conv1 = nn.Conv2D(64,
64,
kernel_size=3,
strides=1,
padding='SAME')
self.e1_conv0 = nn.Conv2D(64,
112,
kernel_size=3,
strides=1,
padding='SAME')
self.e1_conv1 = nn.Conv2D(112,
112,
kernel_size=3,
strides=1,
padding='SAME')
self.e2_conv0 = nn.Conv2D(112,
192,
kernel_size=3,
strides=1,
padding='SAME')
self.e2_conv1 = nn.Conv2D(192,
192,
kernel_size=3,
strides=1,
padding='SAME')
self.e3_conv0 = nn.Conv2D(192,
336,
kernel_size=3,
strides=1,
padding='SAME')
self.e3_conv1 = nn.Conv2D(336,
336,
kernel_size=3,
strides=1,
padding='SAME')
self.e4_conv0 = nn.Conv2D(336,
512,
kernel_size=3,
strides=1,
padding='SAME')
self.e4_conv1 = nn.Conv2D(512,
512,
kernel_size=3,
strides=1,
padding='SAME')
self.center_conv0 = nn.Conv2D(512,
512,
kernel_size=3,
strides=1,
padding='SAME')
self.center_conv1 = nn.Conv2D(512,
512,
kernel_size=3,
strides=1,
padding='SAME')
self.center_conv2 = nn.Conv2D(512,
512,
kernel_size=3,
strides=1,
padding='SAME')
self.center_conv3 = nn.Conv2D(512,
512,
kernel_size=3,
strides=1,
padding='SAME')
self.d4_conv0 = nn.Conv2D(1024,
512,
kernel_size=3,
strides=1,
padding='SAME')
self.d4_conv1 = nn.Conv2D(512,
512,
kernel_size=3,
strides=1,
padding='SAME')
self.d3_conv0 = nn.Conv2D(848,
512,
kernel_size=3,
strides=1,
padding='SAME')
self.d3_conv1 = nn.Conv2D(512,
512,
kernel_size=3,
strides=1,
padding='SAME')
self.d2_conv0 = nn.Conv2D(704,
288,
kernel_size=3,
strides=1,
padding='SAME')
self.d2_conv1 = nn.Conv2D(288,
288,
kernel_size=3,
strides=1,
padding='SAME')
self.d1_conv0 = nn.Conv2D(400,
160,
kernel_size=3,
strides=1,
padding='SAME')
self.d1_conv1 = nn.Conv2D(160,
160,
kernel_size=3,
strides=1,
padding='SAME')
self.d0_conv0 = nn.Conv2D(224,
96,
kernel_size=3,
strides=1,
padding='SAME')
self.d0_conv1 = nn.Conv2D(96,
96,
kernel_size=3,
strides=1,
padding='SAME')
self.out1x_conv0 = nn.Conv2D(96,
48,
kernel_size=3,
strides=1,
padding='SAME')
self.out1x_conv1 = nn.Conv2D(48,
3,
kernel_size=3,
strides=1,
padding='SAME')
self.dec2x_conv0 = nn.Conv2D(96,
96,
kernel_size=3,
strides=1,
padding='SAME')
self.dec2x_conv1 = nn.Conv2D(96,
96,
kernel_size=3,
strides=1,
padding='SAME')
self.out2x_conv0 = nn.Conv2D(96,
48,
kernel_size=3,
strides=1,
padding='SAME')
self.out2x_conv1 = nn.Conv2D(48,
3,
kernel_size=3,
strides=1,
padding='SAME')
self.dec4x_conv0 = nn.Conv2D(96,
72,
kernel_size=3,
strides=1,
padding='SAME')
self.dec4x_conv1 = nn.Conv2D(72,
72,
kernel_size=3,
strides=1,
padding='SAME')
self.out4x_conv0 = nn.Conv2D(72,
36,
kernel_size=3,
strides=1,
padding='SAME')
self.out4x_conv1 = nn.Conv2D(36,
3,
kernel_size=3,
strides=1,
padding='SAME')
def forward(self, inp):
bgr, param, param1 = inp
x = self.conv1(bgr)
a = self.dense1(param)
a = tf.reshape(a, (-1, 1, 1, 64))
b = self.dense2(param1)
b = tf.reshape(b, (-1, 1, 1, 64))
x = tf.nn.leaky_relu(x + a + b, 0.1)
x = tf.nn.leaky_relu(self.e0_conv0(x), 0.1)
x = e0 = tf.nn.leaky_relu(self.e0_conv1(x), 0.1)
x = tf.nn.avg_pool(x, [1, 2, 2, 1], [1, 2, 2, 1], "VALID")
x = tf.nn.leaky_relu(self.e1_conv0(x), 0.1)
x = e1 = tf.nn.leaky_relu(self.e1_conv1(x), 0.1)
x = tf.nn.avg_pool(x, [1, 2, 2, 1], [1, 2, 2, 1], "VALID")
x = tf.nn.leaky_relu(self.e2_conv0(x), 0.1)
x = e2 = tf.nn.leaky_relu(self.e2_conv1(x), 0.1)
x = tf.nn.avg_pool(x, [1, 2, 2, 1], [1, 2, 2, 1], "VALID")
x = tf.nn.leaky_relu(self.e3_conv0(x), 0.1)
x = e3 = tf.nn.leaky_relu(self.e3_conv1(x), 0.1)
x = tf.nn.avg_pool(x, [1, 2, 2, 1], [1, 2, 2, 1], "VALID")
x = tf.nn.leaky_relu(self.e4_conv0(x), 0.1)
x = e4 = tf.nn.leaky_relu(self.e4_conv1(x), 0.1)
x = tf.nn.avg_pool(x, [1, 2, 2, 1], [1, 2, 2, 1], "VALID")
x = tf.nn.leaky_relu(self.center_conv0(x), 0.1)
x = tf.nn.leaky_relu(self.center_conv1(x), 0.1)
x = tf.nn.leaky_relu(self.center_conv2(x), 0.1)
x = tf.nn.leaky_relu(self.center_conv3(x), 0.1)
x = tf.concat([nn.resize2d_bilinear(x), e4], -1)
x = tf.nn.leaky_relu(self.d4_conv0(x), 0.1)
x = tf.nn.leaky_relu(self.d4_conv1(x), 0.1)
x = tf.concat([nn.resize2d_bilinear(x), e3], -1)
x = tf.nn.leaky_relu(self.d3_conv0(x), 0.1)
x = tf.nn.leaky_relu(self.d3_conv1(x), 0.1)
x = tf.concat([nn.resize2d_bilinear(x), e2], -1)
x = tf.nn.leaky_relu(self.d2_conv0(x), 0.1)
x = tf.nn.leaky_relu(self.d2_conv1(x), 0.1)
x = tf.concat([nn.resize2d_bilinear(x), e1], -1)
x = tf.nn.leaky_relu(self.d1_conv0(x), 0.1)
x = tf.nn.leaky_relu(self.d1_conv1(x), 0.1)
x = tf.concat([nn.resize2d_bilinear(x), e0], -1)
x = tf.nn.leaky_relu(self.d0_conv0(x), 0.1)
x = d0 = tf.nn.leaky_relu(self.d0_conv1(x), 0.1)
x = tf.nn.leaky_relu(self.out1x_conv0(x), 0.1)
x = self.out1x_conv1(x)
out1x = bgr + tf.nn.tanh(x)
x = d0
x = tf.nn.leaky_relu(self.dec2x_conv0(x), 0.1)
x = tf.nn.leaky_relu(self.dec2x_conv1(x), 0.1)
x = d2x = nn.resize2d_bilinear(x)
x = tf.nn.leaky_relu(self.out2x_conv0(x), 0.1)
x = self.out2x_conv1(x)
out2x = nn.resize2d_bilinear(out1x) + tf.nn.tanh(x)
x = d2x
x = tf.nn.leaky_relu(self.dec4x_conv0(x), 0.1)
x = tf.nn.leaky_relu(self.dec4x_conv1(x), 0.1)
x = d4x = nn.resize2d_bilinear(x)
x = tf.nn.leaky_relu(self.out4x_conv0(x), 0.1)
x = self.out4x_conv1(x)
out4x = nn.resize2d_bilinear(out2x) + tf.nn.tanh(x)
return out4x
model_path = Path(__file__).parent / "FaceEnhancer.npy"
if not model_path.exists():
raise Exception("Unable to load FaceEnhancer.npy")
with tf.device('/CPU:0' if place_model_on_cpu else '/GPU:0'):
self.model = FaceEnhancer()
self.model.load_weights(model_path)
with tf.device('/CPU:0' if run_on_cpu else '/GPU:0'):
self.model.build_for_run([
(tf.float32, nn.get4Dshape(192, 192, 3)),
(tf.float32, (
None,
1,
)),
(tf.float32, (
None,
1,
)),
])
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
self.model_data_format = "NCHW" if len(devices) != 0 and not self.is_debug() else "NHWC"
nn.initialize(data_format=self.model_data_format)
tf = nn.tf
self.resolution = resolution = self.options['resolution']
self.face_type = {'h' : FaceType.HALF,
'mf' : FaceType.MID_FULL,
'f' : FaceType.FULL,
'wf' : FaceType.WHOLE_FACE,
'head' : FaceType.HEAD}[ self.options['face_type'] ]
eyes_prio = self.options['eyes_prio']
archi_split = self.options['archi'].split('-')
if len(archi_split) == 2:
archi_type, archi_opts = archi_split
elif len(archi_split) == 1:
archi_type, archi_opts = archi_split[0], None
ae_dims = self.options['ae_dims']
e_dims = self.options['e_dims']
d_dims = self.options['d_dims']
d_mask_dims = self.options['d_mask_dims']
self.pretrain = self.options['pretrain']
if self.pretrain_just_disabled:
self.set_iter(0)
self.gan_power = gan_power = 0.0 if self.pretrain else self.options['gan_power']
random_warp = False if self.pretrain else self.options['random_warp']
if self.pretrain:
self.options_show_override['gan_power'] = 0.0
self.options_show_override['random_warp'] = False
self.options_show_override['lr_dropout'] = 'n'
self.options_show_override['face_style_power'] = 0.0
self.options_show_override['bg_style_power'] = 0.0
self.options_show_override['uniform_yaw'] = True
masked_training = self.options['masked_training']
import dfl
dfl.load_config()
masked_training = dfl.get_config("masked_training", "1") == "1"
ct_mode = self.options['ct_mode']
if ct_mode == 'none':
ct_mode = None
models_opt_on_gpu = False if len(devices) == 0 else self.options['models_opt_on_gpu']
models_opt_device = '/GPU:0' if models_opt_on_gpu and self.is_training else '/CPU:0'
optimizer_vars_on_cpu = models_opt_device=='/CPU:0'
input_ch=3
bgr_shape = nn.get4Dshape(resolution,resolution,input_ch)
mask_shape = nn.get4Dshape(resolution,resolution,1)
self.model_filename_list = []
with tf.device ('/CPU:0'):
#Place holders on CPU
self.warped_src = tf.placeholder (nn.floatx, bgr_shape)
self.warped_dst = tf.placeholder (nn.floatx, bgr_shape)
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)
# Initializing model classes
model_archi = nn.DeepFakeArchi(resolution, opts=archi_opts)
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))
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)))
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')
self.model_filename_list += [ [self.encoder, 'encoder.npy' ],
[self.inter, 'inter.npy' ],
[self.decoder_src, 'decoder_src.npy'],
[self.decoder_dst, 'decoder_dst.npy'] ]
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.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))
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
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'],
[self.inter_AB, 'inter_AB.npy'],
[self.inter_B , 'inter_B.npy'],
[self.decoder , 'decoder.npy'] ]
if self.is_training:
if gan_power != 0:
self.D_src = nn.UNetPatchDiscriminator(patch_size=resolution//16, in_ch=input_ch, name="D_src")
self.model_filename_list += [ [self.D_src, 'D_src_v2.npy'] ]
# Initialize optimizers
lr=5e-5
lr_dropout = 0.3 if self.options['lr_dropout'] in ['y','cpu'] and not self.pretrain else 1.0
clipnorm = 1.0 if self.options['clipgrad'] else 0.0
if 'df' in archi_type:
self.src_dst_trainable_weights = self.encoder.get_weights() + self.inter.get_weights() + self.decoder_src.get_weights() + self.decoder_dst.get_weights()
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.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.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.initialize_variables ( self.D_src.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')#+self.D_src_x2.get_weights()
self.model_filename_list += [ (self.D_src_dst_opt, 'D_src_v2_opt.npy') ]
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices) )
bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
self.set_batch_size( gpu_count*bs_per_gpu)
# Compute losses per GPU
gpu_pred_src_src_list = []
gpu_pred_dst_dst_list = []
gpu_pred_src_dst_list = []
gpu_pred_src_srcm_list = []
gpu_pred_dst_dstm_list = []
gpu_pred_src_dstm_list = []
gpu_src_losses = []
gpu_dst_losses = []
gpu_G_loss_gvs = []
gpu_D_code_loss_gvs = []
gpu_D_src_dst_loss_gvs = []
for gpu_id in range(gpu_count):
with tf.device( f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0' ):
with tf.device(f'/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_warped_src = self.warped_src [batch_slice,:,:,:]
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,:,:,:]
# process model tensors
if 'df' in archi_type:
gpu_src_code = self.inter(self.encoder(gpu_warped_src))
gpu_dst_code = self.inter(self.encoder(gpu_warped_dst))
gpu_pred_src_src, gpu_pred_src_srcm = self.decoder_src(gpu_src_code)
gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code)
elif 'liae' in archi_type:
gpu_src_code = self.encoder (gpu_warped_src)
gpu_src_inter_AB_code = self.inter_AB (gpu_src_code)
gpu_src_code = tf.concat([gpu_src_inter_AB_code,gpu_src_inter_AB_code], nn.conv2d_ch_axis )
gpu_dst_code = self.encoder (gpu_warped_dst)
gpu_dst_inter_B_code = self.inter_B (gpu_dst_code)
gpu_dst_inter_AB_code = self.inter_AB (gpu_dst_code)
gpu_dst_code = tf.concat([gpu_dst_inter_B_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis )
gpu_src_dst_code = tf.concat([gpu_dst_inter_AB_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis )
gpu_pred_src_src, gpu_pred_src_srcm = self.decoder(gpu_src_code)
gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder(gpu_dst_code)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
gpu_pred_src_src_list.append(gpu_pred_src_src)
gpu_pred_dst_dst_list.append(gpu_pred_dst_dst)
gpu_pred_src_dst_list.append(gpu_pred_src_dst)
gpu_pred_src_srcm_list.append(gpu_pred_src_srcm)
gpu_pred_dst_dstm_list.append(gpu_pred_dst_dstm)
gpu_pred_src_dstm_list.append(gpu_pred_src_dstm)
# unpack masks from one combined mask
gpu_target_srcm = tf.clip_by_value (gpu_target_srcm_all, 0, 1)
gpu_target_dstm = tf.clip_by_value (gpu_target_dstm_all, 0, 1)
gpu_target_srcm_eyes = tf.clip_by_value (gpu_target_srcm_all-1, 0, 1)
gpu_target_dstm_eyes = tf.clip_by_value (gpu_target_dstm_all-1, 0, 1)
gpu_target_srcm_blur = nn.gaussian_blur(gpu_target_srcm, max(1, resolution // 32) )
gpu_target_srcm_blur = tf.clip_by_value(gpu_target_srcm_blur, 0, 0.5) * 2
gpu_target_dstm_blur = nn.gaussian_blur(gpu_target_dstm, max(1, resolution // 32) )
gpu_target_dstm_style_blur = gpu_target_dstm_blur #default style mask is 0.5 on boundary
gpu_target_dstm_blur = tf.clip_by_value(gpu_target_dstm_blur, 0, 0.5) * 2
gpu_target_dst_masked = gpu_target_dst*gpu_target_dstm_blur
gpu_target_dst_style_masked = gpu_target_dst*gpu_target_dstm_style_blur
gpu_target_dst_style_anti_masked = gpu_target_dst*(1.0 - gpu_target_dstm_style_blur)
gpu_target_src_masked_opt = gpu_target_src*gpu_target_srcm_blur if masked_training else gpu_target_src
gpu_target_dst_masked_opt = gpu_target_dst_masked if masked_training else gpu_target_dst
gpu_pred_src_src_masked_opt = gpu_pred_src_src*gpu_target_srcm_blur if masked_training else gpu_pred_src_src
gpu_pred_dst_dst_masked_opt = gpu_pred_dst_dst*gpu_target_dstm_blur if masked_training else gpu_pred_dst_dst
gpu_psd_target_dst_style_masked = gpu_pred_src_dst*gpu_target_dstm_style_blur
gpu_psd_target_dst_style_anti_masked = gpu_pred_src_dst*(1.0 - gpu_target_dstm_style_blur)
if resolution < 256:
gpu_src_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
else:
gpu_src_loss = tf.reduce_mean ( 5*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
gpu_src_loss += tf.reduce_mean ( 5*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/23.2)), axis=[1])
gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_src_masked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3])
if eyes_prio:
gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_srcm_eyes - gpu_pred_src_src*gpu_target_srcm_eyes ), axis=[1,2,3])
gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] )
face_style_power = self.options['face_style_power'] / 100.0
if face_style_power != 0 and not self.pretrain:
gpu_src_loss += nn.style_loss(gpu_psd_target_dst_style_masked, gpu_target_dst_style_masked, gaussian_blur_radius=resolution//16, loss_weight=10000*face_style_power)
bg_style_power = self.options['bg_style_power'] / 100.0
if bg_style_power != 0 and not self.pretrain:
gpu_src_loss += tf.reduce_mean( (10*bg_style_power)*nn.dssim( gpu_psd_target_dst_style_anti_masked, gpu_target_dst_style_anti_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
gpu_src_loss += tf.reduce_mean( (10*bg_style_power)*tf.square(gpu_psd_target_dst_style_anti_masked - gpu_target_dst_style_anti_masked), axis=[1,2,3] )
if resolution < 256:
gpu_dst_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
else:
gpu_dst_loss = tf.reduce_mean ( 5*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
gpu_dst_loss += tf.reduce_mean ( 5*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/23.2) ), axis=[1])
gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3])
if eyes_prio:
gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_dstm_eyes - gpu_pred_dst_dst*gpu_target_dstm_eyes ), axis=[1,2,3])
gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] )
gpu_src_losses += [gpu_src_loss]
gpu_dst_losses += [gpu_dst_loss]
gpu_G_loss = gpu_src_loss + gpu_dst_loss
def DLoss(labels,logits):
return tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits), axis=[1,2,3])
if self.options['true_face_power'] != 0:
gpu_src_code_d = self.code_discriminator( gpu_src_code )
gpu_src_code_d_ones = tf.ones_like (gpu_src_code_d)
gpu_src_code_d_zeros = tf.zeros_like(gpu_src_code_d)
gpu_dst_code_d = self.code_discriminator( gpu_dst_code )
gpu_dst_code_d_ones = tf.ones_like(gpu_dst_code_d)
gpu_G_loss += self.options['true_face_power']*DLoss(gpu_src_code_d_ones, gpu_src_code_d)
gpu_D_code_loss = (DLoss(gpu_src_code_d_ones , gpu_dst_code_d) + \
DLoss(gpu_src_code_d_zeros, gpu_src_code_d) ) * 0.5
gpu_D_code_loss_gvs += [ nn.gradients (gpu_D_code_loss, self.code_discriminator.get_weights() ) ]
if gan_power != 0:
gpu_pred_src_src_d, \
gpu_pred_src_src_d2 = self.D_src(gpu_pred_src_src_masked_opt)
gpu_pred_src_src_d_ones = tf.ones_like (gpu_pred_src_src_d)
gpu_pred_src_src_d_zeros = tf.zeros_like(gpu_pred_src_src_d)
gpu_pred_src_src_d2_ones = tf.ones_like (gpu_pred_src_src_d2)
gpu_pred_src_src_d2_zeros = tf.zeros_like(gpu_pred_src_src_d2)
gpu_target_src_d, \
gpu_target_src_d2 = self.D_src(gpu_target_src_masked_opt)
gpu_target_src_d_ones = tf.ones_like(gpu_target_src_d)
gpu_target_src_d2_ones = tf.ones_like(gpu_target_src_d2)
gpu_D_src_dst_loss = (DLoss(gpu_target_src_d_ones , gpu_target_src_d) + \
DLoss(gpu_pred_src_src_d_zeros , gpu_pred_src_src_d) ) * 0.5 + \
(DLoss(gpu_target_src_d2_ones , gpu_target_src_d2) + \
DLoss(gpu_pred_src_src_d2_zeros , gpu_pred_src_src_d2) ) * 0.5
gpu_D_src_dst_loss_gvs += [ nn.gradients (gpu_D_src_dst_loss, self.D_src.get_weights() ) ]#+self.D_src_x2.get_weights()
gpu_G_loss += gan_power*(DLoss(gpu_pred_src_src_d_ones, gpu_pred_src_src_d) + \
DLoss(gpu_pred_src_src_d2_ones, gpu_pred_src_src_d2))
gpu_G_loss_gvs += [ nn.gradients ( gpu_G_loss, self.src_dst_trainable_weights ) ]
# Average losses and gradients, and create optimizer update ops
with tf.device (models_opt_device):
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)
pred_src_srcm = nn.concat(gpu_pred_src_srcm_list, 0)
pred_dst_dstm = nn.concat(gpu_pred_dst_dstm_list, 0)
pred_src_dstm = nn.concat(gpu_pred_src_dstm_list, 0)
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))
if self.options['true_face_power'] != 0:
D_loss_gv_op = self.D_code_opt.get_update_op (nn.average_gv_list(gpu_D_code_loss_gvs))
if gan_power != 0:
src_D_src_dst_loss_gv_op = self.D_src_dst_opt.get_update_op (nn.average_gv_list(gpu_D_src_dst_loss_gvs) )
# Initializing training and view functions
def src_dst_train(warped_src, target_src, target_srcm_all, \
warped_dst, target_dst, target_dstm_all):
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.warped_dst :warped_dst,
self.target_dst :target_dst,
self.target_dstm_all:target_dstm_all,
})
return s, d
self.src_dst_train = src_dst_train
if self.options['true_face_power'] != 0:
def D_train(warped_src, warped_dst):
nn.tf_sess.run ([D_loss_gv_op], feed_dict={self.warped_src: warped_src, self.warped_dst: warped_dst})
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):
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.warped_dst :warped_dst,
self.target_dst :target_dst,
self.target_dstm_all:target_dstm_all})
self.D_src_dst_train = D_src_dst_train
def AE_view(warped_src, warped_dst):
return nn.tf_sess.run ( [pred_src_src, pred_dst_dst, pred_dst_dstm, pred_src_dst, pred_src_dstm],
feed_dict={self.warped_src:warped_src,
self.warped_dst:warped_dst})
self.AE_view = AE_view
else:
# Initializing merge function
with tf.device( f'/GPU:0' if len(devices) != 0 else f'/CPU:0'):
if 'df' in archi_type:
gpu_dst_code = self.inter(self.encoder(self.warped_dst))
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code)
_, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
elif 'liae' in archi_type:
gpu_dst_code = self.encoder (self.warped_dst)
gpu_dst_inter_B_code = self.inter_B (gpu_dst_code)
gpu_dst_inter_AB_code = self.inter_AB (gpu_dst_code)
gpu_dst_code = tf.concat([gpu_dst_inter_B_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis)
gpu_src_dst_code = tf.concat([gpu_dst_inter_AB_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
_, gpu_pred_dst_dstm = self.decoder(gpu_dst_code)
def AE_merge( warped_dst):
return nn.tf_sess.run ( [gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm], feed_dict={self.warped_dst:warped_dst})
self.AE_merge = AE_merge
# Loading/initializing all models/optimizers weights
for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"):
if self.pretrain_just_disabled:
do_init = False
if 'df' in archi_type:
if model == self.inter:
do_init = True
elif 'liae' in archi_type:
if model == self.inter_AB or model == self.inter_B:
do_init = True
else:
do_init = self.is_first_run()
if not do_init:
do_init = not model.load_weights( self.get_strpath_storage_for_file(filename) )
if do_init:
model.init_weights()
# initializing sample generators
if self.is_training:
training_data_src_path = self.training_data_src_path if not self.pretrain else self.get_pretraining_data_path()
training_data_dst_path = self.training_data_dst_path if not self.pretrain else self.get_pretraining_data_path()
random_ct_samples_path=training_data_dst_path if ct_mode is not None and not self.pretrain else None
cpu_count = min(multiprocessing.cpu_count(), 8)
src_generators_count = cpu_count // 2
dst_generators_count = cpu_count // 2
if ct_mode is not None:
src_generators_count = int(src_generators_count * 1.5)
self.set_training_data_generators ([
SampleGeneratorFace(training_data_src_path, random_ct_samples_path=random_ct_samples_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
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_EYES, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
],
uniform_yaw_distribution=self.options['uniform_yaw'] or self.pretrain,
generators_count=src_generators_count ),
SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
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, '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, '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,
generators_count=dst_generators_count )
])
self.last_src_samples_loss = []
self.last_dst_samples_loss = []
if self.pretrain_just_disabled:
self.update_sample_for_preview(force_new=True)
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
self.model_data_format = "NCHW" if len(devices) != 0 and not self.is_debug() else "NHWC"
nn.initialize(data_format=self.model_data_format)
tf = nn.tf
self.resolution = resolution = self.options['resolution']
self.face_type = {'h' : FaceType.HALF,
'mf' : FaceType.MID_FULL,
'f' : FaceType.FULL,
'wf' : FaceType.WHOLE_FACE}[ self.options['face_type'] ]
eyes_prio = self.options['eyes_prio']
archi = self.options['archi']
is_hd = 'hd' in archi
ae_dims = self.options['ae_dims']
e_dims = self.options['e_dims']
d_dims = self.options['d_dims']
d_mask_dims = self.options['d_mask_dims']
self.pretrain = self.options['pretrain']
if self.pretrain_just_disabled:
self.set_iter(0)
self.gan_power = gan_power = self.options['gan_power'] if not self.pretrain else 0.0
masked_training = self.options['masked_training']
ct_mode = self.options['ct_mode']
if ct_mode == 'none':
ct_mode = None
models_opt_on_gpu = False if len(devices) == 0 else self.options['models_opt_on_gpu']
models_opt_device = '/GPU:0' if models_opt_on_gpu and self.is_training else '/CPU:0'
optimizer_vars_on_cpu = models_opt_device=='/CPU:0'
input_ch=3
bgr_shape = nn.get4Dshape(resolution,resolution,input_ch)
mask_shape = nn.get4Dshape(resolution,resolution,1)
self.model_filename_list = []
with tf.device ('/CPU:0'):
#Place holders on CPU
self.warped_src = tf.placeholder (nn.floatx, bgr_shape)
self.warped_dst = tf.placeholder (nn.floatx, bgr_shape)
self.src_code_in = tf.placeholder (nn.floatx, (None,256) )
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)
# Initializing model classes
model_archi = nn.DeepFakeArchi(resolution, mod='uhd' if 'uhd' in archi else None)
with tf.device (models_opt_device):
if 'df' in archi:
self.encoder = model_archi.Encoder(in_ch=input_ch, e_ch=e_dims, is_hd=is_hd, name='encoder')
encoder_out_ch = self.encoder.compute_output_channels ( (nn.floatx, bgr_shape))
self.inter = model_archi.Inter (in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims, is_hd=is_hd, name='inter')
inter_out_ch = self.inter.compute_output_channels ( (nn.floatx, (None,encoder_out_ch)))
self.decoder_src = model_archi.Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd=is_hd, name='decoder_src')
self.decoder_dst = model_archi.Decoder(in_ch=inter_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd=is_hd, name='decoder_dst')
self.model_filename_list += [ [self.encoder, 'encoder.npy' ],
[self.inter, 'inter.npy' ],
[self.decoder_src, 'decoder_src.npy'],
[self.decoder_dst, 'decoder_dst.npy'] ]
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.model_filename_list += [ [self.code_discriminator, 'code_discriminator.npy'] ]
elif 'liae' in archi:
self.encoder = model_archi.Encoder(in_ch=input_ch, e_ch=e_dims, is_hd=is_hd, name='encoder')
encoder_out_ch = self.encoder.compute_output_channels ( (nn.floatx, bgr_shape))
self.inter_AB = model_archi.Inter(in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims*2, is_hd=is_hd, name='inter_AB')
self.inter_B = model_archi.Inter(in_ch=encoder_out_ch, ae_ch=ae_dims, ae_out_ch=ae_dims*2, is_hd=is_hd, 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
self.decoder = model_archi.Decoder(in_ch=inters_out_ch, d_ch=d_dims, d_mask_ch=d_mask_dims, is_hd=is_hd, name='decoder')
self.model_filename_list += [ [self.encoder, 'encoder.npy'],
[self.inter_AB, 'inter_AB.npy'],
[self.inter_B , 'inter_B.npy'],
[self.decoder , 'decoder.npy'] ]
if self.is_training:
if gan_power != 0:
self.D_src = nn.PatchDiscriminator(patch_size=resolution//16, in_ch=input_ch, name="D_src")
self.D_dst = nn.PatchDiscriminator(patch_size=resolution//16, in_ch=input_ch, name="D_dst")
self.model_filename_list += [ [self.D_src, 'D_src.npy'] ]
self.model_filename_list += [ [self.D_dst, 'D_dst.npy'] ]
# Initialize optimizers
lr=5e-5
lr_dropout = 0.3 if self.options['lr_dropout'] and not self.pretrain else 1.0
clipnorm = 1.0 if self.options['clipgrad'] else 0.0
self.src_dst_opt = nn.RMSprop(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='src_dst_opt')
self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ]
if 'df' in archi:
self.src_dst_trainable_weights = self.encoder.get_weights() + self.inter.get_weights() + self.decoder_src.get_weights() + self.decoder_dst.get_weights()
elif 'liae' in archi:
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.initialize_variables (self.src_dst_trainable_weights, vars_on_cpu=optimizer_vars_on_cpu)
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.initialize_variables ( self.code_discriminator.get_weights(), vars_on_cpu=optimizer_vars_on_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.initialize_variables ( self.D_src.get_weights()+self.D_dst.get_weights(), vars_on_cpu=optimizer_vars_on_cpu)
self.model_filename_list += [ (self.D_src_dst_opt, 'D_src_dst_opt.npy') ]
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices) )
bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
self.set_batch_size( gpu_count*bs_per_gpu)
# Compute losses per GPU
gpu_src_latent_code_list = []
gpu_pred_src_src_list = []
gpu_pred_dst_dst_list = []
gpu_pred_src_dst_list = []
gpu_pred_src_srcm_list = []
gpu_pred_dst_dstm_list = []
gpu_pred_src_dstm_list = []
gpu_src_losses = []
gpu_dst_losses = []
gpu_G_loss_gvs = []
gpu_D_code_loss_gvs = []
gpu_D_src_dst_loss_gvs = []
for gpu_id in range(gpu_count):
with tf.device( f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0' ):
with tf.device(f'/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_warped_src = self.warped_src [batch_slice,:,:,:]
gpu_src_code_in = self.src_code_in[batch_slice,:]
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,:,:,:]
# process model tensors
if 'df' in archi:
gpu_src_latent_code = self.inter.dense1(self.encoder(gpu_warped_src))
gpu_src_in_code = self.inter.fd(gpu_src_code_in)
gpu_src_code = self.inter(self.encoder(gpu_warped_src))
gpu_dst_code = self.inter(self.encoder(gpu_warped_dst))
gpu_pred_src_src, gpu_pred_src_srcm = self.decoder_src(gpu_src_in_code)
#gpu_pred_src_src, gpu_pred_src_srcm = self.decoder_src(gpu_src_code)
gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code)
elif 'liae' in archi:
gpu_src_code = self.encoder (gpu_warped_src)
gpu_src_inter_AB_code = self.inter_AB (gpu_src_code)
gpu_src_code = tf.concat([gpu_src_inter_AB_code,gpu_src_inter_AB_code], nn.conv2d_ch_axis )
gpu_dst_code = self.encoder (gpu_warped_dst)
gpu_dst_inter_B_code = self.inter_B (gpu_dst_code)
gpu_dst_inter_AB_code = self.inter_AB (gpu_dst_code)
gpu_dst_code = tf.concat([gpu_dst_inter_B_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis )
gpu_src_dst_code = tf.concat([gpu_dst_inter_AB_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis )
gpu_pred_src_src, gpu_pred_src_srcm = self.decoder(gpu_src_code)
gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder(gpu_dst_code)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
gpu_src_latent_code_list.append(gpu_src_latent_code)
gpu_pred_src_src_list.append(gpu_pred_src_src)
gpu_pred_dst_dst_list.append(gpu_pred_dst_dst)
gpu_pred_src_dst_list.append(gpu_pred_src_dst)
gpu_pred_src_srcm_list.append(gpu_pred_src_srcm)
gpu_pred_dst_dstm_list.append(gpu_pred_dst_dstm)
gpu_pred_src_dstm_list.append(gpu_pred_src_dstm)
# unpack masks from one combined mask
gpu_target_srcm = tf.clip_by_value (gpu_target_srcm_all, 0, 1)
gpu_target_dstm = tf.clip_by_value (gpu_target_dstm_all, 0, 1)
gpu_target_srcm_eyes = tf.clip_by_value (gpu_target_srcm_all-1, 0, 1)
gpu_target_dstm_eyes = tf.clip_by_value (gpu_target_dstm_all-1, 0, 1)
gpu_target_srcm_blur = nn.gaussian_blur(gpu_target_srcm, max(1, resolution // 32) )
gpu_target_dstm_blur = nn.gaussian_blur(gpu_target_dstm, max(1, resolution // 32) )
gpu_target_dst_masked = gpu_target_dst*gpu_target_dstm_blur
gpu_target_dst_anti_masked = gpu_target_dst*(1.0 - gpu_target_dstm_blur)
gpu_target_src_masked_opt = gpu_target_src*gpu_target_srcm_blur if masked_training else gpu_target_src
gpu_target_dst_masked_opt = gpu_target_dst_masked if masked_training else gpu_target_dst
gpu_pred_src_src_masked_opt = gpu_pred_src_src*gpu_target_srcm_blur if masked_training else gpu_pred_src_src
gpu_pred_dst_dst_masked_opt = gpu_pred_dst_dst*gpu_target_dstm_blur if masked_training else gpu_pred_dst_dst
gpu_psd_target_dst_masked = gpu_pred_src_dst*gpu_target_dstm_blur
gpu_psd_target_dst_anti_masked = gpu_pred_src_dst*(1.0 - gpu_target_dstm_blur)
gpu_src_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_src_masked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3])
if eyes_prio:
gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_srcm_eyes - gpu_pred_src_src*gpu_target_srcm_eyes ), axis=[1,2,3])
gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] )
face_style_power = self.options['face_style_power'] / 100.0
if face_style_power != 0 and not self.pretrain:
gpu_src_loss += nn.style_loss(gpu_psd_target_dst_masked, gpu_target_dst_masked, gaussian_blur_radius=resolution//16, loss_weight=10000*face_style_power)
bg_style_power = self.options['bg_style_power'] / 100.0
if bg_style_power != 0 and not self.pretrain:
gpu_src_loss += tf.reduce_mean( (10*bg_style_power)*nn.dssim(gpu_psd_target_dst_anti_masked, gpu_target_dst_anti_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
gpu_src_loss += tf.reduce_mean( (10*bg_style_power)*tf.square( gpu_psd_target_dst_anti_masked - gpu_target_dst_anti_masked), axis=[1,2,3] )
gpu_dst_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3])
if eyes_prio:
gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_dstm_eyes - gpu_pred_dst_dst*gpu_target_dstm_eyes ), axis=[1,2,3])
gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] )
gpu_src_losses += [gpu_src_loss]
gpu_dst_losses += [gpu_dst_loss]
gpu_G_loss = gpu_src_loss + gpu_dst_loss
def DLoss(labels,logits):
return tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=labels, logits=logits), axis=[1,2,3])
if self.options['true_face_power'] != 0:
gpu_src_code_d = self.code_discriminator( gpu_src_code )
gpu_src_code_d_ones = tf.ones_like (gpu_src_code_d)
gpu_src_code_d_zeros = tf.zeros_like(gpu_src_code_d)
gpu_dst_code_d = self.code_discriminator( gpu_dst_code )
gpu_dst_code_d_ones = tf.ones_like(gpu_dst_code_d)
gpu_G_loss += self.options['true_face_power']*DLoss(gpu_src_code_d_ones, gpu_src_code_d)
gpu_D_code_loss = (DLoss(gpu_src_code_d_ones , gpu_dst_code_d) + \
DLoss(gpu_src_code_d_zeros, gpu_src_code_d) ) * 0.5
gpu_D_code_loss_gvs += [ nn.gradients (gpu_D_code_loss, self.code_discriminator.get_weights() ) ]
if gan_power != 0:
gpu_pred_src_src_d = self.D_src(gpu_pred_src_src_masked_opt)
gpu_pred_src_src_d_ones = tf.ones_like (gpu_pred_src_src_d)
gpu_pred_src_src_d_zeros = tf.zeros_like(gpu_pred_src_src_d)
gpu_target_src_d = self.D_src(gpu_target_src_masked_opt)
gpu_target_src_d_ones = tf.ones_like(gpu_target_src_d)
gpu_pred_dst_dst_d = self.D_dst(gpu_pred_dst_dst_masked_opt)
gpu_pred_dst_dst_d_ones = tf.ones_like (gpu_pred_dst_dst_d)
gpu_pred_dst_dst_d_zeros = tf.zeros_like(gpu_pred_dst_dst_d)
gpu_target_dst_d = self.D_dst(gpu_target_dst_masked_opt)
gpu_target_dst_d_ones = tf.ones_like(gpu_target_dst_d)
gpu_D_src_dst_loss = (DLoss(gpu_target_src_d_ones , gpu_target_src_d) + \
DLoss(gpu_pred_src_src_d_zeros, gpu_pred_src_src_d) ) * 0.5 + \
(DLoss(gpu_target_dst_d_ones , gpu_target_dst_d) + \
DLoss(gpu_pred_dst_dst_d_zeros, gpu_pred_dst_dst_d) ) * 0.5
gpu_D_src_dst_loss_gvs += [ nn.gradients (gpu_D_src_dst_loss, self.D_src.get_weights()+self.D_dst.get_weights() ) ]
gpu_G_loss += gan_power*(DLoss(gpu_pred_src_src_d_ones, gpu_pred_src_src_d) + DLoss(gpu_pred_dst_dst_d_ones, gpu_pred_dst_dst_d))
gpu_G_loss_gvs += [ nn.gradients ( gpu_G_loss, self.src_dst_trainable_weights ) ]
# Average losses and gradients, and create optimizer update ops
with tf.device (models_opt_device):
src_latent_code = nn.concat(gpu_src_latent_code_list, 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)
pred_src_srcm = nn.concat(gpu_pred_src_srcm_list, 0)
pred_dst_dstm = nn.concat(gpu_pred_dst_dstm_list, 0)
pred_src_dstm = nn.concat(gpu_pred_src_dstm_list, 0)
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))
if self.options['true_face_power'] != 0:
D_loss_gv_op = self.D_code_opt.get_update_op (nn.average_gv_list(gpu_D_code_loss_gvs))
if gan_power != 0:
src_D_src_dst_loss_gv_op = self.D_src_dst_opt.get_update_op (nn.average_gv_list(gpu_D_src_dst_loss_gvs) )
# Initializing training and view functions
def src_dst_train(warped_src, target_src, target_srcm_all, \
warped_dst, target_dst, target_dstm_all):
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.warped_dst :warped_dst,
self.target_dst :target_dst,
self.target_dstm_all:target_dstm_all,
})
return s, d
self.src_dst_train = src_dst_train
if self.options['true_face_power'] != 0:
def D_train(warped_src, warped_dst):
nn.tf_sess.run ([D_loss_gv_op], feed_dict={self.warped_src: warped_src, self.warped_dst: warped_dst})
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):
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.warped_dst :warped_dst,
self.target_dst :target_dst,
self.target_dstm_all:target_dstm_all})
self.D_src_dst_train = D_src_dst_train
def AE_get_latent(warped_src):
return nn.tf_sess.run ( src_latent_code, feed_dict={self.warped_src:warped_src})
self.AE_get_latent = AE_get_latent
def AE_view_src(warped_src, src_code_in):
return nn.tf_sess.run ( pred_src_src,
feed_dict={self.warped_src:warped_src, self.src_code_in:src_code_in })
self.AE_view_src = AE_view_src
def AE_view(warped_src, warped_dst):
return nn.tf_sess.run ( [pred_src_src, pred_dst_dst, pred_dst_dstm, pred_src_dst, pred_src_dstm],
feed_dict={self.warped_src:warped_src,
self.warped_dst:warped_dst})
self.AE_view = AE_view
else:
# Initializing merge function
with tf.device( f'/GPU:0' if len(devices) != 0 else f'/CPU:0'):
if 'df' in archi:
gpu_dst_code = self.inter(self.encoder(self.warped_dst))
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code)
_, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
elif 'liae' in archi:
gpu_dst_code = self.encoder (self.warped_dst)
gpu_dst_inter_B_code = self.inter_B (gpu_dst_code)
gpu_dst_inter_AB_code = self.inter_AB (gpu_dst_code)
gpu_dst_code = tf.concat([gpu_dst_inter_B_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis)
gpu_src_dst_code = tf.concat([gpu_dst_inter_AB_code,gpu_dst_inter_AB_code], nn.conv2d_ch_axis)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
_, gpu_pred_dst_dstm = self.decoder(gpu_dst_code)
def AE_merge( warped_dst):
return nn.tf_sess.run ( [gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm], feed_dict={self.warped_dst:warped_dst})
self.AE_merge = AE_merge
# Loading/initializing all models/optimizers weights
for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"):
if self.pretrain_just_disabled:
do_init = False
if 'df' in archi:
if model == self.inter:
do_init = True
elif 'liae' in archi:
if model == self.inter_AB:
do_init = True
else:
do_init = self.is_first_run()
if not do_init:
do_init = not model.load_weights( self.get_strpath_storage_for_file(filename) )
if do_init:
model.init_weights()
# initializing sample generators
if self.is_training:
training_data_src_path = self.training_data_src_path if not self.pretrain else self.get_pretraining_data_path()
training_data_dst_path = self.training_data_dst_path if not self.pretrain else self.get_pretraining_data_path()
random_ct_samples_path=training_data_dst_path if ct_mode is not None and not self.pretrain else None
cpu_count = min(multiprocessing.cpu_count(), 8)
src_generators_count = cpu_count // 2
dst_generators_count = cpu_count // 2
if ct_mode is not None:
src_generators_count = int(src_generators_count * 1.5)
self.set_training_data_generators ([
SampleGeneratorFace(training_data_src_path, random_ct_samples_path=random_ct_samples_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':self.options['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':False, '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_EYES, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
{'sample_type': SampleProcessor.SampleType.PITCH_YAW_ROLL, 'resolution': resolution},
],
generators_count=src_generators_count ),
SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':self.options['random_warp'], 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, '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, '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},
],
generators_count=dst_generators_count )
])
self.last_src_samples_loss = []
self.last_dst_samples_loss = []
if self.pretrain_just_disabled:
self.update_sample_for_preview(force_new=True)
class PRD(nn.ModelBase):
def on_build(self, ae_ch):
self.dense1 = nn.Dense( ae_ch+1, 1024 )
self.dense2 = nn.Dense( 1024, 2048 )
self.dense3 = nn.Dense( 2048, 4096 )
self.dense4 = nn.Dense( 4096, 4096 )
self.dense5 = nn.Dense( 4096, ae_ch )
def forward(self, inp, yaw_in):
x = tf.concat( [inp, yaw_in], -1 )
x = self.dense1(x)
x = self.dense2(x)
x = self.dense3(x)
x = self.dense4(x)
x = self.dense5(x)
return x
with tf.device( f'/GPU:0'):
prd_model = PRD(256, name='PRD')
prd_model.init_weights()
prd_in = tf.placeholder (nn.floatx, (None,256) )
prd_targ = tf.placeholder (nn.floatx, (None,256) )
yaw_diff_in = tf.placeholder (nn.floatx, (None,1) )
prd_out = prd_model(prd_in, yaw_diff_in)
loss = tf.reduce_sum ( tf.abs (prd_out - prd_targ) )
loss_gvs = nn.gradients (loss, prd_model.get_weights() )
prd_opt = nn.RMSprop(lr=5e-6, lr_dropout=0.3, name='prd_opt')
prd_opt.initialize_variables(prd_model.get_weights())
prd_opt.init_weights()
loss_gv_op = prd_opt.get_update_op (loss_gvs)
s_gen, _ = self.get_training_data_generators()
bs = self.get_batch_size()
for n in range(1000):
warped_src, target_src, target_srcm_all, src_pyr = s_gen.generate_next()
sl = self.AE_get_latent(target_src)
prd_in_np = []
prd_targ_np = []
yaw_diff_in_np = []
for i in range(bs):
prd_in_np += [sl[i]]
j = i
while j == i:
j = np.random.randint(bs)
prd_targ_np += [ sl[j] ]
yaw_diff_in_np += [ np.float32( [ src_pyr[j][1]-src_pyr[i][1] ] ) ]
prd_loss, _ = nn.tf_sess.run([loss, loss_gv_op], feed_dict={prd_in:prd_in_np, prd_targ:prd_targ_np, yaw_diff_in:yaw_diff_in_np} )
print(f'{n} loss = {prd_loss}')
warped_src, target_src, target_srcm_all, src_pyr = s_gen.generate_next()
sl = self.AE_get_latent(target_src)
yaw_diff_in_np = np.float32( [ [-0.4] ] *bs )
new_sl = nn.tf_sess.run(prd_out, feed_dict={prd_in:sl, yaw_diff_in:yaw_diff_in_np} )
new_target_src = self.AE_view_src( target_src, new_sl )
target_src = np.clip( nn.to_data_format( target_src ,"NHWC", self.model_data_format), 0.0, 1.0)
new_target_src = np.clip( nn.to_data_format( new_target_src ,"NHWC", self.model_data_format), 0.0, 1.0)
for i in range(bs):
screen = np.concatenate ( (target_src[i], new_target_src[i]), 1 )
cv2.imshow("", (screen*255).astype(np.uint8) )
cv2.waitKey(0)
import code
code.interact(local=dict(globals(), **locals()))
def apply_xseg(input_path, model_path):
if not input_path.exists():
raise ValueError(f'{input_path} not found. Please ensure it exists.')
if not model_path.exists():
raise ValueError(f'{model_path} not found. Please ensure it exists.')
face_type = None
model_dat = model_path / 'XSeg_data.dat'
if model_dat.exists():
dat = pickle.loads( model_dat.read_bytes() )
dat_options = dat.get('options', None)
if dat_options is not None:
face_type = dat_options.get('face_type', None)
if face_type is None:
face_type = io.input_str ("XSeg model face type", 'same', ['h','mf','f','wf','head','same'], help_message="Specify face type of trained XSeg model. For example if XSeg model trained as WF, but faceset is HEAD, specify WF to apply xseg only on WF part of HEAD. Default is 'same'").lower()
if face_type == 'same':
face_type = None
if face_type is not None:
face_type = {'h' : FaceType.HALF,
'mf' : FaceType.MID_FULL,
'f' : FaceType.FULL,
'wf' : FaceType.WHOLE_FACE,
'head' : FaceType.HEAD}[face_type]
io.log_info(f'Applying trained XSeg model to {input_path.name}/ folder.')
device_config = nn.DeviceConfig.ask_choose_device(choose_only_one=True)
nn.initialize(device_config)
xseg = XSegNet(name='XSeg',
load_weights=True,
weights_file_root=model_path,
data_format=nn.data_format,
raise_on_no_model_files=True)
xseg_res = xseg.get_resolution()
images_paths = pathex.get_image_paths(input_path, return_Path_class=True)
for filepath in io.progress_bar_generator(images_paths, "Processing"):
dflimg = DFLIMG.load(filepath)
if dflimg is None or not dflimg.has_data():
io.log_info(f'{filepath} is not a DFLIMG')
continue
img = cv2_imread(filepath).astype(np.float32) / 255.0
h,w,c = img.shape
img_face_type = FaceType.fromString( dflimg.get_face_type() )
if face_type is not None and img_face_type != face_type:
lmrks = dflimg.get_source_landmarks()
fmat = LandmarksProcessor.get_transform_mat(lmrks, w, face_type)
imat = LandmarksProcessor.get_transform_mat(lmrks, w, img_face_type)
g_p = LandmarksProcessor.transform_points (np.float32([(0,0),(w,0),(0,w) ]), fmat, True)
g_p2 = LandmarksProcessor.transform_points (g_p, imat)
mat = cv2.getAffineTransform( g_p2, np.float32([(0,0),(w,0),(0,w) ]) )
img = cv2.warpAffine(img, mat, (w, w), cv2.INTER_LANCZOS4)
img = cv2.resize(img, (xseg_res, xseg_res), interpolation=cv2.INTER_LANCZOS4)
else:
if w != xseg_res:
img = cv2.resize( img, (xseg_res,xseg_res), interpolation=cv2.INTER_LANCZOS4 )
if len(img.shape) == 2:
img = img[...,None]
mask = xseg.extract(img)
if face_type is not None and img_face_type != face_type:
mask = cv2.resize(mask, (w, w), interpolation=cv2.INTER_LANCZOS4)
mask = cv2.warpAffine( mask, mat, (w,w), np.zeros( (h,w,c), dtype=np.float), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4)
mask = cv2.resize(mask, (xseg_res, xseg_res), interpolation=cv2.INTER_LANCZOS4)
mask[mask < 0.5]=0
mask[mask >= 0.5]=1
dflimg.set_xseg_mask(mask)
dflimg.save()
def __init__(self,
name,
resolution,
face_type_str,
load_weights=True,
weights_file_root=None,
training=False,
place_model_on_cpu=False):
nn.initialize(data_format="NHWC")
tf = nn.tf
class Ternaus(nn.ModelBase):
def on_build(self, in_ch, ch):
self.features_0 = nn.Conv2D(in_ch,
ch,
kernel_size=3,
padding='SAME')
self.blurpool_0 = nn.BlurPool(filt_size=3)
self.features_3 = nn.Conv2D(ch,
ch * 2,
kernel_size=3,
padding='SAME')
self.blurpool_3 = nn.BlurPool(filt_size=3)
self.features_6 = nn.Conv2D(ch * 2,
ch * 4,
kernel_size=3,
padding='SAME')
self.features_8 = nn.Conv2D(ch * 4,
ch * 4,
kernel_size=3,
padding='SAME')
self.blurpool_8 = nn.BlurPool(filt_size=3)
self.features_11 = nn.Conv2D(ch * 4,
ch * 8,
kernel_size=3,
padding='SAME')
self.features_13 = nn.Conv2D(ch * 8,
ch * 8,
kernel_size=3,
padding='SAME')
self.blurpool_13 = nn.BlurPool(filt_size=3)
self.features_16 = nn.Conv2D(ch * 8,
ch * 8,
kernel_size=3,
padding='SAME')
self.features_18 = nn.Conv2D(ch * 8,
ch * 8,
kernel_size=3,
padding='SAME')
self.blurpool_18 = nn.BlurPool(filt_size=3)
self.conv_center = nn.Conv2D(ch * 8,
ch * 8,
kernel_size=3,
padding='SAME')
self.conv1_up = nn.Conv2DTranspose(ch * 8,
ch * 4,
kernel_size=3,
padding='SAME')
self.conv1 = nn.Conv2D(ch * 12,
ch * 8,
kernel_size=3,
padding='SAME')
self.conv2_up = nn.Conv2DTranspose(ch * 8,
ch * 4,
kernel_size=3,
padding='SAME')
self.conv2 = nn.Conv2D(ch * 12,
ch * 8,
kernel_size=3,
padding='SAME')
self.conv3_up = nn.Conv2DTranspose(ch * 8,
ch * 2,
kernel_size=3,
padding='SAME')
self.conv3 = nn.Conv2D(ch * 6,
ch * 4,
kernel_size=3,
padding='SAME')
self.conv4_up = nn.Conv2DTranspose(ch * 4,
ch,
kernel_size=3,
padding='SAME')
self.conv4 = nn.Conv2D(ch * 3,
ch * 2,
kernel_size=3,
padding='SAME')
self.conv5_up = nn.Conv2DTranspose(ch * 2,
ch // 2,
kernel_size=3,
padding='SAME')
self.conv5 = nn.Conv2D(ch // 2 + ch,
ch,
kernel_size=3,
padding='SAME')
self.out_conv = nn.Conv2D(ch, 1, kernel_size=3, padding='SAME')
def forward(self, inp):
x, = inp
x = x0 = tf.nn.relu(self.features_0(x))
x = self.blurpool_0(x)
x = x1 = tf.nn.relu(self.features_3(x))
x = self.blurpool_3(x)
x = tf.nn.relu(self.features_6(x))
x = x2 = tf.nn.relu(self.features_8(x))
x = self.blurpool_8(x)
x = tf.nn.relu(self.features_11(x))
x = x3 = tf.nn.relu(self.features_13(x))
x = self.blurpool_13(x)
x = tf.nn.relu(self.features_16(x))
x = x4 = tf.nn.relu(self.features_18(x))
x = self.blurpool_18(x)
x = self.conv_center(x)
x = tf.nn.relu(self.conv1_up(x))
x = tf.concat([x, x4], -1)
x = tf.nn.relu(self.conv1(x))
x = tf.nn.relu(self.conv2_up(x))
x = tf.concat([x, x3], -1)
x = tf.nn.relu(self.conv2(x))
x = tf.nn.relu(self.conv3_up(x))
x = tf.concat([x, x2], -1)
x = tf.nn.relu(self.conv3(x))
x = tf.nn.relu(self.conv4_up(x))
x = tf.concat([x, x1], -1)
x = tf.nn.relu(self.conv4(x))
x = tf.nn.relu(self.conv5_up(x))
x = tf.concat([x, x0], -1)
x = tf.nn.relu(self.conv5(x))
x = tf.nn.sigmoid(self.out_conv(x))
return x
if weights_file_root is not None:
weights_file_root = Path(weights_file_root)
else:
weights_file_root = Path(__file__).parent
self.weights_path = weights_file_root / (
'%s_%d_%s.npy' % (name, resolution, face_type_str))
e = tf.device('/CPU:0') if place_model_on_cpu else None
if e is not None: e.__enter__()
self.net = Ternaus(3, 64, name='Ternaus')
if load_weights:
self.net.load_weights(self.weights_path)
else:
self.net.init_weights()
if e is not None: e.__exit__(None, None, None)
self.net.build_for_run([(tf.float32,
nn.get4Dshape(resolution, resolution, 3))])
if training:
raise Exception("training not supported yet")
"""
if training:
try:
with open( Path(__file__).parent / 'vgg11_enc_weights.npy', 'rb' ) as f:
d = pickle.loads (f.read())
for i in [0,3,6,8,11,13,16,18]:
s = 'features.%d' % i
self.model.get_layer (s).set_weights ( d[s] )
except:
io.log_err("Unable to load VGG11 pretrained weights from vgg11_enc_weights.npy")
conv_weights_list = []
for layer in self.model.layers:
if 'CA.' in layer.name:
conv_weights_list += [layer.weights[0]] #Conv2D kernel_weights
CAInitializerMP ( conv_weights_list )
"""
"""
def __init__(self, is_training=False,
saved_models_path=None,
training_data_src_path=None,
training_data_dst_path=None,
pretraining_data_path=None,
pretrained_model_path=None,
no_preview=False,
force_model_name=None,
force_gpu_idxs=None,
cpu_only=False,
debug=False,
force_model_class_name=None,
**kwargs):
self.is_training = is_training
self.saved_models_path = saved_models_path
self.training_data_src_path = training_data_src_path
self.training_data_dst_path = training_data_dst_path
self.pretraining_data_path = pretraining_data_path
self.pretrained_model_path = pretrained_model_path
self.no_preview = no_preview
self.debug = debug
self.model_class_name = model_class_name = Path(inspect.getmodule(self).__file__).parent.name.rsplit("_", 1)[1]
if force_model_class_name is None:
if force_model_name is not None:
self.model_name = force_model_name
else:
while True:
# gather all model dat files
saved_models_names = []
for filepath in pathex.get_file_paths(saved_models_path):
filepath_name = filepath.name
if filepath_name.endswith(f'{model_class_name}_data.dat'):
saved_models_names += [ (filepath_name.split('_')[0], os.path.getmtime(filepath)) ]
# sort by modified datetime
saved_models_names = sorted(saved_models_names, key=operator.itemgetter(1), reverse=True )
saved_models_names = [ x[0] for x in saved_models_names ]
if len(saved_models_names) != 0:
io.log_info ("Choose one of saved models, or enter a name to create a new model.")
io.log_info ("[r] : rename")
io.log_info ("[d] : delete")
io.log_info ("")
for i, model_name in enumerate(saved_models_names):
s = f"[{i}] : {model_name} "
if i == 0:
s += "- latest"
io.log_info (s)
inp = io.input_str(f"", "0", show_default_value=False )
model_idx = -1
try:
model_idx = np.clip ( int(inp), 0, len(saved_models_names)-1 )
except:
pass
if model_idx == -1:
if len(inp) == 1:
is_rename = inp[0] == 'r'
is_delete = inp[0] == 'd'
if is_rename or is_delete:
if len(saved_models_names) != 0:
if is_rename:
name = io.input_str(f"Enter the name of the model you want to rename")
elif is_delete:
name = io.input_str(f"Enter the name of the model you want to delete")
if name in saved_models_names:
if is_rename:
new_model_name = io.input_str(f"Enter new name of the model")
for filepath in pathex.get_paths(saved_models_path):
filepath_name = filepath.name
model_filename, remain_filename = filepath_name.split('_', 1)
if model_filename == name:
if is_rename:
new_filepath = filepath.parent / ( new_model_name + '_' + remain_filename )
filepath.rename (new_filepath)
elif is_delete:
filepath.unlink()
continue
self.model_name = inp
else:
self.model_name = saved_models_names[model_idx]
else:
self.model_name = io.input_str(f"No saved models found. Enter a name of a new model", "new")
self.model_name = self.model_name.replace('_', ' ')
break
self.model_name = self.model_name + '_' + self.model_class_name
else:
self.model_name = force_model_class_name
self.iter = 0
self.options = {}
self.loss_history = []
self.sample_for_preview = None
self.choosed_gpu_indexes = None
model_data = {}
self.model_data_path = Path( self.get_strpath_storage_for_file('data.dat') )
if self.model_data_path.exists():
io.log_info (f"Loading {self.model_name} model...")
model_data = pickle.loads ( self.model_data_path.read_bytes() )
self.iter = model_data.get('iter',0)
if self.iter != 0:
self.options = model_data['options']
self.loss_history = model_data.get('loss_history', [])
self.sample_for_preview = model_data.get('sample_for_preview', None)
self.choosed_gpu_indexes = model_data.get('choosed_gpu_indexes', None)
if self.is_first_run():
io.log_info ("\nModel first run.")
self.device_config = nn.DeviceConfig.GPUIndexes( force_gpu_idxs or nn.ask_choose_device_idxs(suggest_best_multi_gpu=True)) \
if not cpu_only else nn.DeviceConfig.CPU()
nn.initialize(self.device_config)
####
self.default_options_path = saved_models_path / f'{self.model_class_name}_default_options.dat'
self.default_options = {}
if self.default_options_path.exists():
try:
self.default_options = pickle.loads ( self.default_options_path.read_bytes() )
except:
pass
self.choose_preview_history = False
self.batch_size = self.load_or_def_option('batch_size', 1)
#####
io.input_skip_pending()
self.on_initialize_options()
if self.is_first_run():
# save as default options only for first run model initialize
self.default_options_path.write_bytes( pickle.dumps (self.options) )
self.autobackup_hour = self.options.get('autobackup_hour', 0)
self.write_preview_history = self.options.get('write_preview_history', False)
self.target_iter = self.options.get('target_iter',0)
self.random_flip = self.options.get('random_flip',True)
self.on_initialize()
self.options['batch_size'] = self.batch_size
if self.is_training:
self.preview_history_path = self.saved_models_path / ( f'{self.get_model_name()}_history' )
self.autobackups_path = self.saved_models_path / ( f'{self.get_model_name()}_autobackups' )
if self.write_preview_history or io.is_colab():
if not self.preview_history_path.exists():
self.preview_history_path.mkdir(exist_ok=True)
else:
if self.iter == 0:
for filename in pathex.get_image_paths(self.preview_history_path):
Path(filename).unlink()
if self.generator_list is None:
raise ValueError( 'You didnt set_training_data_generators()')
else:
for i, generator in enumerate(self.generator_list):
if not isinstance(generator, SampleGeneratorBase):
raise ValueError('training data generator is not subclass of SampleGeneratorBase')
self.update_sample_for_preview(choose_preview_history=self.choose_preview_history)
if self.autobackup_hour != 0:
self.autobackup_start_time = time.time()
if not self.autobackups_path.exists():
self.autobackups_path.mkdir(exist_ok=True)
io.log_info( self.get_summary_text() )
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
self.model_data_format = "NCHW" if len(
device_config.devices) != 0 and not self.is_debug() else "NHWC"
nn.initialize(data_format=self.model_data_format)
tf = nn.tf
conv_kernel_initializer = nn.initializers.ca()
class Downscale(nn.ModelBase):
def __init__(self,
in_ch,
out_ch,
kernel_size=5,
dilations=1,
subpixel=True,
use_activator=True,
*kwargs):
self.in_ch = in_ch
self.out_ch = out_ch
self.kernel_size = kernel_size
self.dilations = dilations
self.subpixel = subpixel
self.use_activator = use_activator
super().__init__(*kwargs)
def on_build(self, *args, **kwargs):
self.conv1 = nn.Conv2D(
self.in_ch,
self.out_ch // (4 if self.subpixel else 1),
kernel_size=self.kernel_size,
strides=1 if self.subpixel else 2,
padding='SAME',
dilations=self.dilations,
kernel_initializer=conv_kernel_initializer)
def forward(self, x):
x = self.conv1(x)
if self.subpixel:
x = nn.space_to_depth(x, 2)
if self.use_activator:
x = tf.nn.leaky_relu(x, 0.1)
return x
def get_out_ch(self):
return (self.out_ch // 4) * 4
class DownscaleBlock(nn.ModelBase):
def on_build(self,
in_ch,
ch,
n_downscales,
kernel_size,
dilations=1,
subpixel=True):
self.downs = []
last_ch = in_ch
for i in range(n_downscales):
cur_ch = ch * (min(2**i, 8))
self.downs.append(
Downscale(last_ch,
cur_ch,
kernel_size=kernel_size,
dilations=dilations,
subpixel=subpixel))
last_ch = self.downs[-1].get_out_ch()
def forward(self, inp):
x = inp
for down in self.downs:
x = down(x)
return x
class Upscale(nn.ModelBase):
def on_build(self, in_ch, out_ch, kernel_size=3):
self.conv1 = nn.Conv2D(
in_ch,
out_ch * 4,
kernel_size=kernel_size,
padding='SAME',
kernel_initializer=conv_kernel_initializer)
def forward(self, x):
x = self.conv1(x)
x = tf.nn.leaky_relu(x, 0.1)
x = nn.depth_to_space(x, 2)
return x
class ResidualBlock(nn.ModelBase):
def on_build(self, ch, kernel_size=3):
self.conv1 = nn.Conv2D(
ch,
ch,
kernel_size=kernel_size,
padding='SAME',
kernel_initializer=conv_kernel_initializer)
self.conv2 = nn.Conv2D(
ch,
ch,
kernel_size=kernel_size,
padding='SAME',
kernel_initializer=conv_kernel_initializer)
def forward(self, inp):
x = self.conv1(inp)
x = tf.nn.leaky_relu(x, 0.2)
x = self.conv2(x)
x = tf.nn.leaky_relu(inp + x, 0.2)
return x
class UpdownResidualBlock(nn.ModelBase):
def on_build(self, ch, inner_ch, kernel_size=3):
self.up = Upscale(ch, inner_ch, kernel_size=kernel_size)
self.res = ResidualBlock(inner_ch, kernel_size=kernel_size)
self.down = Downscale(inner_ch,
ch,
kernel_size=kernel_size,
use_activator=False)
def forward(self, inp):
x = self.up(inp)
x = upx = self.res(x)
x = self.down(x)
x = x + inp
x = tf.nn.leaky_relu(x, 0.2)
return x, upx
class Encoder(nn.ModelBase):
def on_build(self, in_ch, e_ch):
self.down1 = DownscaleBlock(in_ch,
e_ch,
n_downscales=5,
kernel_size=5,
dilations=1,
subpixel=False)
def forward(self, inp):
x = nn.flatten(self.down1(inp))
return x
class Inter(nn.ModelBase):
def __init__(self, in_ch, lowest_dense_res, ae_ch, ae_out_ch,
**kwargs):
self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch = in_ch, lowest_dense_res, ae_ch, ae_out_ch
super().__init__(**kwargs)
def on_build(self):
in_ch, lowest_dense_res, ae_ch, ae_out_ch = self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch
self.dense1 = nn.Dense(in_ch, ae_ch)
self.dense2 = nn.Dense(
ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch)
self.upscale1 = Upscale(ae_out_ch, ae_out_ch * 2)
def forward(self, inp):
x = self.dense1(inp)
x = self.dense2(x)
x = nn.reshape_4D(x, lowest_dense_res, lowest_dense_res,
self.ae_out_ch)
x = self.upscale1(x)
return x
def get_out_ch(self):
return self.ae_out_ch
class Decoder(nn.ModelBase):
def on_build(self, in_ch, d_ch, d_mask_ch):
self.upscale0 = Upscale(in_ch, d_ch * 8, kernel_size=3)
self.upscale1 = Upscale(d_ch * 8, d_ch * 4, kernel_size=3)
self.upscale2 = Upscale(d_ch * 4, d_ch * 2, kernel_size=3)
self.upscale3 = Upscale(d_ch * 2, d_ch * 1, kernel_size=3)
self.res0 = ResidualBlock(d_ch * 8, kernel_size=3)
self.res1 = ResidualBlock(d_ch * 4, kernel_size=3)
self.res2 = ResidualBlock(d_ch * 2, kernel_size=3)
self.res3 = ResidualBlock(d_ch * 1, kernel_size=3)
self.out_conv = nn.Conv2D(
d_ch * 1,
3,
kernel_size=1,
padding='SAME',
kernel_initializer=conv_kernel_initializer)
self.upscalem0 = Upscale(in_ch, d_mask_ch * 8, kernel_size=3)
self.upscalem1 = Upscale(d_mask_ch * 8,
d_mask_ch * 4,
kernel_size=3)
self.upscalem2 = Upscale(d_mask_ch * 4,
d_mask_ch * 2,
kernel_size=3)
self.upscalem3 = Upscale(d_mask_ch * 2,
d_mask_ch * 1,
kernel_size=3)
self.out_convm = nn.Conv2D(
d_mask_ch * 1,
1,
kernel_size=1,
padding='SAME',
kernel_initializer=conv_kernel_initializer)
"""
def get_weights_ex(self, include_mask):
# Call internal get_weights in order to initialize inner logic
self.get_weights()
weights = self.upscale0.get_weights() + self.upscale1.get_weights() + self.upscale2.get_weights() \
+ self.res0.get_weights() + self.res1.get_weights() + self.res2.get_weights() + self.out_conv.get_weights()
if include_mask:
weights += self.upscalem0.get_weights() + self.upscalem1.get_weights() + self.upscalem2.get_weights() \
+ self.out_convm.get_weights()
return weights
"""
def get_weights_ex(self, include_mask):
# Call internal get_weights in order to initialize inner logic
self.get_weights()
weights = self.upscale0.get_weights() + self.upscale1.get_weights() + self.upscale2.get_weights() + self.upscale3.get_weights()\
+ self.res0.get_weights() + self.res1.get_weights() + self.res2.get_weights() + self.res3.get_weights() + self.out_conv.get_weights()
if include_mask:
weights += self.upscalem0.get_weights() + self.upscalem1.get_weights() + self.upscalem2.get_weights() + self.upscalem3.get_weights() \
+ self.out_convm.get_weights()
return weights
def forward(self, inp):
z = inp
x = self.upscale0(z)
x = self.res0(x)
x = self.upscale1(x)
x = self.res1(x)
x = self.upscale2(x)
x = self.res2(x)
x = self.upscale3(x)
x = self.res3(x)
m = self.upscalem0(z)
m = self.upscalem1(m)
m = self.upscalem2(m)
m = self.upscalem3(m)
return tf.nn.sigmoid(self.out_conv(x)), \
tf.nn.sigmoid(self.out_convm(m))
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
self.resolution = resolution = 448
self.learn_mask = learn_mask = True
eyes_prio = True
ae_dims = self.options['ae_dims']
e_dims = self.options['e_dims']
d_dims = self.options['d_dims']
d_mask_dims = self.options['d_mask_dims']
self.pretrain = False
self.pretrain_just_disabled = False
if self.pretrain_just_disabled:
self.set_iter(0)
self.gan_power = gan_power = self.options[
'gan_power'] if not self.pretrain else 0.0
masked_training = True
models_opt_on_gpu = False if len(devices) == 0 else True if len(
devices) > 1 else self.options['models_opt_on_gpu']
models_opt_device = '/GPU:0' if models_opt_on_gpu and self.is_training else '/CPU:0'
optimizer_vars_on_cpu = models_opt_device == '/CPU:0'
input_ch = 3
output_ch = 3
bgr_shape = nn.get4Dshape(resolution, resolution, input_ch)
mask_shape = nn.get4Dshape(resolution, resolution, 1)
lowest_dense_res = resolution // 32
self.model_filename_list = []
with tf.device('/CPU:0'):
#Place holders on CPU
self.warped_src = tf.placeholder(nn.floatx, bgr_shape)
self.warped_dst = tf.placeholder(nn.floatx, bgr_shape)
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)
# Initializing model classes
with tf.device(models_opt_device):
self.encoder = Encoder(in_ch=input_ch, e_ch=e_dims, name='encoder')
encoder_out_ch = self.encoder.compute_output_channels(
(nn.floatx, bgr_shape))
self.inter = Inter(in_ch=encoder_out_ch,
lowest_dense_res=lowest_dense_res,
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)))
self.decoder_src = Decoder(in_ch=inter_out_ch,
d_ch=d_dims,
d_mask_ch=d_mask_dims,
name='decoder_src')
self.decoder_dst = Decoder(in_ch=inter_out_ch,
d_ch=d_dims,
d_mask_ch=d_mask_dims,
name='decoder_dst')
self.model_filename_list += [[self.encoder, 'encoder.npy'],
[self.inter, 'inter.npy'],
[self.decoder_src, 'decoder_src.npy'],
[self.decoder_dst, 'decoder_dst.npy']]
if self.is_training:
if gan_power != 0:
self.D_src = nn.PatchDiscriminator(patch_size=resolution //
16,
in_ch=output_ch,
base_ch=256,
name="D_src")
self.D_dst = nn.PatchDiscriminator(patch_size=resolution //
16,
in_ch=output_ch,
base_ch=256,
name="D_dst")
self.model_filename_list += [[self.D_src, 'D_src.npy']]
self.model_filename_list += [[self.D_dst, 'D_dst.npy']]
# Initialize optimizers
lr = 5e-5
clipnorm = 1.0 if self.options['clipgrad'] else 0.0
self.src_dst_opt = nn.RMSprop(lr=lr,
clipnorm=clipnorm,
name='src_dst_opt')
self.model_filename_list += [(self.src_dst_opt,
'src_dst_opt.npy')]
self.src_dst_all_trainable_weights = self.encoder.get_weights(
) + self.inter.get_weights() + self.decoder_src.get_weights(
) + self.decoder_dst.get_weights()
self.src_dst_trainable_weights = self.encoder.get_weights(
) + self.inter.get_weights() + self.decoder_src.get_weights_ex(
learn_mask) + self.decoder_dst.get_weights_ex(learn_mask)
self.src_dst_opt.initialize_variables(
self.src_dst_all_trainable_weights,
vars_on_cpu=optimizer_vars_on_cpu)
if gan_power != 0:
self.D_src_dst_opt = nn.RMSprop(lr=lr,
clipnorm=clipnorm,
name='D_src_dst_opt')
self.D_src_dst_opt.initialize_variables(
self.D_src.get_weights() + self.D_dst.get_weights(),
vars_on_cpu=optimizer_vars_on_cpu)
self.model_filename_list += [(self.D_src_dst_opt,
'D_src_dst_opt.npy')]
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices))
bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
self.set_batch_size(gpu_count * bs_per_gpu)
# Compute losses per GPU
gpu_pred_src_src_list = []
gpu_pred_dst_dst_list = []
gpu_pred_src_dst_list = []
gpu_pred_src_srcm_list = []
gpu_pred_dst_dstm_list = []
gpu_pred_src_dstm_list = []
gpu_src_losses = []
gpu_dst_losses = []
gpu_G_loss_gvs = []
gpu_D_code_loss_gvs = []
gpu_D_src_dst_loss_gvs = []
for gpu_id in range(gpu_count):
with tf.device(
f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0'):
with tf.device(f'/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_warped_src = self.warped_src[batch_slice, :, :, :]
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, :, :, :]
# process model tensors
gpu_src_code = self.inter(self.encoder(gpu_warped_src))
gpu_dst_code = self.inter(self.encoder(gpu_warped_dst))
gpu_pred_src_src, gpu_pred_src_srcm = self.decoder_src(
gpu_src_code)
gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder_dst(
gpu_dst_code)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(
gpu_dst_code)
gpu_pred_src_src_list.append(gpu_pred_src_src)
gpu_pred_dst_dst_list.append(gpu_pred_dst_dst)
gpu_pred_src_dst_list.append(gpu_pred_src_dst)
gpu_pred_src_srcm_list.append(gpu_pred_src_srcm)
gpu_pred_dst_dstm_list.append(gpu_pred_dst_dstm)
gpu_pred_src_dstm_list.append(gpu_pred_src_dstm)
# unpack masks from one combined mask
gpu_target_srcm = tf.clip_by_value(gpu_target_srcm_all, 0,
1)
gpu_target_dstm = tf.clip_by_value(gpu_target_dstm_all, 0,
1)
gpu_target_srcm_eyes = tf.clip_by_value(
gpu_target_srcm_all - 1, 0, 1)
gpu_target_dstm_eyes = tf.clip_by_value(
gpu_target_dstm_all - 1, 0, 1)
gpu_target_srcm_blur = nn.gaussian_blur(
gpu_target_srcm, max(1, resolution // 32))
gpu_target_dstm_blur = nn.gaussian_blur(
gpu_target_dstm, max(1, resolution // 32))
gpu_target_dst_masked = gpu_target_dst * gpu_target_dstm_blur
gpu_target_dst_anti_masked = gpu_target_dst * (
1.0 - gpu_target_dstm_blur)
gpu_target_src_masked_opt = gpu_target_src * gpu_target_srcm_blur if masked_training else gpu_target_src
gpu_target_dst_masked_opt = gpu_target_dst_masked if masked_training else gpu_target_dst
gpu_pred_src_src_masked_opt = gpu_pred_src_src * gpu_target_srcm_blur if masked_training else gpu_pred_src_src
gpu_pred_dst_dst_masked_opt = gpu_pred_dst_dst * gpu_target_dstm_blur if masked_training else gpu_pred_dst_dst
gpu_psd_target_dst_masked = gpu_pred_src_dst * gpu_target_dstm_blur
gpu_psd_target_dst_anti_masked = gpu_pred_src_dst * (
1.0 - gpu_target_dstm_blur)
gpu_src_loss = tf.reduce_mean(
10 * nn.dssim(gpu_target_src_masked_opt,
gpu_pred_src_src_masked_opt,
max_val=1.0,
filter_size=int(resolution / 11.6)),
axis=[1])
gpu_src_loss += tf.reduce_mean(
10 * tf.square(gpu_target_src_masked_opt -
gpu_pred_src_src_masked_opt),
axis=[1, 2, 3])
if eyes_prio:
gpu_src_loss += tf.reduce_mean(
300 *
tf.abs(gpu_target_src * gpu_target_srcm_eyes -
gpu_pred_src_src * gpu_target_srcm_eyes),
axis=[1, 2, 3])
if learn_mask:
gpu_src_loss += tf.reduce_mean(
10 *
tf.square(gpu_target_srcm - gpu_pred_src_srcm),
axis=[1, 2, 3])
gpu_dst_loss = tf.reduce_mean(
10 * nn.dssim(gpu_target_dst_masked_opt,
gpu_pred_dst_dst_masked_opt,
max_val=1.0,
filter_size=int(resolution / 11.6)),
axis=[1])
gpu_dst_loss += tf.reduce_mean(
10 * tf.square(gpu_target_dst_masked_opt -
gpu_pred_dst_dst_masked_opt),
axis=[1, 2, 3])
if eyes_prio:
gpu_dst_loss += tf.reduce_mean(
300 *
tf.abs(gpu_target_dst * gpu_target_dstm_eyes -
gpu_pred_dst_dst * gpu_target_dstm_eyes),
axis=[1, 2, 3])
if learn_mask:
gpu_dst_loss += tf.reduce_mean(
10 *
tf.square(gpu_target_dstm - gpu_pred_dst_dstm),
axis=[1, 2, 3])
gpu_src_losses += [gpu_src_loss]
gpu_dst_losses += [gpu_dst_loss]
gpu_G_loss = gpu_src_loss + gpu_dst_loss
def DLoss(labels, logits):
return tf.reduce_mean(
tf.nn.sigmoid_cross_entropy_with_logits(
labels=labels, logits=logits),
axis=[1, 2, 3])
if gan_power != 0:
gpu_pred_src_src_d = self.D_src(
gpu_pred_src_src_masked_opt)
gpu_pred_src_src_d_ones = tf.ones_like(
gpu_pred_src_src_d)
gpu_pred_src_src_d_zeros = tf.zeros_like(
gpu_pred_src_src_d)
gpu_target_src_d = self.D_src(
gpu_target_src_masked_opt)
gpu_target_src_d_ones = tf.ones_like(gpu_target_src_d)
gpu_pred_dst_dst_d = self.D_dst(
gpu_pred_dst_dst_masked_opt)
gpu_pred_dst_dst_d_ones = tf.ones_like(
gpu_pred_dst_dst_d)
gpu_pred_dst_dst_d_zeros = tf.zeros_like(
gpu_pred_dst_dst_d)
gpu_target_dst_d = self.D_dst(
gpu_target_dst_masked_opt)
gpu_target_dst_d_ones = tf.ones_like(gpu_target_dst_d)
gpu_D_src_dst_loss = (DLoss(gpu_target_src_d_ones , gpu_target_src_d) + \
DLoss(gpu_pred_src_src_d_zeros, gpu_pred_src_src_d) ) * 0.5 + \
(DLoss(gpu_target_dst_d_ones , gpu_target_dst_d) + \
DLoss(gpu_pred_dst_dst_d_zeros, gpu_pred_dst_dst_d) ) * 0.5
gpu_D_src_dst_loss_gvs += [
nn.gradients(
gpu_D_src_dst_loss,
self.D_src.get_weights() +
self.D_dst.get_weights())
]
gpu_G_loss += gan_power * (
DLoss(gpu_pred_src_src_d_ones,
gpu_pred_src_src_d) +
DLoss(gpu_pred_dst_dst_d_ones, gpu_pred_dst_dst_d))
gpu_G_loss_gvs += [
nn.gradients(gpu_G_loss,
self.src_dst_trainable_weights)
]
# Average losses and gradients, and create optimizer update ops
with tf.device(models_opt_device):
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)
pred_src_srcm = nn.concat(gpu_pred_src_srcm_list, 0)
pred_dst_dstm = nn.concat(gpu_pred_dst_dstm_list, 0)
pred_src_dstm = nn.concat(gpu_pred_src_dstm_list, 0)
src_loss = nn.average_tensor_list(gpu_src_losses)
dst_loss = nn.average_tensor_list(gpu_dst_losses)
src_dst_loss_gv_op = self.src_dst_opt.get_update_op(
nn.average_gv_list(gpu_G_loss_gvs))
if gan_power != 0:
src_D_src_dst_loss_gv_op = self.D_src_dst_opt.get_update_op(
nn.average_gv_list(gpu_D_src_dst_loss_gvs))
# Initializing training and view functions
def src_dst_train(warped_src, target_src, target_srcm_all, \
warped_dst, target_dst, target_dstm_all):
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.warped_dst: warped_dst,
self.target_dst: target_dst,
self.target_dstm_all: target_dstm_all,
})
s = np.mean(s)
d = np.mean(d)
return s, d
self.src_dst_train = src_dst_train
if gan_power != 0:
def D_src_dst_train(warped_src, target_src, target_srcm_all, \
warped_dst, target_dst, target_dstm_all):
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.warped_dst: warped_dst,
self.target_dst: target_dst,
self.target_dstm_all: target_dstm_all
})
self.D_src_dst_train = D_src_dst_train
if learn_mask:
def AE_view(warped_src, warped_dst):
return nn.tf_sess.run([
pred_src_src, pred_dst_dst, pred_dst_dstm,
pred_src_dst, pred_src_dstm
],
feed_dict={
self.warped_src: warped_src,
self.warped_dst: warped_dst
})
else:
def AE_view(warped_src, warped_dst):
return nn.tf_sess.run(
[pred_src_src, pred_dst_dst, pred_src_dst],
feed_dict={
self.warped_src: warped_src,
self.warped_dst: warped_dst
})
self.AE_view = AE_view
else:
# Initializing merge function
with tf.device(f'/GPU:0' if len(devices) != 0 else f'/CPU:0'):
gpu_dst_code = self.inter(self.encoder(self.warped_dst))
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(
gpu_dst_code)
_, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
if learn_mask:
def AE_merge(warped_dst):
return nn.tf_sess.run([
gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm
],
feed_dict={
self.warped_dst: warped_dst
})
else:
def AE_merge(warped_dst):
return nn.tf_sess.run(
[gpu_pred_src_dst],
feed_dict={self.warped_dst: warped_dst})
self.AE_merge = AE_merge
# Loading/initializing all models/optimizers weights
for model, filename in io.progress_bar_generator(
self.model_filename_list, "Initializing models"):
if self.pretrain_just_disabled:
do_init = False
if model == self.inter:
do_init = True
else:
do_init = self.is_first_run()
if not do_init:
do_init = not model.load_weights(
self.get_strpath_storage_for_file(filename))
if do_init:
model.init_weights()
# initializing sample generators
if self.is_training:
t = SampleProcessor.Types
face_type = t.FACE_TYPE_HEAD
training_data_src_path = self.training_data_src_path if not self.pretrain else self.get_pretraining_data_path(
)
training_data_dst_path = self.training_data_dst_path if not self.pretrain else self.get_pretraining_data_path(
)
t_img_warped = t.IMG_WARPED_TRANSFORMED
cpu_count = min(multiprocessing.cpu_count(), 8)
src_generators_count = cpu_count // 2
dst_generators_count = cpu_count // 2
self.set_training_data_generators([
SampleGeneratorFace(
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=[
{
'types': (t_img_warped, face_type, t.MODE_BGR),
'data_format': nn.data_format,
'resolution': resolution
},
{
'types':
(t.IMG_TRANSFORMED, face_type, t.MODE_BGR),
'data_format': nn.data_format,
'resolution': resolution
},
{
'types': (t.IMG_TRANSFORMED, face_type,
t.MODE_FACE_MASK_ALL_EYES_HULL),
'data_format':
nn.data_format,
'resolution':
resolution
},
],
generators_count=src_generators_count),
SampleGeneratorFace(
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=[
{
'types': (t_img_warped, face_type, t.MODE_BGR),
'data_format': nn.data_format,
'resolution': resolution
},
{
'types':
(t.IMG_TRANSFORMED, face_type, t.MODE_BGR),
'data_format': nn.data_format,
'resolution': resolution
},
{
'types': (t.IMG_TRANSFORMED, face_type,
t.MODE_FACE_MASK_ALL_EYES_HULL),
'data_format':
nn.data_format,
'resolution':
resolution
},
],
generators_count=dst_generators_count)
])
if self.pretrain_just_disabled:
self.update_sample_for_preview(force_new=True)
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
self.model_data_format = "NCHW" if len(
devices) != 0 and not self.is_debug() else "NHWC"
nn.initialize(data_format=self.model_data_format)
tf = nn.tf
class EncBlock(nn.ModelBase):
def on_build(self, in_ch, out_ch, level):
self.zero_level = level == 0
self.conv1 = nn.Conv2D(in_ch,
out_ch,
kernel_size=3,
padding='SAME')
self.conv2 = nn.Conv2D(
out_ch,
out_ch,
kernel_size=4 if self.zero_level else 3,
padding='VALID' if self.zero_level else 'SAME')
def forward(self, x):
x = tf.nn.leaky_relu(self.conv1(x), 0.2)
x = tf.nn.leaky_relu(self.conv2(x), 0.2)
if not self.zero_level:
x = nn.max_pool(x)
#if self.zero_level:
return x
class DecBlock(nn.ModelBase):
def on_build(self, in_ch, out_ch, level):
self.zero_level = level == 0
self.conv1 = nn.Conv2D(
in_ch,
out_ch,
kernel_size=4 if self.zero_level else 3,
padding=3 if self.zero_level else 'SAME')
self.conv2 = nn.Conv2D(out_ch,
out_ch,
kernel_size=3,
padding='SAME')
def forward(self, x):
if not self.zero_level:
x = nn.upsample2d(x)
x = tf.nn.leaky_relu(self.conv1(x), 0.2)
x = tf.nn.leaky_relu(self.conv2(x), 0.2)
return x
class InterBlock(nn.ModelBase):
def on_build(self, in_ch, out_ch, level):
self.zero_level = level == 0
self.dense1 = nn.Dense()
def forward(self, x):
x = tf.nn.leaky_relu(self.conv1(x), 0.2)
x = tf.nn.leaky_relu(self.conv2(x), 0.2)
if not self.zero_level:
x = nn.max_pool(x)
#if self.zero_level:
return x
class FromRGB(nn.ModelBase):
def on_build(self, out_ch):
self.conv1 = nn.Conv2D(3,
out_ch,
kernel_size=1,
padding='SAME')
def forward(self, x):
return tf.nn.leaky_relu(self.conv1(x), 0.2)
class ToRGB(nn.ModelBase):
def on_build(self, in_ch):
self.conv = nn.Conv2D(in_ch, 3, kernel_size=1, padding='SAME')
self.convm = nn.Conv2D(in_ch, 1, kernel_size=1, padding='SAME')
def forward(self, x):
return tf.nn.sigmoid(self.conv(x)), tf.nn.sigmoid(
self.convm(x))
ed_dims = 16
ae_res = 4
level_chs = {
i - 1: v
for i, v in enumerate([
np.clip(ed_dims * (2**i), 0, 512)
for i in range(self.stage_max + 2)
][::-1])
}
ae_ch = level_chs[0]
class Encoder(nn.ModelBase):
def on_build(self, e_ch, levels):
self.enc_blocks = {}
self.from_rgbs = {}
self.dense_norm = nn.DenseNorm()
for level in range(levels, -1, -1):
self.from_rgbs[level] = FromRGB(level_chs[level])
if level != 0:
self.enc_blocks[level] = EncBlock(
level_chs[level], level_chs[level - 1], level)
self.ae_dense1 = nn.Dense(ae_res * ae_res * ae_ch, 256)
self.ae_dense2 = nn.Dense(256, ae_res * ae_res * ae_ch)
def forward(self, stage, inp, prev_inp=None, alpha=None):
x = inp
for level in range(stage, -1, -1):
if stage in self.from_rgbs:
if level == stage:
x = self.from_rgbs[level](x)
elif level == stage - 1:
x = x * alpha + self.from_rgbs[level](prev_inp) * (
1 - alpha)
if level != 0:
x = self.enc_blocks[level](x)
x = nn.flatten(x)
x = self.dense_norm(x)
x = self.ae_dense1(x)
x = self.ae_dense2(x)
x = nn.reshape_4D(x, ae_res, ae_res, ae_ch)
return x
def get_stage_weights(self, stage):
self.get_weights()
weights = []
for level in range(stage, -1, -1):
if stage in self.from_rgbs:
if level == stage or level == stage - 1:
weights.append(self.from_rgbs[level].get_weights())
if level != 0:
weights.append(
self.enc_blocks[level].get_weights())
weights.append(self.ae_dense1.get_weights())
weights.append(self.ae_dense2.get_weights())
if len(weights) == 0:
return []
elif len(weights) == 1:
return weights[0]
else:
return sum(weights[1:], weights[0])
class Decoder(nn.ModelBase):
def on_build(self, levels_range):
self.dec_blocks = {}
self.to_rgbs = {}
for level in range(levels_range[0], levels_range[1] + 1):
self.to_rgbs[level] = ToRGB(level_chs[level])
if level != 0:
self.dec_blocks[level] = DecBlock(
level_chs[level - 1], level_chs[level], level)
def forward(self, stage, inp, alpha=None, inter=None):
x = inp
for level in range(stage + 1):
if level in self.to_rgbs:
if level == stage and stage > 0:
prev_level = level - 1
#prev_x, prev_xm = (inter.to_rgbs[prev_level] if inter is not None and prev_level in inter.to_rgbs else self.to_rgbs[prev_level])(x)
prev_x, prev_xm = self.to_rgbs[prev_level](x)
prev_x = nn.upsample2d(prev_x)
prev_xm = nn.upsample2d(prev_xm)
if level != 0:
x = self.dec_blocks[level](x)
if level == stage:
x, xm = self.to_rgbs[level](x)
if stage > 0:
x = x * alpha + prev_x * (1 - alpha)
xm = xm * alpha + prev_xm * (1 - alpha)
return x, xm
return x
def get_stage_weights(self, stage):
# Call internal get_weights in order to initialize inner logic
self.get_weights()
weights = []
for level in range(stage + 1):
if level in self.to_rgbs:
if level != 0:
weights.append(
self.dec_blocks[level].get_weights())
if level == stage or level == stage - 1:
weights.append(self.to_rgbs[level].get_weights())
if len(weights) == 0:
return []
elif len(weights) == 1:
return weights[0]
else:
return sum(weights[1:], weights[0])
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
self.stage = stage = self.options['stage']
self.start_stage_iter = self.options.get('start_stage_iter', 0)
self.target_stage_iter = self.options.get('target_stage_iter', 0)
resolution = self.options['resolution']
stage_resolutions = [2**(i + 2) for i in range(self.stage_max + 1)]
stage_resolution = stage_resolutions[stage]
prev_stage = stage - 1 if stage != 0 else stage
prev_stage_resolution = stage_resolutions[
stage - 1] if stage != 0 else stage_resolution
self.pretrain = False
self.pretrain_just_disabled = False
masked_training = True
models_opt_on_gpu = len(devices) == 1 and devices[0].total_mem_gb >= 4
models_opt_device = '/GPU:0' if models_opt_on_gpu and self.is_training else '/CPU:0'
optimizer_vars_on_cpu = models_opt_device == '/CPU:0'
input_nc = 3
output_nc = 3
prev_bgr_shape = nn.get4Dshape(prev_stage_resolution,
prev_stage_resolution, output_nc)
bgr_shape = nn.get4Dshape(stage_resolution, stage_resolution,
output_nc)
mask_shape = nn.get4Dshape(stage_resolution, stage_resolution, 1)
self.model_filename_list = []
with tf.device('/CPU:0'):
#Place holders on CPU
self.prev_warped_src = tf.placeholder(tf.float32, prev_bgr_shape)
self.warped_src = tf.placeholder(tf.float32, bgr_shape)
self.prev_warped_dst = tf.placeholder(tf.float32, prev_bgr_shape)
self.warped_dst = tf.placeholder(tf.float32, bgr_shape)
self.target_src = tf.placeholder(tf.float32, bgr_shape)
self.target_dst = tf.placeholder(tf.float32, bgr_shape)
self.target_srcm = tf.placeholder(tf.float32, mask_shape)
self.target_dstm = tf.placeholder(tf.float32, mask_shape)
self.alpha_t = tf.placeholder(tf.float32, (None, 1, 1, 1))
# Initializing model classes
with tf.device(models_opt_device):
self.encoder = Encoder(e_ch=ed_dims,
levels=self.stage_max,
name='encoder')
#self.inter = Decoder(d_ch=ed_dims, total_levels=self.stage_max, levels_range=[0,2], name='inter')
self.decoder_src = Decoder(levels_range=[0, self.stage_max],
name='decoder_src')
self.decoder_dst = Decoder(levels_range=[0, self.stage_max],
name='decoder_dst')
self.model_filename_list += [
[self.encoder, 'encoder.npy'],
#[self.inter, 'inter.npy' ],
[self.decoder_src, 'decoder_src.npy'],
[self.decoder_dst, 'decoder_dst.npy']
]
if self.is_training:
self.src_dst_all_weights = self.encoder.get_weights(
) + self.decoder_src.get_weights(
) + self.decoder_dst.get_weights()
self.src_dst_trainable_weights = self.encoder.get_stage_weights(stage) \
+ self.decoder_src.get_stage_weights(stage) \
+ self.decoder_dst.get_stage_weights(stage)
# Initialize optimizers
self.src_dst_opt = nn.RMSprop(lr=2e-4,
lr_dropout=1.0,
name='src_dst_opt')
self.src_dst_opt.initialize_variables(
self.src_dst_all_weights,
vars_on_cpu=optimizer_vars_on_cpu)
self.model_filename_list += [(self.src_dst_opt,
'src_dst_opt.npy')]
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices))
bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
self.set_batch_size(gpu_count * bs_per_gpu)
# Compute losses per GPU
gpu_pred_src_src_list = []
gpu_pred_dst_dst_list = []
gpu_pred_src_dst_list = []
gpu_pred_src_srcm_list = []
gpu_pred_dst_dstm_list = []
gpu_pred_src_dstm_list = []
gpu_src_losses = []
gpu_dst_losses = []
gpu_src_dst_loss_gvs = []
for gpu_id in range(gpu_count):
with tf.device(
f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0'):
batch_slice = slice(gpu_id * bs_per_gpu,
(gpu_id + 1) * bs_per_gpu)
with tf.device(f'/CPU:0'):
# slice on CPU, otherwise all batch data will be transfered to GPU first
gpu_prev_warped_src = self.prev_warped_src[
batch_slice, :, :, :]
gpu_warped_src = self.warped_src[batch_slice, :, :, :]
gpu_prev_warped_dst = self.prev_warped_dst[
batch_slice, :, :, :]
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 = self.target_srcm[
batch_slice, :, :, :]
gpu_target_dstm = self.target_dstm[
batch_slice, :, :, :]
gpu_alpha_t = self.alpha_t[batch_slice, :, :, :]
# process model tensors
#gpu_src_code = self.inter(stage, self.encoder(stage, gpu_warped_src, gpu_prev_warped_src, gpu_alpha_t), gpu_alpha_t )
#gpu_dst_code = self.inter(stage, self.encoder(stage, gpu_warped_dst, gpu_prev_warped_dst, gpu_alpha_t), gpu_alpha_t )
gpu_src_code = self.encoder(stage, gpu_warped_src,
gpu_prev_warped_src,
gpu_alpha_t)
gpu_dst_code = self.encoder(stage, gpu_warped_dst,
gpu_prev_warped_dst,
gpu_alpha_t)
gpu_pred_src_src, gpu_pred_src_srcm = self.decoder_src(
stage, gpu_src_code, gpu_alpha_t) #, inter=self.inter)
gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder_dst(
stage, gpu_dst_code, gpu_alpha_t) #, inter=self.inter)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(
stage, gpu_dst_code, gpu_alpha_t) #, inter=self.inter)
gpu_pred_src_src_list.append(gpu_pred_src_src)
gpu_pred_dst_dst_list.append(gpu_pred_dst_dst)
gpu_pred_src_dst_list.append(gpu_pred_src_dst)
gpu_pred_src_srcm_list.append(gpu_pred_src_srcm)
gpu_pred_dst_dstm_list.append(gpu_pred_dst_dstm)
gpu_pred_src_dstm_list.append(gpu_pred_src_dstm)
gpu_target_srcm_blur = nn.gaussian_blur(
gpu_target_srcm, max(1, stage_resolution // 32))
gpu_target_dstm_blur = nn.gaussian_blur(
gpu_target_dstm, max(1, stage_resolution // 32))
gpu_target_dst_masked = gpu_target_dst * gpu_target_dstm_blur
gpu_target_dst_anti_masked = gpu_target_dst * (
1.0 - gpu_target_dstm_blur)
gpu_target_srcmasked_opt = gpu_target_src * gpu_target_srcm_blur if masked_training else gpu_target_src
gpu_target_dst_masked_opt = gpu_target_dst_masked if masked_training else gpu_target_dst
gpu_pred_src_src_masked_opt = gpu_pred_src_src * gpu_target_srcm_blur if masked_training else gpu_pred_src_src
gpu_pred_dst_dst_masked_opt = gpu_pred_dst_dst * gpu_target_dstm_blur if masked_training else gpu_pred_dst_dst
gpu_psd_target_dst_masked = gpu_pred_src_dst * gpu_target_dstm_blur
gpu_psd_target_dst_anti_masked = gpu_pred_src_dst * (
1.0 - gpu_target_dstm_blur)
gpu_src_loss = tf.reduce_mean(
10 * tf.square(gpu_target_srcmasked_opt -
gpu_pred_src_src_masked_opt),
axis=[1, 2, 3])
gpu_src_loss += tf.reduce_mean(
tf.square(gpu_target_srcm - gpu_pred_src_srcm),
axis=[1, 2, 3])
if stage_resolution >= 16:
gpu_src_loss += tf.reduce_mean(
5 *
nn.dssim(gpu_target_srcmasked_opt,
gpu_pred_src_src_masked_opt,
max_val=1.0,
filter_size=int(stage_resolution / 11.6)),
axis=[1])
if stage_resolution >= 32:
gpu_src_loss += tf.reduce_mean(
5 *
nn.dssim(gpu_target_srcmasked_opt,
gpu_pred_src_src_masked_opt,
max_val=1.0,
filter_size=int(stage_resolution / 23.2)),
axis=[1])
gpu_dst_loss = tf.reduce_mean(
10 * tf.square(gpu_target_dst_masked_opt -
gpu_pred_dst_dst_masked_opt),
axis=[1, 2, 3])
gpu_dst_loss += tf.reduce_mean(
tf.square(gpu_target_dstm - gpu_pred_dst_dstm),
axis=[1, 2, 3])
if stage_resolution >= 16:
gpu_dst_loss += tf.reduce_mean(
5 *
nn.dssim(gpu_target_dst_masked_opt,
gpu_pred_dst_dst_masked_opt,
max_val=1.0,
filter_size=int(stage_resolution / 11.6)),
axis=[1])
if stage_resolution >= 32:
gpu_dst_loss += tf.reduce_mean(
5 *
nn.dssim(gpu_target_dst_masked_opt,
gpu_pred_dst_dst_masked_opt,
max_val=1.0,
filter_size=int(stage_resolution / 23.2)),
axis=[1])
gpu_src_losses += [gpu_src_loss]
gpu_dst_losses += [gpu_dst_loss]
gpu_src_dst_loss = gpu_src_loss + gpu_dst_loss
gpu_src_dst_loss_gvs += [
nn.gradients(gpu_src_dst_loss,
self.src_dst_trainable_weights)
]
# Average losses and gradients, and create optimizer update ops
with tf.device(models_opt_device):
if gpu_count == 1:
pred_src_src = gpu_pred_src_src_list[0]
pred_dst_dst = gpu_pred_dst_dst_list[0]
pred_src_dst = gpu_pred_src_dst_list[0]
pred_src_srcm = gpu_pred_src_srcm_list[0]
pred_dst_dstm = gpu_pred_dst_dstm_list[0]
pred_src_dstm = gpu_pred_src_dstm_list[0]
src_loss = gpu_src_losses[0]
dst_loss = gpu_dst_losses[0]
src_dst_loss_gv = gpu_src_dst_loss_gvs[0]
else:
pred_src_src = tf.concat(gpu_pred_src_src_list, 0)
pred_dst_dst = tf.concat(gpu_pred_dst_dst_list, 0)
pred_src_dst = tf.concat(gpu_pred_src_dst_list, 0)
pred_src_srcm = tf.concat(gpu_pred_src_srcm_list, 0)
pred_dst_dstm = tf.concat(gpu_pred_dst_dstm_list, 0)
pred_src_dstm = tf.concat(gpu_pred_src_dstm_list, 0)
src_loss = nn.average_tensor_list(gpu_src_losses)
dst_loss = nn.average_tensor_list(gpu_dst_losses)
src_dst_loss_gv = nn.average_gv_list(gpu_src_dst_loss_gvs)
src_dst_loss_gv_op = self.src_dst_opt.get_update_op(
src_dst_loss_gv)
# Initializing training and view functions
def get_alpha(batch_size):
alpha = 0
if self.stage != 0:
alpha = (self.iter - self.start_stage_iter) / (
self.target_stage_iter - self.start_stage_iter)
alpha = np.clip(alpha, 0, 1)
alpha = np.array([alpha], nn.floatx.as_numpy_dtype).reshape(
(1, 1, 1, 1))
alpha = np.repeat(alpha, batch_size, 0)
return alpha
def src_dst_train(prev_warped_src, warped_src, target_src, target_srcm, \
prev_warped_dst, warped_dst, target_dst, target_dstm):
s, d, _ = nn.tf_sess.run(
[src_loss, dst_loss, src_dst_loss_gv_op],
feed_dict={
self.prev_warped_src: prev_warped_src,
self.warped_src: warped_src,
self.target_src: target_src,
self.target_srcm: target_srcm,
self.prev_warped_dst: prev_warped_dst,
self.warped_dst: warped_dst,
self.target_dst: target_dst,
self.target_dstm: target_dstm,
self.alpha_t: get_alpha(prev_warped_src.shape[0])
})
s = np.mean(s)
d = np.mean(d)
return s, d
self.src_dst_train = src_dst_train
def AE_view(prev_warped_src, warped_src, prev_warped_dst,
warped_dst):
return nn.tf_sess.run(
[
pred_src_src, pred_dst_dst, pred_dst_dstm,
pred_src_dst, pred_src_dstm
],
feed_dict={
self.prev_warped_src: prev_warped_src,
self.warped_src: warped_src,
self.prev_warped_dst: prev_warped_dst,
self.warped_dst: warped_dst,
self.alpha_t: get_alpha(prev_warped_src.shape[0])
})
self.AE_view = AE_view
else:
# Initializing merge function
with tf.device(f'/GPU:0' if len(devices) != 0 else f'/CPU:0'):
gpu_dst_code = self.inter(self.encoder(self.warped_dst))
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(
gpu_dst_code, stage=stage)
_, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code,
stage=stage)
def AE_merge(warped_dst):
return nn.tf_sess.run(
[gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm],
feed_dict={self.warped_dst: warped_dst})
self.AE_merge = AE_merge
# Loading/initializing all models/optimizers weights
for model, filename in io.progress_bar_generator(
self.model_filename_list, "Initializing models"):
do_init = self.is_first_run()
if self.pretrain_just_disabled:
if model == self.inter:
do_init = True
if not do_init:
do_init = not model.load_weights(
self.get_strpath_storage_for_file(filename))
if do_init:
model.init_weights()
# initializing sample generators
if self.is_training:
self.face_type = FaceType.FULL
training_data_src_path = self.training_data_src_path if not self.pretrain else self.get_pretraining_data_path(
)
training_data_dst_path = self.training_data_dst_path if not self.pretrain else self.get_pretraining_data_path(
)
cpu_count = multiprocessing.cpu_count()
src_generators_count = cpu_count // 2
dst_generators_count = cpu_count - src_generators_count
self.set_training_data_generators([
SampleGeneratorFace(
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': True,
'transform': True,
'channel_type': SampleProcessor.ChannelType.BGR,
'face_type': self.face_type,
'data_format': nn.data_format,
'resolution': resolution,
'nearest_resize_to': prev_stage_resolution
},
{
'sample_type':
SampleProcessor.SampleType.FACE_IMAGE,
'warp': True,
'transform': True,
'channel_type': SampleProcessor.ChannelType.BGR,
'face_type': self.face_type,
'data_format': nn.data_format,
'resolution': resolution,
'nearest_resize_to': stage_resolution
},
{
'sample_type':
SampleProcessor.SampleType.FACE_IMAGE,
'warp': False,
'transform': True,
'channel_type': SampleProcessor.ChannelType.BGR,
'face_type': self.face_type,
'data_format': nn.data_format,
'resolution': resolution,
'nearest_resize_to': prev_stage_resolution
},
{
'sample_type':
SampleProcessor.SampleType.FACE_IMAGE,
'warp': False,
'transform': True,
'channel_type': SampleProcessor.ChannelType.BGR,
'face_type': self.face_type,
'data_format': nn.data_format,
'resolution': resolution,
'nearest_resize_to': stage_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,
'nearest_resize_to': stage_resolution
},
],
generators_count=src_generators_count),
SampleGeneratorFace(
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': True,
'transform': True,
'channel_type': SampleProcessor.ChannelType.BGR,
'face_type': self.face_type,
'data_format': nn.data_format,
'resolution': resolution,
'nearest_resize_to': prev_stage_resolution
},
{
'sample_type':
SampleProcessor.SampleType.FACE_IMAGE,
'warp': True,
'transform': True,
'channel_type': SampleProcessor.ChannelType.BGR,
'face_type': self.face_type,
'data_format': nn.data_format,
'resolution': resolution,
'nearest_resize_to': stage_resolution
},
{
'sample_type':
SampleProcessor.SampleType.FACE_IMAGE,
'warp': False,
'transform': True,
'channel_type': SampleProcessor.ChannelType.BGR,
'face_type': self.face_type,
'data_format': nn.data_format,
'resolution': resolution,
'nearest_resize_to': prev_stage_resolution
},
{
'sample_type':
SampleProcessor.SampleType.FACE_IMAGE,
'warp': False,
'transform': True,
'channel_type': SampleProcessor.ChannelType.BGR,
'face_type': self.face_type,
'data_format': nn.data_format,
'resolution': resolution,
'nearest_resize_to': stage_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,
'nearest_resize_to': stage_resolution
},
],
generators_count=dst_generators_count)
])
self.last_samples = None
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
self.model_data_format = "NCHW"
nn.initialize(data_format=self.model_data_format)
tf = nn.tf
input_ch=3
resolution = self.resolution = self.options['resolution']
e_dims = self.options['e_dims']
ae_dims = self.options['ae_dims']
inter_dims = self.inter_dims = self.options['inter_dims']
inter_res = self.inter_res = resolution // 32
d_dims = self.options['d_dims']
d_mask_dims = self.options['d_mask_dims']
face_type = self.face_type = {'f' : FaceType.FULL,
'wf' : FaceType.WHOLE_FACE,
'head' : FaceType.HEAD}[ self.options['face_type'] ]
morph_factor = self.options['morph_factor']
gan_power = self.gan_power = self.options['gan_power']
random_warp = self.options['random_warp']
blur_out_mask = self.options['blur_out_mask']
ct_mode = self.options['ct_mode']
if ct_mode == 'none':
ct_mode = None
use_fp16 = False
if self.is_exporting:
use_fp16 = io.input_bool ("Export quantized?", False, help_message='Makes the exported model faster. If you have problems, disable this option.')
conv_dtype = tf.float16 if use_fp16 else tf.float32
class Downscale(nn.ModelBase):
def on_build(self, in_ch, out_ch, kernel_size=5 ):
self.conv1 = nn.Conv2D( in_ch, out_ch, kernel_size=kernel_size, strides=2, padding='SAME', dtype=conv_dtype)
def forward(self, x):
return tf.nn.leaky_relu(self.conv1(x), 0.1)
class Upscale(nn.ModelBase):
def on_build(self, in_ch, out_ch, kernel_size=3 ):
self.conv1 = nn.Conv2D(in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype)
def forward(self, x):
x = nn.depth_to_space(tf.nn.leaky_relu(self.conv1(x), 0.1), 2)
return x
class ResidualBlock(nn.ModelBase):
def on_build(self, ch, kernel_size=3 ):
self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype)
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', dtype=conv_dtype)
def forward(self, inp):
x = self.conv1(inp)
x = tf.nn.leaky_relu(x, 0.2)
x = self.conv2(x)
x = tf.nn.leaky_relu(inp+x, 0.2)
return x
class Encoder(nn.ModelBase):
def on_build(self):
self.down1 = Downscale(input_ch, e_dims, kernel_size=5)
self.res1 = ResidualBlock(e_dims)
self.down2 = Downscale(e_dims, e_dims*2, kernel_size=5)
self.down3 = Downscale(e_dims*2, e_dims*4, kernel_size=5)
self.down4 = Downscale(e_dims*4, e_dims*8, kernel_size=5)
self.down5 = Downscale(e_dims*8, e_dims*8, kernel_size=5)
self.res5 = ResidualBlock(e_dims*8)
self.dense1 = nn.Dense( (( resolution//(2**5) )**2) * e_dims*8, ae_dims )
def forward(self, x):
if use_fp16:
x = tf.cast(x, tf.float16)
x = self.down1(x)
x = self.res1(x)
x = self.down2(x)
x = self.down3(x)
x = self.down4(x)
x = self.down5(x)
x = self.res5(x)
if use_fp16:
x = tf.cast(x, tf.float32)
x = nn.pixel_norm(nn.flatten(x), axes=-1)
x = self.dense1(x)
return x
class Inter(nn.ModelBase):
def on_build(self):
self.dense2 = nn.Dense(ae_dims, inter_res * inter_res * inter_dims)
def forward(self, inp):
x = inp
x = self.dense2(x)
x = nn.reshape_4D (x, inter_res, inter_res, inter_dims)
return x
class Decoder(nn.ModelBase):
def on_build(self ):
self.upscale0 = Upscale(inter_dims, d_dims*8, kernel_size=3)
self.upscale1 = Upscale(d_dims*8, d_dims*8, kernel_size=3)
self.upscale2 = Upscale(d_dims*8, d_dims*4, kernel_size=3)
self.upscale3 = Upscale(d_dims*4, d_dims*2, kernel_size=3)
self.res0 = ResidualBlock(d_dims*8, kernel_size=3)
self.res1 = ResidualBlock(d_dims*8, kernel_size=3)
self.res2 = ResidualBlock(d_dims*4, kernel_size=3)
self.res3 = ResidualBlock(d_dims*2, kernel_size=3)
self.upscalem0 = Upscale(inter_dims, d_mask_dims*8, kernel_size=3)
self.upscalem1 = Upscale(d_mask_dims*8, d_mask_dims*8, kernel_size=3)
self.upscalem2 = Upscale(d_mask_dims*8, d_mask_dims*4, kernel_size=3)
self.upscalem3 = Upscale(d_mask_dims*4, d_mask_dims*2, kernel_size=3)
self.upscalem4 = Upscale(d_mask_dims*2, d_mask_dims*1, kernel_size=3)
self.out_convm = nn.Conv2D( d_mask_dims*1, 1, kernel_size=1, padding='SAME', dtype=conv_dtype)
self.out_conv = nn.Conv2D( d_dims*2, 3, kernel_size=1, padding='SAME', dtype=conv_dtype)
self.out_conv1 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype)
self.out_conv2 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype)
self.out_conv3 = nn.Conv2D( d_dims*2, 3, kernel_size=3, padding='SAME', dtype=conv_dtype)
def forward(self, z):
if use_fp16:
z = tf.cast(z, tf.float16)
x = self.upscale0(z)
x = self.res0(x)
x = self.upscale1(x)
x = self.res1(x)
x = self.upscale2(x)
x = self.res2(x)
x = self.upscale3(x)
x = self.res3(x)
x = tf.nn.sigmoid( nn.depth_to_space(tf.concat( (self.out_conv(x),
self.out_conv1(x),
self.out_conv2(x),
self.out_conv3(x)), nn.conv2d_ch_axis), 2) )
m = self.upscalem0(z)
m = self.upscalem1(m)
m = self.upscalem2(m)
m = self.upscalem3(m)
m = self.upscalem4(m)
m = tf.nn.sigmoid(self.out_convm(m))
if use_fp16:
x = tf.cast(x, tf.float32)
m = tf.cast(m, tf.float32)
return x, m
models_opt_on_gpu = False if len(devices) == 0 else self.options['models_opt_on_gpu']
models_opt_device = nn.tf_default_device_name if models_opt_on_gpu and self.is_training else '/CPU:0'
optimizer_vars_on_cpu = models_opt_device=='/CPU:0'
bgr_shape = self.bgr_shape = nn.get4Dshape(resolution,resolution,input_ch)
mask_shape = nn.get4Dshape(resolution,resolution,1)
self.model_filename_list = []
with tf.device ('/CPU:0'):
#Place holders on CPU
self.warped_src = tf.placeholder (nn.floatx, bgr_shape, name='warped_src')
self.warped_dst = tf.placeholder (nn.floatx, bgr_shape, name='warped_dst')
self.target_src = tf.placeholder (nn.floatx, bgr_shape, name='target_src')
self.target_dst = tf.placeholder (nn.floatx, bgr_shape, name='target_dst')
self.target_srcm = tf.placeholder (nn.floatx, mask_shape, name='target_srcm')
self.target_srcm_em = tf.placeholder (nn.floatx, mask_shape, name='target_srcm_em')
self.target_dstm = tf.placeholder (nn.floatx, mask_shape, name='target_dstm')
self.target_dstm_em = tf.placeholder (nn.floatx, mask_shape, name='target_dstm_em')
self.morph_value_t = tf.placeholder (nn.floatx, (1,), name='morph_value_t')
# Initializing model classes
with tf.device (models_opt_device):
self.encoder = Encoder(name='encoder')
self.inter_src = Inter(name='inter_src')
self.inter_dst = Inter(name='inter_dst')
self.decoder = Decoder(name='decoder')
self.model_filename_list += [ [self.encoder, 'encoder.npy'],
[self.inter_src, 'inter_src.npy'],
[self.inter_dst , 'inter_dst.npy'],
[self.decoder , 'decoder.npy'] ]
if self.is_training:
# Initialize optimizers
clipnorm = 1.0 if self.options['clipgrad'] else 0.0
lr_dropout = 0.3 if self.options['lr_dropout'] in ['y','cpu'] else 1.0
self.G_weights = self.encoder.get_weights() + self.decoder.get_weights()
#if random_warp:
# self.G_weights += self.inter_src.get_weights() + self.inter_dst.get_weights()
self.src_dst_opt = nn.AdaBelief(lr=5e-5, lr_dropout=lr_dropout, clipnorm=clipnorm, name='src_dst_opt')
self.src_dst_opt.initialize_variables (self.G_weights, vars_on_cpu=optimizer_vars_on_cpu)
self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ]
if gan_power != 0:
self.GAN = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], name="GAN")
self.GAN_opt = nn.AdaBelief(lr=5e-5, lr_dropout=lr_dropout, clipnorm=clipnorm, name='GAN_opt')
self.GAN_opt.initialize_variables ( self.GAN.get_weights(), vars_on_cpu=optimizer_vars_on_cpu)
self.model_filename_list += [ [self.GAN, 'GAN.npy'],
[self.GAN_opt, 'GAN_opt.npy'] ]
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices) )
bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
self.set_batch_size( gpu_count*bs_per_gpu)
# Compute losses per GPU
gpu_pred_src_src_list = []
gpu_pred_dst_dst_list = []
gpu_pred_src_dst_list = []
gpu_pred_src_srcm_list = []
gpu_pred_dst_dstm_list = []
gpu_pred_src_dstm_list = []
gpu_src_losses = []
gpu_dst_losses = []
gpu_G_loss_gradients = []
gpu_GAN_loss_gradients = []
def DLossOnes(logits):
return tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.ones_like(logits), logits=logits), axis=[1,2,3])
def DLossZeros(logits):
return tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.zeros_like(logits), logits=logits), axis=[1,2,3])
for gpu_id in range(gpu_count):
with tf.device( f'/{devices[gpu_id].tf_dev_type}:{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_warped_src = self.warped_src [batch_slice,:,:,:]
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 = self.target_srcm[batch_slice,:,:,:]
gpu_target_srcm_em = self.target_srcm_em[batch_slice,:,:,:]
gpu_target_dstm = self.target_dstm[batch_slice,:,:,:]
gpu_target_dstm_em = self.target_dstm_em[batch_slice,:,:,:]
# process model tensors
gpu_src_code = self.encoder (gpu_warped_src)
gpu_dst_code = self.encoder (gpu_warped_dst)
gpu_src_inter_src_code, gpu_src_inter_dst_code = self.inter_src (gpu_src_code), self.inter_dst (gpu_src_code)
gpu_dst_inter_src_code, gpu_dst_inter_dst_code = self.inter_src (gpu_dst_code), self.inter_dst (gpu_dst_code)
inter_dims_bin = int(inter_dims*morph_factor)
with tf.device(f'/CPU:0'):
inter_rnd_binomial = tf.stack([tf.random.shuffle(tf.concat([tf.tile(tf.constant([1], tf.float32), ( inter_dims_bin, )),
tf.tile(tf.constant([0], tf.float32), ( inter_dims-inter_dims_bin, ))], 0 )) for _ in range(bs_per_gpu)], 0)
inter_rnd_binomial = tf.stop_gradient(inter_rnd_binomial[...,None,None])
gpu_src_code = gpu_src_inter_src_code * inter_rnd_binomial + gpu_src_inter_dst_code * (1-inter_rnd_binomial)
gpu_dst_code = gpu_dst_inter_dst_code
inter_dims_slice = tf.cast(inter_dims*self.morph_value_t[0], tf.int32)
gpu_src_dst_code = tf.concat( (tf.slice(gpu_dst_inter_src_code, [0,0,0,0], [-1, inter_dims_slice , inter_res, inter_res]),
tf.slice(gpu_dst_inter_dst_code, [0,inter_dims_slice,0,0], [-1,inter_dims-inter_dims_slice, inter_res,inter_res]) ), 1 )
gpu_pred_src_src, gpu_pred_src_srcm = self.decoder(gpu_src_code)
gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder(gpu_dst_code)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
gpu_pred_src_src_list.append(gpu_pred_src_src), gpu_pred_src_srcm_list.append(gpu_pred_src_srcm)
gpu_pred_dst_dst_list.append(gpu_pred_dst_dst), gpu_pred_dst_dstm_list.append(gpu_pred_dst_dstm)
gpu_pred_src_dst_list.append(gpu_pred_src_dst), gpu_pred_src_dstm_list.append(gpu_pred_src_dstm)
gpu_target_srcm_anti = 1-gpu_target_srcm
gpu_target_dstm_anti = 1-gpu_target_dstm
gpu_target_srcm_gblur = nn.gaussian_blur(gpu_target_srcm, resolution // 32)
gpu_target_dstm_gblur = nn.gaussian_blur(gpu_target_dstm, resolution // 32)
gpu_target_srcm_blur = tf.clip_by_value(gpu_target_srcm_gblur, 0, 0.5) * 2
gpu_target_dstm_blur = tf.clip_by_value(gpu_target_dstm_gblur, 0, 0.5) * 2
gpu_target_srcm_anti_blur = 1.0-gpu_target_srcm_blur
gpu_target_dstm_anti_blur = 1.0-gpu_target_dstm_blur
if blur_out_mask:
sigma = resolution / 128
x = nn.gaussian_blur(gpu_target_src*gpu_target_srcm_anti, sigma)
y = 1-nn.gaussian_blur(gpu_target_srcm, sigma)
y = tf.where(tf.equal(y, 0), tf.ones_like(y), y)
gpu_target_src = gpu_target_src*gpu_target_srcm + (x/y)*gpu_target_srcm_anti
x = nn.gaussian_blur(gpu_target_dst*gpu_target_dstm_anti, sigma)
y = 1-nn.gaussian_blur(gpu_target_dstm, sigma)
y = tf.where(tf.equal(y, 0), tf.ones_like(y), y)
gpu_target_dst = gpu_target_dst*gpu_target_dstm + (x/y)*gpu_target_dstm_anti
gpu_target_src_masked = gpu_target_src*gpu_target_srcm_blur
gpu_target_dst_masked = gpu_target_dst*gpu_target_dstm_blur
gpu_target_src_anti_masked = gpu_target_src*gpu_target_srcm_anti_blur
gpu_target_dst_anti_masked = gpu_target_dst*gpu_target_dstm_anti_blur
gpu_pred_src_src_masked = gpu_pred_src_src*gpu_target_srcm_blur
gpu_pred_dst_dst_masked = gpu_pred_dst_dst*gpu_target_dstm_blur
gpu_pred_src_src_anti_masked = gpu_pred_src_src*gpu_target_srcm_anti_blur
gpu_pred_dst_dst_anti_masked = gpu_pred_dst_dst*gpu_target_dstm_anti_blur
# Structural loss
gpu_src_loss = tf.reduce_mean (5*nn.dssim(gpu_target_src_masked, gpu_pred_src_src_masked, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
gpu_src_loss += tf.reduce_mean (5*nn.dssim(gpu_target_src_masked, gpu_pred_src_src_masked, max_val=1.0, filter_size=int(resolution/23.2)), axis=[1])
gpu_dst_loss = tf.reduce_mean (5*nn.dssim(gpu_target_dst_masked, gpu_pred_dst_dst_masked, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
gpu_dst_loss += tf.reduce_mean (5*nn.dssim(gpu_target_dst_masked, gpu_pred_dst_dst_masked, max_val=1.0, filter_size=int(resolution/23.2) ), axis=[1])
# Pixel loss
gpu_src_loss += tf.reduce_mean (10*tf.square(gpu_target_src_masked-gpu_pred_src_src_masked), axis=[1,2,3])
gpu_dst_loss += tf.reduce_mean (10*tf.square(gpu_target_dst_masked-gpu_pred_dst_dst_masked), axis=[1,2,3])
# Eyes+mouth prio loss
gpu_src_loss += tf.reduce_mean (300*tf.abs (gpu_target_src*gpu_target_srcm_em-gpu_pred_src_src*gpu_target_srcm_em), axis=[1,2,3])
gpu_dst_loss += tf.reduce_mean (300*tf.abs (gpu_target_dst*gpu_target_dstm_em-gpu_pred_dst_dst*gpu_target_dstm_em), axis=[1,2,3])
# Mask loss
gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] )
gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] )
gpu_src_losses += [gpu_src_loss]
gpu_dst_losses += [gpu_dst_loss]
gpu_G_loss = gpu_src_loss + gpu_dst_loss
# dst-dst background weak loss
gpu_G_loss += tf.reduce_mean(0.1*tf.square(gpu_pred_dst_dst_anti_masked-gpu_target_dst_anti_masked),axis=[1,2,3] )
gpu_G_loss += 0.000001*nn.total_variation_mse(gpu_pred_dst_dst_anti_masked)
if gan_power != 0:
gpu_pred_src_src_d, gpu_pred_src_src_d2 = self.GAN(gpu_pred_src_src_masked)
gpu_pred_dst_dst_d, gpu_pred_dst_dst_d2 = self.GAN(gpu_pred_dst_dst_masked)
gpu_target_src_d, gpu_target_src_d2 = self.GAN(gpu_target_src_masked)
gpu_target_dst_d, gpu_target_dst_d2 = self.GAN(gpu_target_dst_masked)
gpu_GAN_loss = (DLossOnes (gpu_target_src_d) + DLossOnes (gpu_target_src_d2) + \
DLossZeros(gpu_pred_src_src_d) + DLossZeros(gpu_pred_src_src_d2) + \
DLossOnes (gpu_target_dst_d) + DLossOnes (gpu_target_dst_d2) + \
DLossZeros(gpu_pred_dst_dst_d) + DLossZeros(gpu_pred_dst_dst_d2)
) * (1.0 / 8)
gpu_GAN_loss_gradients += [ nn.gradients (gpu_GAN_loss, self.GAN.get_weights() ) ]
gpu_G_loss += (DLossOnes(gpu_pred_src_src_d) + DLossOnes(gpu_pred_src_src_d2) + \
DLossOnes(gpu_pred_dst_dst_d) + DLossOnes(gpu_pred_dst_dst_d2)
) * gan_power
# Minimal src-src-bg rec with total_variation_mse to suppress random bright dots from gan
gpu_G_loss += 0.000001*nn.total_variation_mse(gpu_pred_src_src)
gpu_G_loss += 0.02*tf.reduce_mean(tf.square(gpu_pred_src_src_anti_masked-gpu_target_src_anti_masked),axis=[1,2,3] )
gpu_G_loss_gradients += [ nn.gradients ( gpu_G_loss, self.G_weights ) ]
# Average losses and gradients, and create optimizer update ops
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)
pred_src_srcm = nn.concat(gpu_pred_src_srcm_list, 0)
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)
train_op = self.src_dst_opt.get_update_op (nn.average_gv_list (gpu_G_loss_gradients))
if gan_power != 0:
GAN_train_op = self.GAN_opt.get_update_op (nn.average_gv_list(gpu_GAN_loss_gradients) )
# Initializing training and view functions
def 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, train_op],
feed_dict={self.warped_src :warped_src,
self.target_src :target_src,
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:target_dstm,
self.target_dstm_em:target_dstm_em,
})
return s, d
self.train = train
if gan_power != 0:
def GAN_train(warped_src, target_src, target_srcm, target_srcm_em, \
warped_dst, target_dst, target_dstm, target_dstm_em, ):
nn.tf_sess.run ([GAN_train_op], feed_dict={self.warped_src :warped_src,
self.target_src :target_src,
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:target_dstm,
self.target_dstm_em:target_dstm_em})
self.GAN_train = GAN_train
def AE_view(warped_src, warped_dst, morph_value):
return nn.tf_sess.run ( [pred_src_src, pred_dst_dst, pred_dst_dstm, pred_src_dst, pred_src_dstm],
feed_dict={self.warped_src:warped_src, self.warped_dst:warped_dst, self.morph_value_t:[morph_value] })
self.AE_view = AE_view
else:
#Initializing merge function
with tf.device( nn.tf_default_device_name if len(devices) != 0 else f'/CPU:0'):
gpu_dst_code = self.encoder (self.warped_dst)
gpu_dst_inter_src_code = self.inter_src (gpu_dst_code)
gpu_dst_inter_dst_code = self.inter_dst (gpu_dst_code)
inter_dims_slice = tf.cast(inter_dims*self.morph_value_t[0], tf.int32)
gpu_src_dst_code = tf.concat( ( tf.slice(gpu_dst_inter_src_code, [0,0,0,0], [-1, inter_dims_slice , inter_res, inter_res]),
tf.slice(gpu_dst_inter_dst_code, [0,inter_dims_slice,0,0], [-1,inter_dims-inter_dims_slice, inter_res,inter_res]) ), 1 )
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
_, gpu_pred_dst_dstm = self.decoder(gpu_dst_inter_dst_code)
def AE_merge(warped_dst, morph_value):
return nn.tf_sess.run ( [gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm], feed_dict={self.warped_dst:warped_dst, self.morph_value_t:[morph_value] })
self.AE_merge = AE_merge
# Loading/initializing all models/optimizers weights
for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"):
do_init = self.is_first_run()
if self.is_training and gan_power != 0 and model == self.GAN:
if self.gan_model_changed:
do_init = True
if not do_init:
do_init = not model.load_weights( self.get_strpath_storage_for_file(filename) )
if do_init:
model.init_weights()
###############
# initializing sample generators
if self.is_training:
training_data_src_path = self.training_data_src_path #if not self.pretrain else self.get_pretraining_data_path()
training_data_dst_path = self.training_data_dst_path #if not self.pretrain else self.get_pretraining_data_path()
random_ct_samples_path=training_data_dst_path if ct_mode is not None else None #and not self.pretrain
cpu_count = multiprocessing.cpu_count()
src_generators_count = cpu_count // 2
dst_generators_count = cpu_count // 2
if ct_mode is not None:
src_generators_count = int(src_generators_count * 1.5)
self.set_training_data_generators ([
SampleGeneratorFace(training_data_src_path, random_ct_samples_path=random_ct_samples_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=self.random_src_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':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':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':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.EYES_MOUTH, 'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution},
],
uniform_yaw_distribution=self.options['uniform_yaw'],# or self.pretrain,
generators_count=src_generators_count ),
SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=self.random_dst_flip),
output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution},
{'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'face_type':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':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.EYES_MOUTH, 'face_type':face_type, 'data_format':nn.data_format, 'resolution': resolution},
],
uniform_yaw_distribution=self.options['uniform_yaw'],# or self.pretrain,
generators_count=dst_generators_count )
])
self.last_src_samples_loss = []
self.last_dst_samples_loss = []
def __init__(self,
name,
resolution,
load_weights=True,
weights_file_root=None,
training=False,
place_model_on_cpu=False,
run_on_cpu=False,
optimizer=None,
data_format="NHWC"):
nn.initialize(data_format=data_format)
tf = nn.tf
self.weights_file_root = Path(
weights_file_root) if weights_file_root is not None else Path(
__file__).parent
with tf.device('/CPU:0'):
#Place holders on CPU
self.input_t = tf.placeholder(
nn.floatx, nn.get4Dshape(resolution, resolution, 3))
self.target_t = tf.placeholder(
nn.floatx, nn.get4Dshape(resolution, resolution, 1))
# Initializing model classes
archi = nn.DFLSegnetArchi()
with tf.device('/CPU:0' if place_model_on_cpu else '/GPU:0'):
self.enc = archi.Encoder(3, 64, name='Encoder')
self.dec = archi.Decoder(64, 1, name='Decoder')
self.enc_dec_weights = self.enc.get_weights(
) + self.dec.get_weights()
model_name = f'{name}_{resolution}'
self.model_filename_list = [
[self.enc, f'{model_name}_enc.npy'],
[self.dec, f'{model_name}_dec.npy'],
]
if training:
if optimizer is None:
raise ValueError(
"Optimizer should be provided for training mode.")
self.opt = optimizer
self.opt.initialize_variables(self.enc_dec_weights,
vars_on_cpu=place_model_on_cpu)
self.model_filename_list += [[self.opt, f'{model_name}_opt.npy']]
else:
with tf.device('/CPU:0' if run_on_cpu else '/GPU:0'):
_, pred = self.dec(self.enc(self.input_t))
def net_run(input_np):
return nn.tf_sess.run([pred],
feed_dict={self.input_t: input_np})[0]
self.net_run = net_run
# Loading/initializing all models/optimizers weights
for model, filename in self.model_filename_list:
do_init = not load_weights
if not do_init:
do_init = not model.load_weights(
self.weights_file_root / filename)
if do_init:
model.init_weights()
def __init__(self,
name,
resolution,
load_weights=True,
weights_file_root=None,
training=False,
place_model_on_cpu=False,
run_on_cpu=False,
optimizer=None,
data_format="NHWC"):
nn.initialize(data_format=data_format)
tf = nn.tf
if weights_file_root is not None:
weights_file_root = Path(weights_file_root)
else:
weights_file_root = Path(__file__).parent
self.weights_file_root = weights_file_root
with tf.device('/CPU:0'):
#Place holders on CPU
self.input_t = tf.placeholder(
nn.floatx, nn.get4Dshape(resolution, resolution, 3))
self.target_t = tf.placeholder(
nn.floatx, nn.get4Dshape(resolution, resolution, 1))
# Initializing model classes
with tf.device('/CPU:0' if place_model_on_cpu else '/GPU:0'):
self.net = nn.Ternaus(3, 64, name='Ternaus')
self.net_weights = self.net.get_weights()
model_name = f'{name}_{resolution}'
self.model_filename_list = [[self.net, f'{model_name}.npy']]
if training:
if optimizer is None:
raise ValueError(
"Optimizer should be provided for traning mode.")
self.opt = optimizer
self.opt.initialize_variables(self.net_weights,
vars_on_cpu=place_model_on_cpu)
self.model_filename_list += [[self.opt, f'{model_name}_opt.npy']]
else:
with tf.device('/CPU:0' if run_on_cpu else '/GPU:0'):
_, pred = self.net([self.input_t])
def net_run(input_np):
return nn.tf_sess.run([pred],
feed_dict={self.input_t: input_np})[0]
self.net_run = net_run
# Loading/initializing all models/optimizers weights
for model, filename in self.model_filename_list:
do_init = not load_weights
if not do_init:
do_init = not model.load_weights(
self.weights_file_root / filename)
if do_init:
model.init_weights()
if model == self.net:
try:
with open(
Path(__file__).parent /
'vgg11_enc_weights.npy', 'rb') as f:
d = pickle.loads(f.read())
for i in [0, 3, 6, 8, 11, 13, 16, 18]:
model.get_layer_by_name('features_%d' %
i).set_weights(
d['features.%d' % i])
except:
io.log_err(
"Unable to load VGG11 pretrained weights from vgg11_enc_weights.npy"
)
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
self.model_data_format = "NCHW" if len(
devices) != 0 and not self.is_debug() else "NHWC"
nn.initialize(data_format=self.model_data_format)
tf = nn.tf
resolution = self.resolution = 96
self.face_type = FaceType.FULL
ae_dims = 128
e_dims = 128
d_dims = 64
self.pretrain = False
self.pretrain_just_disabled = False
masked_training = True
models_opt_on_gpu = len(devices) >= 1 and all(
[dev.total_mem_gb >= 4 for dev in devices])
models_opt_device = '/GPU:0' if models_opt_on_gpu and self.is_training else '/CPU:0'
optimizer_vars_on_cpu = models_opt_device == '/CPU:0'
input_ch = 3
bgr_shape = nn.get4Dshape(resolution, resolution, input_ch)
mask_shape = nn.get4Dshape(resolution, resolution, 1)
self.model_filename_list = []
model_archi = nn.DeepFakeArchi(resolution, mod='quick')
with tf.device('/CPU:0'):
#Place holders on CPU
self.warped_src = tf.placeholder(nn.floatx, bgr_shape)
self.warped_dst = tf.placeholder(nn.floatx, bgr_shape)
self.target_src = tf.placeholder(nn.floatx, bgr_shape)
self.target_dst = tf.placeholder(nn.floatx, bgr_shape)
self.target_srcm = tf.placeholder(nn.floatx, mask_shape)
self.target_dstm = tf.placeholder(nn.floatx, mask_shape)
# 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))
self.inter = model_archi.Inter(in_ch=encoder_out_ch,
ae_ch=ae_dims,
ae_out_ch=ae_dims,
d_ch=d_dims,
name='inter')
inter_out_ch = self.inter.compute_output_channels(
(nn.floatx, (None, encoder_out_ch)))
self.decoder_src = model_archi.Decoder(in_ch=inter_out_ch,
d_ch=d_dims,
name='decoder_src')
self.decoder_dst = model_archi.Decoder(in_ch=inter_out_ch,
d_ch=d_dims,
name='decoder_dst')
self.model_filename_list += [[self.encoder, 'encoder.npy'],
[self.inter, 'inter.npy'],
[self.decoder_src, 'decoder_src.npy'],
[self.decoder_dst, 'decoder_dst.npy']]
if self.is_training:
self.src_dst_trainable_weights = self.encoder.get_weights(
) + self.inter.get_weights() + self.decoder_src.get_weights(
) + self.decoder_dst.get_weights()
# Initialize optimizers
self.src_dst_opt = nn.RMSprop(lr=2e-4,
lr_dropout=0.3,
name='src_dst_opt')
self.src_dst_opt.initialize_variables(
self.src_dst_trainable_weights,
vars_on_cpu=optimizer_vars_on_cpu)
self.model_filename_list += [(self.src_dst_opt,
'src_dst_opt.npy')]
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices))
bs_per_gpu = max(1, 4 // gpu_count)
self.set_batch_size(gpu_count * bs_per_gpu)
# Compute losses per GPU
gpu_pred_src_src_list = []
gpu_pred_dst_dst_list = []
gpu_pred_src_dst_list = []
gpu_pred_src_srcm_list = []
gpu_pred_dst_dstm_list = []
gpu_pred_src_dstm_list = []
gpu_src_losses = []
gpu_dst_losses = []
gpu_src_dst_loss_gvs = []
for gpu_id in range(gpu_count):
with tf.device(
f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0'):
batch_slice = slice(gpu_id * bs_per_gpu,
(gpu_id + 1) * bs_per_gpu)
with tf.device(f'/CPU:0'):
# slice on CPU, otherwise all batch data will be transfered to GPU first
gpu_warped_src = self.warped_src[batch_slice, :, :, :]
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 = self.target_srcm[
batch_slice, :, :, :]
gpu_target_dstm = self.target_dstm[
batch_slice, :, :, :]
# process model tensors
gpu_src_code = self.inter(self.encoder(gpu_warped_src))
gpu_dst_code = self.inter(self.encoder(gpu_warped_dst))
gpu_pred_src_src, gpu_pred_src_srcm = self.decoder_src(
gpu_src_code)
gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder_dst(
gpu_dst_code)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(
gpu_dst_code)
gpu_pred_src_src_list.append(gpu_pred_src_src)
gpu_pred_dst_dst_list.append(gpu_pred_dst_dst)
gpu_pred_src_dst_list.append(gpu_pred_src_dst)
gpu_pred_src_srcm_list.append(gpu_pred_src_srcm)
gpu_pred_dst_dstm_list.append(gpu_pred_dst_dstm)
gpu_pred_src_dstm_list.append(gpu_pred_src_dstm)
gpu_target_srcm_blur = nn.gaussian_blur(
gpu_target_srcm, max(1, resolution // 32))
gpu_target_dstm_blur = nn.gaussian_blur(
gpu_target_dstm, max(1, resolution // 32))
gpu_target_dst_masked = gpu_target_dst * gpu_target_dstm_blur
gpu_target_dst_anti_masked = gpu_target_dst * (
1.0 - gpu_target_dstm_blur)
gpu_target_src_masked_opt = gpu_target_src * gpu_target_srcm_blur if masked_training else gpu_target_src
gpu_target_dst_masked_opt = gpu_target_dst_masked if masked_training else gpu_target_dst
gpu_pred_src_src_masked_opt = gpu_pred_src_src * gpu_target_srcm_blur if masked_training else gpu_pred_src_src
gpu_pred_dst_dst_masked_opt = gpu_pred_dst_dst * gpu_target_dstm_blur if masked_training else gpu_pred_dst_dst
gpu_psd_target_dst_masked = gpu_pred_src_dst * gpu_target_dstm_blur
gpu_psd_target_dst_anti_masked = gpu_pred_src_dst * (
1.0 - gpu_target_dstm_blur)
gpu_src_loss = tf.reduce_mean(
10 * nn.dssim(gpu_target_src_masked_opt,
gpu_pred_src_src_masked_opt,
max_val=1.0,
filter_size=int(resolution / 11.6)),
axis=[1])
gpu_src_loss += tf.reduce_mean(
10 * tf.square(gpu_target_src_masked_opt -
gpu_pred_src_src_masked_opt),
axis=[1, 2, 3])
gpu_src_loss += tf.reduce_mean(
10 * tf.square(gpu_target_srcm - gpu_pred_src_srcm),
axis=[1, 2, 3])
gpu_dst_loss = tf.reduce_mean(
10 * nn.dssim(gpu_target_dst_masked_opt,
gpu_pred_dst_dst_masked_opt,
max_val=1.0,
filter_size=int(resolution / 11.6)),
axis=[1])
gpu_dst_loss += tf.reduce_mean(
10 * tf.square(gpu_target_dst_masked_opt -
gpu_pred_dst_dst_masked_opt),
axis=[1, 2, 3])
gpu_dst_loss += tf.reduce_mean(
10 * tf.square(gpu_target_dstm - gpu_pred_dst_dstm),
axis=[1, 2, 3])
gpu_src_losses += [gpu_src_loss]
gpu_dst_losses += [gpu_dst_loss]
gpu_G_loss = gpu_src_loss + gpu_dst_loss
gpu_src_dst_loss_gvs += [
nn.gradients(gpu_G_loss,
self.src_dst_trainable_weights)
]
# Average losses and gradients, and create optimizer update ops
with tf.device(models_opt_device):
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)
pred_src_srcm = nn.concat(gpu_pred_src_srcm_list, 0)
pred_dst_dstm = nn.concat(gpu_pred_dst_dstm_list, 0)
pred_src_dstm = nn.concat(gpu_pred_src_dstm_list, 0)
src_loss = nn.average_tensor_list(gpu_src_losses)
dst_loss = nn.average_tensor_list(gpu_dst_losses)
src_dst_loss_gv = nn.average_gv_list(gpu_src_dst_loss_gvs)
src_dst_loss_gv_op = self.src_dst_opt.get_update_op(
src_dst_loss_gv)
# Initializing training and view functions
def src_dst_train(warped_src, target_src, target_srcm, \
warped_dst, target_dst, target_dstm):
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: target_srcm,
self.warped_dst: warped_dst,
self.target_dst: target_dst,
self.target_dstm: target_dstm,
})
s = np.mean(s)
d = np.mean(d)
return s, d
self.src_dst_train = src_dst_train
def AE_view(warped_src, warped_dst):
return nn.tf_sess.run([
pred_src_src, pred_dst_dst, pred_dst_dstm, pred_src_dst,
pred_src_dstm
],
feed_dict={
self.warped_src: warped_src,
self.warped_dst: warped_dst
})
self.AE_view = AE_view
else:
# Initializing merge function
with tf.device(f'/GPU:0' if len(devices) != 0 else f'/CPU:0'):
gpu_dst_code = self.inter(self.encoder(self.warped_dst))
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(
gpu_dst_code)
_, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
def AE_merge(warped_dst):
return nn.tf_sess.run(
[gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm],
feed_dict={self.warped_dst: warped_dst})
self.AE_merge = AE_merge
# Loading/initializing all models/optimizers weights
for model, filename in io.progress_bar_generator(
self.model_filename_list, "Initializing models"):
if self.pretrain_just_disabled:
do_init = False
if model == self.inter:
do_init = True
else:
do_init = self.is_first_run()
if not do_init:
do_init = not model.load_weights(
self.get_strpath_storage_for_file(filename))
if do_init and self.pretrained_model_path is not None:
pretrained_filepath = self.pretrained_model_path / filename
if pretrained_filepath.exists():
do_init = not model.load_weights(pretrained_filepath)
if do_init:
model.init_weights()
# initializing sample generators
if self.is_training:
training_data_src_path = self.training_data_src_path if not self.pretrain else self.get_pretraining_data_path(
)
training_data_dst_path = self.training_data_dst_path if not self.pretrain else self.get_pretraining_data_path(
)
cpu_count = min(multiprocessing.cpu_count(), 8)
src_generators_count = cpu_count // 2
dst_generators_count = cpu_count // 2
self.set_training_data_generators([
SampleGeneratorFace(
training_data_src_path,
debug=self.is_debug(),
batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(
random_flip=True if self.pretrain else False),
output_sample_types=[{
'sample_type':
SampleProcessor.SampleType.FACE_IMAGE,
'warp':
True,
'transform':
True,
'channel_type':
SampleProcessor.ChannelType.BGR,
'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,
'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
}],
generators_count=src_generators_count),
SampleGeneratorFace(
training_data_dst_path,
debug=self.is_debug(),
batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(
random_flip=True if self.pretrain else False),
output_sample_types=[{
'sample_type':
SampleProcessor.SampleType.FACE_IMAGE,
'warp':
True,
'transform':
True,
'channel_type':
SampleProcessor.ChannelType.BGR,
'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,
'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
}],
generators_count=dst_generators_count)
])
self.last_samples = None
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
self.model_data_format = "NCHW" if self.is_exporting or (len(
device_config.devices) != 0 and not self.is_debug()) else "NHWC"
nn.initialize(data_format=self.model_data_format)
tf = nn.tf
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
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 nn.tf_default_device_name
bgr_shape = nn.get4Dshape(resolution, resolution, 3)
mask_shape = nn.get4Dshape(resolution, resolution, 1)
# Initializing model classes
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)
self.pretrain = self.options['pretrain']
if self.pretrain_just_disabled:
self.set_iter(0)
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices))
bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
self.set_batch_size(gpu_count * bs_per_gpu)
# Compute losses per GPU
gpu_pred_list = []
gpu_losses = []
gpu_loss_gvs = []
for gpu_id in range(gpu_count):
with tf.device(f'/{devices[gpu_id].tf_dev_type}:{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, :, :, :]
# process model tensors
gpu_pred_logits_t, gpu_pred_t = self.model.flow(
gpu_input_t, pretrain=self.pretrain)
gpu_pred_list.append(gpu_pred_t)
if self.pretrain:
# Structural loss
gpu_loss = tf.reduce_mean(
5 * nn.dssim(gpu_target_t,
gpu_pred_t,
max_val=1.0,
filter_size=int(resolution / 11.6)),
axis=[1])
gpu_loss += tf.reduce_mean(
5 * nn.dssim(gpu_target_t,
gpu_pred_t,
max_val=1.0,
filter_size=int(resolution / 23.2)),
axis=[1])
# Pixel loss
gpu_loss += tf.reduce_mean(
10 * tf.square(gpu_target_t - gpu_pred_t),
axis=[1, 2, 3])
else:
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]
gpu_loss_gvs += [
nn.gradients(gpu_loss, self.model.get_weights())
]
# 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 = 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
if self.pretrain:
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
})
return l
else:
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
})
return l
self.train = train
def view(input_np):
return nn.tf_sess.run([pred],
feed_dict={self.model.input_t: input_np})
self.view = view
# initializing sample generators
cpu_count = min(multiprocessing.cpu_count(), 8)
src_dst_generators_count = cpu_count // 2
src_generators_count = cpu_count // 2
dst_generators_count = cpu_count // 2
if self.pretrain:
pretrain_gen = SampleGeneratorFace(
self.get_pretraining_data_path(),
debug=self.is_debug(),
batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(
random_flip=True),
output_sample_types=[
{
'sample_type':
SampleProcessor.SampleType.FACE_IMAGE,
'warp': True,
'transform': True,
'channel_type': SampleProcessor.ChannelType.BGR,
'face_type': self.face_type,
'data_format': nn.data_format,
'resolution': resolution
},
{
'sample_type':
SampleProcessor.SampleType.FACE_IMAGE,
'warp': True,
'transform': True,
'channel_type': SampleProcessor.ChannelType.G,
'face_type': self.face_type,
'data_format': nn.data_format,
'resolution': resolution
},
],
uniform_yaw_distribution=False,
generators_count=cpu_count)
self.set_training_data_generators([pretrain_gen])
else:
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,
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)
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])
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
self.model_data_format = "NCHW" if len(
device_config.devices) != 0 and not self.is_debug() else "NHWC"
nn.initialize(data_format=self.model_data_format)
tf = nn.tf
conv_kernel_initializer = nn.initializers.ca()
class Downscale(nn.ModelBase):
def __init__(self,
in_ch,
out_ch,
kernel_size=3,
dilations=1,
use_activator=True,
*kwargs):
self.in_ch = in_ch
self.out_ch = out_ch
self.kernel_size = kernel_size
self.dilations = dilations
self.use_activator = use_activator
super().__init__(*kwargs)
def on_build(self, *args, **kwargs):
self.conv1 = nn.Conv2D(
self.in_ch,
self.out_ch,
kernel_size=self.kernel_size,
strides=2,
padding='SAME',
dilations=self.dilations,
kernel_initializer=conv_kernel_initializer)
def forward(self, x):
x = self.conv1(x)
if self.use_activator:
x = tf.nn.leaky_relu(x, 0.1)
return x
class Upscale(nn.ModelBase):
def on_build(self, in_ch, out_ch, kernel_size=3):
self.conv1 = nn.Conv2D(
in_ch,
out_ch * 4,
kernel_size=kernel_size,
padding='SAME',
kernel_initializer=conv_kernel_initializer)
def forward(self, x):
x = self.conv1(x)
x = tf.nn.leaky_relu(x, 0.1)
x = nn.depth_to_space(x, 2)
return x
class ResidualBlock(nn.ModelBase):
def on_build(self, ch, mod=1, kernel_size=3):
self.conv1 = nn.Conv2D(
ch,
ch * mod,
kernel_size=kernel_size,
padding='SAME',
kernel_initializer=conv_kernel_initializer)
self.conv2 = nn.Conv2D(
ch * mod,
ch,
kernel_size=kernel_size,
padding='SAME',
kernel_initializer=conv_kernel_initializer)
def forward(self, inp):
x = self.conv1(inp)
x = tf.nn.leaky_relu(x, 0.1)
x = self.conv2(x)
x = inp + x
x = tf.nn.leaky_relu(x, 0.1)
return x
class Encoder(nn.ModelBase):
def on_build(self, in_ch, e_ch):
self.conv1 = Downscale(in_ch, e_ch)
self.conv2 = Downscale(e_ch, e_ch * 2)
self.conv3 = Downscale(e_ch * 2, e_ch * 4)
self.conv4 = Downscale(e_ch * 4, e_ch * 8)
self.conv5 = Downscale(e_ch * 8, e_ch * 16)
self.conv6 = Downscale(e_ch * 16, e_ch * 32)
self.conv7 = Downscale(e_ch * 32, e_ch * 64)
self.res1 = ResidualBlock(e_ch)
self.res2 = ResidualBlock(e_ch * 2)
self.res3 = ResidualBlock(e_ch * 4)
self.res4 = ResidualBlock(e_ch * 8)
self.res5 = ResidualBlock(e_ch * 16)
self.res6 = ResidualBlock(e_ch * 32)
self.res7 = ResidualBlock(e_ch * 64)
def forward(self, inp):
x = self.conv1(inp)
x = self.res1(x)
x = self.conv2(x)
x = self.res2(x)
x = self.conv3(x)
x = self.res3(x)
x = self.conv4(x)
x = self.res4(x)
x = self.conv5(x)
x = self.res5(x)
x = self.conv6(x)
x = self.res6(x)
x = self.conv7(x)
x = self.res7(x)
return x
class Inter(nn.ModelBase):
def __init__(self, in_ch, ae_ch, **kwargs):
self.in_ch, self.ae_ch = in_ch, ae_ch
super().__init__(**kwargs)
def on_build(self):
in_ch, ae_ch = self.in_ch, self.ae_ch
self.dense_conv1 = nn.Conv2D(
in_ch,
64,
kernel_size=1,
padding='SAME',
kernel_initializer=conv_kernel_initializer)
self.dense_conv2 = nn.Conv2D(
64,
in_ch,
kernel_size=1,
padding='SAME',
kernel_initializer=conv_kernel_initializer)
self.conv7 = Upscale(in_ch, in_ch // 2)
self.conv6 = Upscale(in_ch // 2, in_ch // 4)
def forward(self, inp):
x = inp
x = self.dense_conv1(x)
x = self.dense_conv2(x)
x = self.conv7(x)
x = self.conv6(x)
return x
class Decoder(nn.ModelBase):
def on_build(self, in_ch):
self.upscale6 = Upscale(in_ch, in_ch // 2)
self.upscale5 = Upscale(in_ch // 2, in_ch // 4)
self.upscale4 = Upscale(in_ch // 4, in_ch // 8)
self.upscale3 = Upscale(in_ch // 8, in_ch // 16)
self.upscale2 = Upscale(in_ch // 16, in_ch // 32)
#self.upscale1 = Upscale(in_ch//32, in_ch//64)
self.out_conv = nn.Conv2D(
in_ch // 32,
3,
kernel_size=1,
padding='SAME',
kernel_initializer=conv_kernel_initializer)
self.res61 = ResidualBlock(in_ch // 2, mod=8)
self.res62 = ResidualBlock(in_ch // 2, mod=8)
self.res63 = ResidualBlock(in_ch // 2, mod=8)
self.res51 = ResidualBlock(in_ch // 4, mod=8)
self.res52 = ResidualBlock(in_ch // 4, mod=8)
self.res53 = ResidualBlock(in_ch // 4, mod=8)
self.res41 = ResidualBlock(in_ch // 8, mod=8)
self.res42 = ResidualBlock(in_ch // 8, mod=8)
self.res43 = ResidualBlock(in_ch // 8, mod=8)
self.res31 = ResidualBlock(in_ch // 16, mod=8)
self.res32 = ResidualBlock(in_ch // 16, mod=8)
self.res33 = ResidualBlock(in_ch // 16, mod=8)
self.res21 = ResidualBlock(in_ch // 32, mod=8)
self.res22 = ResidualBlock(in_ch // 32, mod=8)
self.res23 = ResidualBlock(in_ch // 32, mod=8)
m_ch = in_ch // 2
self.upscalem6 = Upscale(in_ch, m_ch // 2)
self.upscalem5 = Upscale(m_ch // 2, m_ch // 4)
self.upscalem4 = Upscale(m_ch // 4, m_ch // 8)
self.upscalem3 = Upscale(m_ch // 8, m_ch // 16)
self.upscalem2 = Upscale(m_ch // 16, m_ch // 32)
#self.upscalem1 = Upscale(m_ch//32, m_ch//64)
self.out_convm = nn.Conv2D(
m_ch // 32,
1,
kernel_size=1,
padding='SAME',
kernel_initializer=conv_kernel_initializer)
def forward(self, inp):
z = inp
x = self.upscale6(z)
x = self.res61(x)
x = self.res62(x)
x = self.res63(x)
x = self.upscale5(x)
x = self.res51(x)
x = self.res52(x)
x = self.res53(x)
x = self.upscale4(x)
x = self.res41(x)
x = self.res42(x)
x = self.res43(x)
x = self.upscale3(x)
x = self.res31(x)
x = self.res32(x)
x = self.res33(x)
x = self.upscale2(x)
x = self.res21(x)
x = self.res22(x)
x = self.res23(x)
#x = self.upscale1 (x)
y = self.upscalem6(z)
y = self.upscalem5(y)
y = self.upscalem4(y)
y = self.upscalem3(y)
y = self.upscalem2(y)
#y = self.upscalem1 (y)
return tf.nn.sigmoid(self.out_conv(x)), \
tf.nn.sigmoid(self.out_convm(y))
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
resolution = self.resolution = 128
ae_dims = 128
e_dims = 16
self.pretrain = False
self.pretrain_just_disabled = False
masked_training = True
models_opt_on_gpu = len(devices) >= 1 and all(
[dev.total_mem_gb >= 4 for dev in devices])
models_opt_device = '/GPU:0' if models_opt_on_gpu and self.is_training else '/CPU:0'
optimizer_vars_on_cpu = models_opt_device == '/CPU:0'
input_ch = 3
output_ch = 3
bgr_shape = nn.get4Dshape(resolution, resolution, input_ch)
mask_shape = nn.get4Dshape(resolution, resolution, 1)
self.model_filename_list = []
with tf.device('/CPU:0'):
#Place holders on CPU
self.warped_src = tf.placeholder(nn.floatx, bgr_shape)
self.warped_dst = tf.placeholder(nn.floatx, bgr_shape)
self.target_src = tf.placeholder(nn.floatx, bgr_shape)
self.target_dst = tf.placeholder(nn.floatx, bgr_shape)
self.target_srcm = tf.placeholder(nn.floatx, mask_shape)
self.target_dstm = tf.placeholder(nn.floatx, mask_shape)
# Initializing model classes
with tf.device(models_opt_device):
self.encoder = Encoder(in_ch=input_ch, e_ch=e_dims, name='encoder')
self.inter = Inter(in_ch=e_dims * 64, ae_ch=ae_dims, name='inter')
self.decoder_src = Decoder(in_ch=e_dims * 16, name='decoder_src')
self.decoder_dst = Decoder(in_ch=e_dims * 16, name='decoder_dst')
self.model_filename_list += [[self.encoder, 'encoder.npy'],
[self.inter, 'inter.npy'],
[self.decoder_src, 'decoder_src.npy'],
[self.decoder_dst, 'decoder_dst.npy']]
if self.is_training:
self.src_dst_trainable_weights = self.encoder.get_weights(
) + self.inter.get_weights() + self.decoder_src.get_weights(
) + self.decoder_dst.get_weights()
# Initialize optimizers
self.src_dst_opt = nn.RMSprop(lr=5e-5, name='src_dst_opt')
self.src_dst_opt.initialize_variables(
self.src_dst_trainable_weights,
vars_on_cpu=optimizer_vars_on_cpu)
self.model_filename_list += [(self.src_dst_opt,
'src_dst_opt.npy')]
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices))
bs_per_gpu = max(1, 4 // gpu_count)
self.set_batch_size(gpu_count * bs_per_gpu)
# Compute losses per GPU
gpu_pred_src_src_list = []
gpu_pred_dst_dst_list = []
gpu_pred_src_dst_list = []
gpu_pred_src_srcm_list = []
gpu_pred_dst_dstm_list = []
gpu_pred_src_dstm_list = []
gpu_src_losses = []
gpu_dst_losses = []
gpu_src_dst_loss_gvs = []
for gpu_id in range(gpu_count):
with tf.device(
f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0'):
batch_slice = slice(gpu_id * bs_per_gpu,
(gpu_id + 1) * bs_per_gpu)
with tf.device(f'/CPU:0'):
# slice on CPU, otherwise all batch data will be transfered to GPU first
gpu_warped_src = self.warped_src[batch_slice, :, :, :]
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 = self.target_srcm[
batch_slice, :, :, :]
gpu_target_dstm = self.target_dstm[
batch_slice, :, :, :]
# process model tensors
gpu_src_code = self.inter(self.encoder(gpu_warped_src))
gpu_dst_code = self.inter(self.encoder(gpu_warped_dst))
gpu_pred_src_src, gpu_pred_src_srcm = self.decoder_src(
gpu_src_code)
gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder_dst(
gpu_dst_code)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(
gpu_dst_code)
gpu_pred_src_src_list.append(gpu_pred_src_src)
gpu_pred_dst_dst_list.append(gpu_pred_dst_dst)
gpu_pred_src_dst_list.append(gpu_pred_src_dst)
gpu_pred_src_srcm_list.append(gpu_pred_src_srcm)
gpu_pred_dst_dstm_list.append(gpu_pred_dst_dstm)
gpu_pred_src_dstm_list.append(gpu_pred_src_dstm)
gpu_target_srcm_blur = nn.gaussian_blur(
gpu_target_srcm, max(1, resolution // 32))
gpu_target_dstm_blur = nn.gaussian_blur(
gpu_target_dstm, max(1, resolution // 32))
gpu_target_dst_masked = gpu_target_dst * gpu_target_dstm_blur
gpu_target_dst_anti_masked = gpu_target_dst * (
1.0 - gpu_target_dstm_blur)
gpu_target_src_masked_opt = gpu_target_src * gpu_target_srcm_blur if masked_training else gpu_target_src
gpu_target_dst_masked_opt = gpu_target_dst_masked if masked_training else gpu_target_dst
gpu_pred_src_src_masked_opt = gpu_pred_src_src * gpu_target_srcm_blur if masked_training else gpu_pred_src_src
gpu_pred_dst_dst_masked_opt = gpu_pred_dst_dst * gpu_target_dstm_blur if masked_training else gpu_pred_dst_dst
gpu_psd_target_dst_masked = gpu_pred_src_dst * gpu_target_dstm_blur
gpu_psd_target_dst_anti_masked = gpu_pred_src_dst * (
1.0 - gpu_target_dstm_blur)
gpu_src_loss = tf.reduce_mean(
10 * nn.dssim(gpu_target_src_masked_opt,
gpu_pred_src_src_masked_opt,
max_val=1.0,
filter_size=int(resolution / 11.6)),
axis=[1])
gpu_src_loss += tf.reduce_mean(
10 * tf.square(gpu_target_src_masked_opt -
gpu_pred_src_src_masked_opt),
axis=[1, 2, 3])
gpu_src_loss += tf.reduce_mean(
10 * tf.square(gpu_target_srcm - gpu_pred_src_srcm),
axis=[1, 2, 3])
gpu_dst_loss = tf.reduce_mean(
10 * nn.dssim(gpu_target_dst_masked_opt,
gpu_pred_dst_dst_masked_opt,
max_val=1.0,
filter_size=int(resolution / 11.6)),
axis=[1])
gpu_dst_loss += tf.reduce_mean(
10 * tf.square(gpu_target_dst_masked_opt -
gpu_pred_dst_dst_masked_opt),
axis=[1, 2, 3])
gpu_dst_loss += tf.reduce_mean(
10 * tf.square(gpu_target_dstm - gpu_pred_dst_dstm),
axis=[1, 2, 3])
gpu_src_losses += [gpu_src_loss]
gpu_dst_losses += [gpu_dst_loss]
gpu_G_loss = gpu_src_loss + gpu_dst_loss
gpu_src_dst_loss_gvs += [
nn.gradients(gpu_G_loss,
self.src_dst_trainable_weights)
]
# Average losses and gradients, and create optimizer update ops
with tf.device(models_opt_device):
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)
pred_src_srcm = nn.concat(gpu_pred_src_srcm_list, 0)
pred_dst_dstm = nn.concat(gpu_pred_dst_dstm_list, 0)
pred_src_dstm = nn.concat(gpu_pred_src_dstm_list, 0)
src_loss = nn.average_tensor_list(gpu_src_losses)
dst_loss = nn.average_tensor_list(gpu_dst_losses)
src_dst_loss_gv = nn.average_gv_list(gpu_src_dst_loss_gvs)
src_dst_loss_gv_op = self.src_dst_opt.get_update_op(
src_dst_loss_gv)
# Initializing training and view functions
def src_dst_train(warped_src, target_src, target_srcm, \
warped_dst, target_dst, target_dstm):
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: target_srcm,
self.warped_dst: warped_dst,
self.target_dst: target_dst,
self.target_dstm: target_dstm,
})
s = np.mean(s)
d = np.mean(d)
return s, d
self.src_dst_train = src_dst_train
def AE_view(warped_src, warped_dst):
return nn.tf_sess.run([
pred_src_src, pred_dst_dst, pred_dst_dstm, pred_src_dst,
pred_src_dstm
],
feed_dict={
self.warped_src: warped_src,
self.warped_dst: warped_dst
})
self.AE_view = AE_view
else:
# Initializing merge function
with tf.device(f'/GPU:0' if len(devices) != 0 else f'/CPU:0'):
gpu_dst_code = self.inter(self.encoder(self.warped_dst))
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(
gpu_dst_code)
_, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
def AE_merge(warped_dst):
return nn.tf_sess.run(
[gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm],
feed_dict={self.warped_dst: warped_dst})
self.AE_merge = AE_merge
# Loading/initializing all models/optimizers weights
for model, filename in io.progress_bar_generator(
self.model_filename_list, "Initializing models"):
if self.pretrain_just_disabled:
do_init = False
if model == self.inter:
do_init = True
else:
do_init = self.is_first_run()
if not do_init:
do_init = not model.load_weights(
self.get_strpath_storage_for_file(filename))
if do_init and self.pretrained_model_path is not None:
pretrained_filepath = self.pretrained_model_path / filename
if pretrained_filepath.exists():
do_init = not model.load_weights(pretrained_filepath)
if do_init:
model.init_weights()
# initializing sample generators
if self.is_training:
t = SampleProcessor.Types
face_type = t.FACE_TYPE_FULL
training_data_src_path = self.training_data_src_path if not self.pretrain else self.get_pretraining_data_path(
)
training_data_dst_path = self.training_data_dst_path if not self.pretrain else self.get_pretraining_data_path(
)
cpu_count = min(multiprocessing.cpu_count(), 8)
src_generators_count = cpu_count // 2
dst_generators_count = cpu_count // 2
self.set_training_data_generators([
SampleGeneratorFace(
training_data_src_path,
debug=self.is_debug(),
batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(
random_flip=True if self.pretrain else False),
output_sample_types=[{
'types':
(t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR),
'data_format':
nn.data_format,
'resolution':
resolution,
}, {
'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR),
'data_format':
nn.data_format,
'resolution':
resolution,
}, {
'types': (t.IMG_TRANSFORMED, face_type,
t.MODE_FACE_MASK_ALL_HULL),
'data_format':
nn.data_format,
'resolution':
resolution
}],
generators_count=src_generators_count),
SampleGeneratorFace(
training_data_dst_path,
debug=self.is_debug(),
batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(
random_flip=True if self.pretrain else False),
output_sample_types=[{
'types':
(t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR),
'data_format':
nn.data_format,
'resolution':
resolution
}, {
'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR),
'data_format':
nn.data_format,
'resolution':
resolution
}, {
'types': (t.IMG_TRANSFORMED, face_type,
t.MODE_FACE_MASK_ALL_HULL),
'data_format':
nn.data_format,
'resolution':
resolution
}],
generators_count=dst_generators_count)
])
self.last_samples = None
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
self.model_data_format = "NCHW" if len(
devices) != 0 and not self.is_debug() else "NHWC"
nn.initialize(data_format=self.model_data_format)
tf = nn.tf
resolution = self.resolution = 96
self.face_type = FaceType.FULL
ae_dims = 256
e_dims = 64
d_dims = 64
self.pretrain = False
self.pretrain_just_disabled = False
masked_training = True
models_opt_on_gpu = len(devices) >= 1 and all(
[dev.total_mem_gb >= 4 for dev in devices])
models_opt_device = '/GPU:0' if models_opt_on_gpu and self.is_training else '/CPU:0'
optimizer_vars_on_cpu = models_opt_device == '/CPU:0'
input_ch = 3
bgr_shape = nn.get4Dshape(resolution, resolution, input_ch)
mask_shape = nn.get4Dshape(resolution, resolution, 1)
self.model_filename_list = []
kernel_initializer = tf.initializers.glorot_uniform()
class Upscale(nn.ModelBase):
def on_build(self, in_ch, out_ch, kernel_size=3):
self.conv1 = nn.Conv2D(in_ch,
out_ch * 4,
kernel_size=kernel_size,
padding='SAME',
kernel_initializer=kernel_initializer)
def forward(self, x):
x = self.conv1(x)
x = tf.nn.leaky_relu(x, 0.1)
x = nn.depth_to_space(x, 2)
return x
class ResidualBlock(nn.ModelBase):
def on_build(self, ch, kernel_size=3):
self.conv1 = nn.Conv2D(ch,
ch,
kernel_size=kernel_size,
padding='SAME',
kernel_initializer=kernel_initializer)
self.conv2 = nn.Conv2D(ch,
ch,
kernel_size=kernel_size,
padding='SAME',
kernel_initializer=kernel_initializer)
def forward(self, inp):
x = self.conv1(inp)
x = tf.nn.leaky_relu(x, 0.2)
x = self.conv2(x)
x = tf.nn.leaky_relu(inp + x, 0.2)
return x
class Encoder(nn.ModelBase):
def on_build(self, in_ch, e_ch):
self.down11 = nn.Conv2D(in_ch,
e_ch,
kernel_size=3,
strides=1,
padding='SAME',
kernel_initializer=kernel_initializer)
self.down12 = nn.Conv2D(e_ch,
e_ch,
kernel_size=3,
strides=1,
padding='SAME',
kernel_initializer=kernel_initializer)
self.down21 = nn.Conv2D(e_ch,
e_ch * 2,
kernel_size=3,
strides=1,
padding='SAME',
kernel_initializer=kernel_initializer)
self.down22 = nn.Conv2D(e_ch * 2,
e_ch * 2,
kernel_size=3,
strides=1,
padding='SAME',
kernel_initializer=kernel_initializer)
self.down31 = nn.Conv2D(e_ch * 2,
e_ch * 4,
kernel_size=3,
strides=1,
padding='SAME',
kernel_initializer=kernel_initializer)
self.down32 = nn.Conv2D(e_ch * 4,
e_ch * 4,
kernel_size=3,
strides=1,
padding='SAME',
kernel_initializer=kernel_initializer)
self.down33 = nn.Conv2D(e_ch * 4,
e_ch * 4,
kernel_size=3,
strides=1,
padding='SAME',
kernel_initializer=kernel_initializer)
self.down41 = nn.Conv2D(e_ch * 4,
e_ch * 8,
kernel_size=3,
strides=1,
padding='SAME',
kernel_initializer=kernel_initializer)
self.down42 = nn.Conv2D(e_ch * 8,
e_ch * 8,
kernel_size=3,
strides=1,
padding='SAME',
kernel_initializer=kernel_initializer)
self.down43 = nn.Conv2D(e_ch * 8,
e_ch * 8,
kernel_size=3,
strides=1,
padding='SAME',
kernel_initializer=kernel_initializer)
self.down51 = nn.Conv2D(e_ch * 8,
e_ch * 8,
kernel_size=3,
strides=1,
padding='SAME',
kernel_initializer=kernel_initializer)
self.down52 = nn.Conv2D(e_ch * 8,
e_ch * 8,
kernel_size=3,
strides=1,
padding='SAME',
kernel_initializer=kernel_initializer)
self.down53 = nn.Conv2D(e_ch * 8,
e_ch * 8,
kernel_size=3,
strides=1,
padding='SAME',
kernel_initializer=kernel_initializer)
def forward(self, inp):
x = inp
x = self.down11(x)
x = self.down12(x)
x = nn.max_pool(x)
x = self.down21(x)
x = self.down22(x)
x = nn.max_pool(x)
x = self.down31(x)
x = self.down32(x)
x = self.down33(x)
x = nn.max_pool(x)
x = self.down41(x)
x = self.down42(x)
x = self.down43(x)
x = nn.max_pool(x)
x = self.down51(x)
x = self.down52(x)
x = self.down53(x)
x = nn.max_pool(x)
x = nn.flatten(x)
return x
class Downscale(nn.ModelBase):
def __init__(self,
in_ch,
out_ch,
kernel_size=5,
dilations=1,
subpixel=True,
use_activator=True,
*kwargs):
self.in_ch = in_ch
self.out_ch = out_ch
self.kernel_size = kernel_size
self.dilations = dilations
self.subpixel = subpixel
self.use_activator = use_activator
super().__init__(*kwargs)
def on_build(self, *args, **kwargs):
self.conv1 = nn.Conv2D(self.in_ch,
self.out_ch //
(4 if self.subpixel else 1),
kernel_size=self.kernel_size,
strides=1 if self.subpixel else 2,
padding='SAME',
dilations=self.dilations,
kernel_initializer=kernel_initializer)
def forward(self, x):
x = self.conv1(x)
if self.subpixel:
x = nn.space_to_depth(x, 2)
if self.use_activator:
x = tf.nn.leaky_relu(x, 0.1)
return x
def get_out_ch(self):
return (self.out_ch // 4) * 4
class DownscaleBlock(nn.ModelBase):
def on_build(self,
in_ch,
ch,
n_downscales,
kernel_size,
dilations=1,
subpixel=True):
self.downs = []
last_ch = in_ch
for i in range(n_downscales):
cur_ch = ch * (min(2**i, 8))
self.downs.append(
Downscale(last_ch,
cur_ch,
kernel_size=kernel_size,
dilations=dilations,
subpixel=subpixel))
last_ch = self.downs[-1].get_out_ch()
def forward(self, inp):
x = inp
for down in self.downs:
x = down(x)
return x
class Upscale(nn.ModelBase):
def on_build(self, in_ch, out_ch, kernel_size=3):
self.conv1 = nn.Conv2D(in_ch,
out_ch * 4,
kernel_size=kernel_size,
padding='SAME',
kernel_initializer=kernel_initializer)
def forward(self, x):
x = self.conv1(x)
x = tf.nn.leaky_relu(x, 0.1)
x = nn.depth_to_space(x, 2)
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,
dilations=1,
subpixel=False)
def forward(self, inp):
x = nn.flatten(self.down1(inp))
return x
class Branch(nn.ModelBase):
def on_build(self, in_ch, ae_ch):
self.dense1 = nn.Dense(in_ch, ae_ch)
def forward(self, inp):
x = self.dense1(inp)
return x
class Classifier(nn.ModelBase):
def on_build(self, in_ch, n_classes):
self.dense1 = nn.Dense(in_ch, 4096)
self.dense2 = nn.Dense(4096, 4096)
self.pitch_dense = nn.Dense(4096, n_classes)
self.yaw_dense = nn.Dense(4096, n_classes)
def forward(self, inp):
x = inp
x = self.dense1(x)
x = self.dense2(x)
return self.pitch_dense(x), self.yaw_dense(x)
lowest_dense_res = resolution // 16
class Inter(nn.ModelBase):
def on_build(self, in_ch, ae_out_ch):
self.ae_out_ch = ae_out_ch
self.dense2 = nn.Dense(
in_ch, lowest_dense_res * lowest_dense_res * ae_out_ch)
self.upscale1 = Upscale(ae_out_ch, ae_out_ch)
def forward(self, inp):
x = inp
x = self.dense2(x)
x = nn.reshape_4D(x, lowest_dense_res, lowest_dense_res,
self.ae_out_ch)
x = self.upscale1(x)
return x
def get_out_ch(self):
return self.ae_out_ch
class Decoder(nn.ModelBase):
def on_build(self, in_ch, d_ch, d_mask_ch):
self.upscale0 = Upscale(in_ch, d_ch * 8, kernel_size=3)
self.upscale1 = Upscale(d_ch * 8, d_ch * 4, kernel_size=3)
self.upscale2 = Upscale(d_ch * 4, d_ch * 2, kernel_size=3)
self.res0 = ResidualBlock(d_ch * 8, kernel_size=3)
self.res1 = ResidualBlock(d_ch * 4, kernel_size=3)
self.res2 = ResidualBlock(d_ch * 2, kernel_size=3)
self.out_conv = nn.Conv2D(
d_ch * 2,
3,
kernel_size=1,
padding='SAME',
kernel_initializer=kernel_initializer)
def forward(self, inp):
z = inp
x = self.upscale0(z)
x = self.res0(x)
x = self.upscale1(x)
x = self.res1(x)
x = self.upscale2(x)
x = self.res2(x)
return tf.nn.sigmoid(self.out_conv(x))
n_pyr_degs = self.n_pyr_degs = 3
n_pyr_classes = self.n_pyr_classes = 180 // self.n_pyr_degs
with tf.device('/CPU:0'):
#Place holders on CPU
self.warped_src = tf.placeholder(nn.floatx, bgr_shape)
self.target_src = tf.placeholder(nn.floatx, bgr_shape)
self.target_dst = tf.placeholder(nn.floatx, bgr_shape)
self.pitches_vector = tf.placeholder(nn.floatx,
(None, n_pyr_classes))
self.yaws_vector = tf.placeholder(nn.floatx, (None, n_pyr_classes))
# Initializing model classes
with tf.device(models_opt_device):
self.encoder = Encoder(in_ch=input_ch, e_ch=e_dims, name='encoder')
encoder_out_ch = self.encoder.compute_output_channels(
(nn.floatx, bgr_shape))
self.bT = Branch(in_ch=encoder_out_ch, ae_ch=ae_dims, name='bT')
self.bP = Branch(in_ch=encoder_out_ch, ae_ch=ae_dims, name='bP')
self.bTC = Classifier(in_ch=ae_dims,
n_classes=self.n_pyr_classes,
name='bTC')
self.bPC = Classifier(in_ch=ae_dims,
n_classes=self.n_pyr_classes,
name='bPC')
self.inter = Inter(in_ch=ae_dims * 2,
ae_out_ch=ae_dims * 2,
name='inter')
self.decoder = Decoder(in_ch=ae_dims * 2,
d_ch=d_dims,
d_mask_ch=d_dims,
name='decoder')
self.model_filename_list += [[self.encoder, 'encoder.npy'],
[self.bT, 'bT.npy'],
[self.bTC, 'bTC.npy'],
[self.bP, 'bP.npy'],
[self.bPC, 'bPC.npy'],
[self.inter, 'inter.npy'],
[self.decoder, 'decoder.npy']]
if self.is_training:
self.all_trainable_weights = self.encoder.get_weights() + \
self.bT.get_weights() +\
self.bTC.get_weights() +\
self.bP.get_weights() +\
self.bPC.get_weights() +\
self.inter.get_weights() +\
self.decoder.get_weights()
# Initialize optimizers
self.src_dst_opt = nn.RMSprop(lr=5e-5, name='src_dst_opt')
self.src_dst_opt.initialize_variables(
self.all_trainable_weights,
vars_on_cpu=optimizer_vars_on_cpu)
self.model_filename_list += [(self.src_dst_opt,
'src_dst_opt.npy')]
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices))
bs_per_gpu = max(1, 32 // gpu_count)
self.set_batch_size(gpu_count * bs_per_gpu)
# Compute losses per GPU
gpu_pred_src_list = []
gpu_pred_dst_list = []
gpu_A_losses = []
gpu_B_losses = []
gpu_C_losses = []
gpu_D_losses = []
gpu_A_loss_gvs = []
gpu_B_loss_gvs = []
gpu_C_loss_gvs = []
gpu_D_loss_gvs = []
for gpu_id in range(gpu_count):
with tf.device(
f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0'):
batch_slice = slice(gpu_id * bs_per_gpu,
(gpu_id + 1) * bs_per_gpu)
with tf.device(f'/CPU:0'):
# slice on CPU, otherwise all batch data will be transfered to GPU first
gpu_warped_src = self.warped_src[batch_slice, :, :, :]
gpu_target_src = self.target_src[batch_slice, :, :, :]
gpu_target_dst = self.target_dst[batch_slice, :, :, :]
gpu_pitches_vector = self.pitches_vector[
batch_slice, :]
gpu_yaws_vector = self.yaws_vector[batch_slice, :]
# process model tensors
gpu_src_enc_code = self.encoder(gpu_warped_src)
gpu_dst_enc_code = self.encoder(gpu_target_dst)
gpu_src_bT_code = self.bT(gpu_src_enc_code)
gpu_src_bT_code_ng = tf.stop_gradient(gpu_src_bT_code)
gpu_src_T_pitch, gpu_src_T_yaw = self.bTC(gpu_src_bT_code)
gpu_dst_bT_code = self.bT(gpu_dst_enc_code)
gpu_src_bP_code = self.bP(gpu_src_enc_code)
gpu_src_P_pitch, gpu_src_P_yaw = self.bPC(gpu_src_bP_code)
def crossentropy(target, output):
output = tf.nn.softmax(output)
output = tf.clip_by_value(output, 1e-7, 1 - 1e-7)
return tf.reduce_sum(target * -tf.log(output),
axis=-1,
keepdims=False)
def negative_crossentropy(n_classes, output):
output = tf.nn.softmax(output)
output = tf.clip_by_value(output, 1e-7, 1 - 1e-7)
return (1.0 / n_classes) * tf.reduce_sum(
tf.log(output), axis=-1, keepdims=False)
gpu_src_bT_code_n = gpu_src_bT_code_ng + tf.random.normal(
tf.shape(gpu_src_bT_code_ng))
gpu_src_bP_code_n = gpu_src_bP_code + tf.random.normal(
tf.shape(gpu_src_bP_code))
gpu_pred_src = self.decoder(
self.inter(
tf.concat([gpu_src_bT_code_ng, gpu_src_bP_code],
axis=-1)))
gpu_pred_src_n = self.decoder(
self.inter(
tf.concat([gpu_src_bT_code_n, gpu_src_bP_code_n],
axis=-1)))
gpu_pred_dst = self.decoder(
self.inter(
tf.concat([gpu_dst_bT_code, gpu_src_bP_code],
axis=-1)))
gpu_A_loss = 1.0*crossentropy(gpu_pitches_vector, gpu_src_T_pitch ) + \
1.0*crossentropy(gpu_yaws_vector, gpu_src_T_yaw )
gpu_B_loss = 0.1*crossentropy(gpu_pitches_vector, gpu_src_P_pitch ) + \
0.1*crossentropy(gpu_yaws_vector, gpu_src_P_yaw )
gpu_C_loss = 0.1*negative_crossentropy( n_pyr_classes, gpu_src_P_pitch ) + \
0.1*negative_crossentropy( n_pyr_classes, gpu_src_P_yaw )
gpu_D_loss = 0.0000001*(\
0.5*tf.reduce_sum(tf.square(gpu_target_src-gpu_pred_src), axis=[1,2,3]) + \
0.5*tf.reduce_sum(tf.square(gpu_target_src-gpu_pred_src_n), axis=[1,2,3]) )
gpu_pred_src_list.append(gpu_pred_src)
gpu_pred_dst_list.append(gpu_pred_dst)
gpu_A_losses += [gpu_A_loss]
gpu_B_losses += [gpu_B_loss]
gpu_C_losses += [gpu_C_loss]
gpu_D_losses += [gpu_D_loss]
A_weights = self.encoder.get_weights(
) + self.bT.get_weights() + self.bTC.get_weights()
B_weights = self.bPC.get_weights()
C_weights = self.encoder.get_weights(
) + self.bP.get_weights()
D_weights = self.inter.get_weights(
) + self.decoder.get_weights()
gpu_A_loss_gvs += [nn.gradients(gpu_A_loss, A_weights)]
gpu_B_loss_gvs += [nn.gradients(gpu_B_loss, B_weights)]
gpu_C_loss_gvs += [nn.gradients(gpu_C_loss, C_weights)]
gpu_D_loss_gvs += [nn.gradients(gpu_D_loss, D_weights)]
# Average losses and gradients, and create optimizer update ops
with tf.device(models_opt_device):
pred_src = nn.concat(gpu_pred_src_list, 0)
pred_dst = nn.concat(gpu_pred_dst_list, 0)
A_loss = nn.average_tensor_list(gpu_A_losses)
B_loss = nn.average_tensor_list(gpu_B_losses)
C_loss = nn.average_tensor_list(gpu_C_losses)
D_loss = nn.average_tensor_list(gpu_D_losses)
A_loss_gv = nn.average_gv_list(gpu_A_loss_gvs)
B_loss_gv = nn.average_gv_list(gpu_B_loss_gvs)
C_loss_gv = nn.average_gv_list(gpu_C_loss_gvs)
D_loss_gv = nn.average_gv_list(gpu_D_loss_gvs)
A_loss_gv_op = self.src_dst_opt.get_update_op(A_loss_gv)
B_loss_gv_op = self.src_dst_opt.get_update_op(B_loss_gv)
C_loss_gv_op = self.src_dst_opt.get_update_op(C_loss_gv)
D_loss_gv_op = self.src_dst_opt.get_update_op(D_loss_gv)
# Initializing training and view functions
def A_train(warped_src, target_src, pitches_vector, yaws_vector):
l, _ = nn.tf_sess.run(
[A_loss, A_loss_gv_op],
feed_dict={
self.warped_src: warped_src,
self.target_src: target_src,
self.pitches_vector: pitches_vector,
self.yaws_vector: yaws_vector
})
return np.mean(l)
self.A_train = A_train
def B_train(warped_src, target_src, pitches_vector, yaws_vector):
l, _ = nn.tf_sess.run(
[B_loss, B_loss_gv_op],
feed_dict={
self.warped_src: warped_src,
self.target_src: target_src,
self.pitches_vector: pitches_vector,
self.yaws_vector: yaws_vector
})
return np.mean(l)
self.B_train = B_train
def C_train(warped_src, target_src, pitches_vector, yaws_vector):
l, _ = nn.tf_sess.run(
[C_loss, C_loss_gv_op],
feed_dict={
self.warped_src: warped_src,
self.target_src: target_src,
self.pitches_vector: pitches_vector,
self.yaws_vector: yaws_vector
})
return np.mean(l)
self.C_train = C_train
def D_train(warped_src, target_src, pitches_vector, yaws_vector):
l, _ = nn.tf_sess.run(
[D_loss, D_loss_gv_op],
feed_dict={
self.warped_src: warped_src,
self.target_src: target_src,
self.pitches_vector: pitches_vector,
self.yaws_vector: yaws_vector
})
return np.mean(l)
self.D_train = D_train
def AE_view(warped_src):
return nn.tf_sess.run([pred_src],
feed_dict={self.warped_src: warped_src})
self.AE_view = AE_view
def AE_view2(warped_src, target_dst):
return nn.tf_sess.run([pred_dst],
feed_dict={
self.warped_src: warped_src,
self.target_dst: target_dst
})
self.AE_view2 = AE_view2
else:
# Initializing merge function
with tf.device(f'/GPU:0' if len(devices) != 0 else f'/CPU:0'):
gpu_dst_code = self.inter(self.encoder(self.warped_dst))
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(
gpu_dst_code)
_, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
def AE_merge(warped_dst):
return nn.tf_sess.run(
[gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm],
feed_dict={self.warped_dst: warped_dst})
self.AE_merge = AE_merge
# Loading/initializing all models/optimizers weights
for model, filename in io.progress_bar_generator(
self.model_filename_list, "Initializing models"):
if self.pretrain_just_disabled:
do_init = False
if model == self.inter:
do_init = True
else:
do_init = self.is_first_run()
if not do_init:
do_init = not model.load_weights(
self.get_strpath_storage_for_file(filename))
if do_init and self.pretrained_model_path is not None:
pretrained_filepath = self.pretrained_model_path / filename
if pretrained_filepath.exists():
do_init = not model.load_weights(pretrained_filepath)
if do_init:
model.init_weights()
# initializing sample generators
if self.is_training:
training_data_src_path = self.training_data_src_path if not self.pretrain else self.get_pretraining_data_path(
)
training_data_dst_path = self.training_data_dst_path if not self.pretrain else self.get_pretraining_data_path(
)
cpu_count = min(multiprocessing.cpu_count(), 8)
src_generators_count = cpu_count // 2
dst_generators_count = cpu_count // 2
self.set_training_data_generators([
SampleGeneratorFace(
training_data_src_path,
debug=self.is_debug(),
batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(
random_flip=True if self.pretrain else False),
output_sample_types=[
{
'sample_type':
SampleProcessor.SampleType.FACE_IMAGE,
'warp': True,
'transform': True,
'channel_type': SampleProcessor.ChannelType.BGR,
'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,
'face_type': self.face_type,
'data_format': nn.data_format,
'resolution': resolution
},
{
'sample_type':
SampleProcessor.SampleType.PITCH_YAW_ROLL_SIGMOID,
'resolution': resolution
},
],
generators_count=src_generators_count),
SampleGeneratorFace(
training_data_dst_path,
debug=self.is_debug(),
batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(
random_flip=True if self.pretrain else False),
output_sample_types=[
{
'sample_type':
SampleProcessor.SampleType.FACE_IMAGE,
'warp': True,
'transform': True,
'channel_type': SampleProcessor.ChannelType.BGR,
'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,
'face_type': self.face_type,
'data_format': nn.data_format,
'resolution': resolution
},
{
'sample_type':
SampleProcessor.SampleType.PITCH_YAW_ROLL_SIGMOID,
'resolution': resolution
},
],
generators_count=dst_generators_count)
])
self.last_samples = None
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
nn.initialize(data_format="NHWC")
tf = nn.tf
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
self.resolution = resolution = 256 #self.options['resolution']
#self.face_type = {'h' : FaceType.HALF,
# 'mf' : FaceType.MID_FULL,
# 'f' : FaceType.FULL,
# 'wf' : FaceType.WHOLE_FACE}[ self.options['face_type'] ]
self.face_type = FaceType.FULL
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 = TernausNet(f'{self.model_name}_FANSeg',
resolution,
FaceType.toString(self.face_type),
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)
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices))
bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
self.set_batch_size(gpu_count * bs_per_gpu)
# Compute losses per GPU
gpu_pred_list = []
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'):
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, :, :, :]
# process model tensors
gpu_pred_logits_t, gpu_pred_t = self.model.net(
[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]
gpu_loss_gvs += [
nn.tf_gradients(gpu_loss, self.model.net_weights)
]
# Average losses and gradients, and create optimizer update ops
with tf.device(models_opt_device):
pred = nn.tf_concat(gpu_pred_list, 0)
loss = tf.reduce_mean(gpu_losses)
loss_gv_op = self.model.opt.get_update_op(
nn.tf_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
})
return l
self.train = train
def view(input_np):
return nn.tf_sess.run([pred],
feed_dict={self.model.input_t: input_np})
self.view = view
# initializing sample generators
training_data_src_path = self.training_data_src_path
training_data_dst_path = self.training_data_dst_path
cpu_count = min(multiprocessing.cpu_count(), 8)
src_generators_count = cpu_count // 2
dst_generators_count = cpu_count // 2
src_generators_count = int(src_generators_count * 1.5)
src_generator = SampleGeneratorFace(
training_data_src_path,
random_ct_samples_path=training_data_src_path,
debug=self.is_debug(),
batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(
random_flip=True),
output_sample_types=[
{
'sample_type': SampleProcessor.SampleType.FACE_IMAGE,
'ct_mode': 'lct',
'warp': True,
'transform': True,
'channel_type': SampleProcessor.ChannelType.BGR,
'face_type': self.face_type,
'motion_blur': (25, 5),
'gaussian_blur': (25, 5),
'data_format': nn.data_format,
'resolution': resolution
},
{
'sample_type': SampleProcessor.SampleType.FACE_MASK,
'warp': True,
'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
},
],
generators_count=src_generators_count)
dst_generator = SampleGeneratorFace(
training_data_dst_path,
debug=self.is_debug(),
batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(
random_flip=True),
output_sample_types=[
{
'sample_type': SampleProcessor.SampleType.FACE_IMAGE,
'warp': False,
'transform': True,
'channel_type': SampleProcessor.ChannelType.BGR,
'face_type': self.face_type,
'motion_blur': (25, 5),
'gaussian_blur': (25, 5),
'data_format': nn.data_format,
'resolution': resolution
},
],
generators_count=dst_generators_count,
raise_on_no_data=False)
if not dst_generator.is_initialized():
io.log_info(
f"\nTo view the model on unseen faces, place any aligned faces in {training_data_dst_path}.\n"
)
self.set_training_data_generators([src_generator, dst_generator])
def __init__(self,
name,
resolution=256,
load_weights=True,
weights_file_root=None,
training=False,
place_model_on_cpu=False,
run_on_cpu=False,
optimizer=None,
data_format="NHWC",
raise_on_no_model_files=False):
self.resolution = resolution
self.weights_file_root = Path(
weights_file_root) if weights_file_root is not None else Path(
__file__).parent
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))
self.target_t = tf.placeholder(
nn.floatx, nn.get4Dshape(resolution, resolution, 1))
# Initializing model classes
with tf.device(
'/CPU:0' if place_model_on_cpu else nn.tf_default_device_name):
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']]
if not training:
with tf.device(
'/CPU:0' if run_on_cpu else nn.tf_default_device_name):
_, pred = self.model(self.input_t)
def net_run(input_np):
return nn.tf_sess.run([pred],
feed_dict={self.input_t: input_np})[0]
self.net_run = net_run
self.initialized = True
# Loading/initializing all models/optimizers weights
for model, filename in self.model_filename_list:
do_init = not load_weights
if not do_init:
model_file_path = self.weights_file_root / filename
do_init = not model.load_weights(model_file_path)
if do_init:
if raise_on_no_model_files:
raise Exception(f'{model_file_path} does not exists.')
if not training:
self.initialized = False
break
if do_init:
model.init_weights()
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
self.model_data_format = "NCHW" if len(device_config.devices) != 0 and not self.is_debug() else "NHWC"
nn.initialize(data_format=self.model_data_format)
tf = nn.tf
conv_kernel_initializer = nn.initializers.ca()
class Downscale(nn.ModelBase):
def __init__(self, in_ch, out_ch, kernel_size=5, dilations=1, subpixel=True, use_activator=True, *kwargs ):
self.in_ch = in_ch
self.out_ch = out_ch
self.kernel_size = kernel_size
self.dilations = dilations
self.subpixel = subpixel
self.use_activator = use_activator
super().__init__(*kwargs)
def on_build(self, *args, **kwargs ):
self.conv1 = nn.Conv2D( self.in_ch,
self.out_ch // (4 if self.subpixel else 1),
kernel_size=self.kernel_size,
strides=1 if self.subpixel else 2,
padding='SAME', dilations=self.dilations, kernel_initializer=conv_kernel_initializer )
def forward(self, x):
x = self.conv1(x)
if self.subpixel:
x = nn.tf_space_to_depth(x, 2)
if self.use_activator:
x = nn.tf_gelu(x)
return x
def get_out_ch(self):
return (self.out_ch // 4) * 4
class DownscaleBlock(nn.ModelBase):
def on_build(self, in_ch, ch, n_downscales, kernel_size, dilations=1, subpixel=True):
self.downs = []
last_ch = in_ch
for i in range(n_downscales):
cur_ch = ch*( min(2**i, 8) )
self.downs.append ( Downscale(last_ch, cur_ch, kernel_size=kernel_size, dilations=dilations, subpixel=subpixel) )
last_ch = self.downs[-1].get_out_ch()
def forward(self, inp):
x = inp
for down in self.downs:
x = down(x)
return x
class Upscale(nn.ModelBase):
def on_build(self, in_ch, out_ch, kernel_size=3 ):
self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
def forward(self, x):
x = self.conv1(x)
x = nn.tf_gelu(x)
x = nn.tf_depth_to_space(x, 2)
return x
class ResidualBlock(nn.ModelBase):
def on_build(self, ch, kernel_size=3 ):
self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME', kernel_initializer=conv_kernel_initializer)
def forward(self, inp):
x = self.conv1(inp)
x = nn.tf_gelu(x)
x = self.conv2(x)
x = inp + x
x = nn.tf_gelu(x)
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 forward(self, inp):
return nn.tf_flatten(self.down1(inp))
class Inter(nn.ModelBase):
def __init__(self, in_ch, lowest_dense_res, ae_ch, ae_out_ch, d_ch, **kwargs):
self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch, self.d_ch = in_ch, lowest_dense_res, ae_ch, ae_out_ch, d_ch
super().__init__(**kwargs)
def on_build(self):
in_ch, lowest_dense_res, ae_ch, ae_out_ch, d_ch = self.in_ch, self.lowest_dense_res, self.ae_ch, self.ae_out_ch, self.d_ch
self.dense1 = nn.Dense( in_ch, ae_ch, kernel_initializer=tf.initializers.orthogonal )
self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch, maxout_features=4, kernel_initializer=tf.initializers.orthogonal )
self.upscale1 = Upscale(ae_out_ch, d_ch*8)
self.res1 = ResidualBlock(d_ch*8)
def forward(self, inp):
x = self.dense1(inp)
x = self.dense2(x)
x = nn.tf_reshape_4D (x, lowest_dense_res, lowest_dense_res, self.ae_out_ch)
x = self.upscale1(x)
x = self.res1(x)
return x
def get_out_ch(self):
return self.ae_out_ch
class Decoder(nn.ModelBase):
def on_build(self, in_ch, d_ch):
self.upscale1 = Upscale(in_ch, d_ch*4)
self.res1 = ResidualBlock(d_ch*4)
self.upscale2 = Upscale(d_ch*4, d_ch*2)
self.res2 = ResidualBlock(d_ch*2)
self.upscale3 = Upscale(d_ch*2, d_ch*1)
self.res3 = ResidualBlock(d_ch*1)
self.upscalem1 = Upscale(in_ch, d_ch)
self.upscalem2 = Upscale(d_ch, d_ch//2)
self.upscalem3 = Upscale(d_ch//2, d_ch//2)
self.out_conv = nn.Conv2D( d_ch*1, 3, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
self.out_convm = nn.Conv2D( d_ch//2, 1, kernel_size=1, padding='SAME', kernel_initializer=conv_kernel_initializer)
def forward(self, inp):
z = inp
x = self.upscale1 (z)
x = self.res1 (x)
x = self.upscale2 (x)
x = self.res2 (x)
x = self.upscale3 (x)
x = self.res3 (x)
y = self.upscalem1 (z)
y = self.upscalem2 (y)
y = self.upscalem3 (y)
return tf.nn.sigmoid(self.out_conv(x)), \
tf.nn.sigmoid(self.out_convm(y))
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
resolution = self.resolution = 96
ae_dims = 128
e_dims = 128
d_dims = 64
self.pretrain = False
self.pretrain_just_disabled = False
masked_training = True
models_opt_on_gpu = len(devices) >= 1 and all([dev.total_mem_gb >= 2 for dev in devices])
models_opt_device = '/GPU:0' if models_opt_on_gpu and self.is_training else '/CPU:0'
optimizer_vars_on_cpu = models_opt_device=='/CPU:0'
input_ch = 3
output_ch = 3
bgr_shape = nn.get4Dshape(resolution,resolution,input_ch)
mask_shape = nn.get4Dshape(resolution,resolution,1)
lowest_dense_res = resolution // 16
self.model_filename_list = []
with tf.device ('/CPU:0'):
#Place holders on CPU
self.warped_src = tf.placeholder (nn.tf_floatx, bgr_shape)
self.warped_dst = tf.placeholder (nn.tf_floatx, bgr_shape)
self.target_src = tf.placeholder (nn.tf_floatx, bgr_shape)
self.target_dst = tf.placeholder (nn.tf_floatx, bgr_shape)
self.target_srcm = tf.placeholder (nn.tf_floatx, mask_shape)
self.target_dstm = tf.placeholder (nn.tf_floatx, mask_shape)
# Initializing model classes
with tf.device (models_opt_device):
self.encoder = Encoder(in_ch=input_ch, e_ch=e_dims, name='encoder')
encoder_out_ch = self.encoder.compute_output_channels ( (nn.tf_floatx, bgr_shape))
self.inter = Inter (in_ch=encoder_out_ch, lowest_dense_res=lowest_dense_res, ae_ch=ae_dims, ae_out_ch=ae_dims, d_ch=d_dims, name='inter')
inter_out_ch = self.inter.compute_output_channels ( (nn.tf_floatx, (None,encoder_out_ch)))
self.decoder_src = Decoder(in_ch=inter_out_ch, d_ch=d_dims, name='decoder_src')
self.decoder_dst = Decoder(in_ch=inter_out_ch, d_ch=d_dims, name='decoder_dst')
self.model_filename_list += [ [self.encoder, 'encoder.npy' ],
[self.inter, 'inter.npy' ],
[self.decoder_src, 'decoder_src.npy'],
[self.decoder_dst, 'decoder_dst.npy'] ]
if self.is_training:
self.src_dst_trainable_weights = self.encoder.get_weights() + self.inter.get_weights() + self.decoder_src.get_weights() + self.decoder_dst.get_weights()
# Initialize optimizers
self.src_dst_opt = nn.TFRMSpropOptimizer(lr=2e-4, lr_dropout=0.3, name='src_dst_opt')
self.src_dst_opt.initialize_variables(self.src_dst_trainable_weights, vars_on_cpu=optimizer_vars_on_cpu )
self.model_filename_list += [ (self.src_dst_opt, 'src_dst_opt.npy') ]
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices) )
bs_per_gpu = max(1, 4 // gpu_count)
self.set_batch_size( gpu_count*bs_per_gpu)
# Compute losses per GPU
gpu_pred_src_src_list = []
gpu_pred_dst_dst_list = []
gpu_pred_src_dst_list = []
gpu_pred_src_srcm_list = []
gpu_pred_dst_dstm_list = []
gpu_pred_src_dstm_list = []
gpu_src_losses = []
gpu_dst_losses = []
gpu_src_dst_loss_gvs = []
for gpu_id in range(gpu_count):
with tf.device( f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0' ):
batch_slice = slice( gpu_id*bs_per_gpu, (gpu_id+1)*bs_per_gpu )
with tf.device(f'/CPU:0'):
# slice on CPU, otherwise all batch data will be transfered to GPU first
gpu_warped_src = self.warped_src [batch_slice,:,:,:]
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 = self.target_srcm[batch_slice,:,:,:]
gpu_target_dstm = self.target_dstm[batch_slice,:,:,:]
# process model tensors
gpu_src_code = self.inter(self.encoder(gpu_warped_src))
gpu_dst_code = self.inter(self.encoder(gpu_warped_dst))
gpu_pred_src_src, gpu_pred_src_srcm = self.decoder_src(gpu_src_code)
gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code)
gpu_pred_src_src_list.append(gpu_pred_src_src)
gpu_pred_dst_dst_list.append(gpu_pred_dst_dst)
gpu_pred_src_dst_list.append(gpu_pred_src_dst)
gpu_pred_src_srcm_list.append(gpu_pred_src_srcm)
gpu_pred_dst_dstm_list.append(gpu_pred_dst_dstm)
gpu_pred_src_dstm_list.append(gpu_pred_src_dstm)
gpu_target_srcm_blur = nn.tf_gaussian_blur(gpu_target_srcm, max(1, resolution // 32) )
gpu_target_dstm_blur = nn.tf_gaussian_blur(gpu_target_dstm, max(1, resolution // 32) )
gpu_target_dst_masked = gpu_target_dst*gpu_target_dstm_blur
gpu_target_dst_anti_masked = gpu_target_dst*(1.0 - gpu_target_dstm_blur)
gpu_target_src_masked_opt = gpu_target_src*gpu_target_srcm_blur if masked_training else gpu_target_src
gpu_target_dst_masked_opt = gpu_target_dst_masked if masked_training else gpu_target_dst
gpu_pred_src_src_masked_opt = gpu_pred_src_src*gpu_target_srcm_blur if masked_training else gpu_pred_src_src
gpu_pred_dst_dst_masked_opt = gpu_pred_dst_dst*gpu_target_dstm_blur if masked_training else gpu_pred_dst_dst
gpu_psd_target_dst_masked = gpu_pred_src_dst*gpu_target_dstm_blur
gpu_psd_target_dst_anti_masked = gpu_pred_src_dst*(1.0 - gpu_target_dstm_blur)
gpu_src_loss = tf.reduce_mean ( 10*nn.tf_dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_src_masked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3])
gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] )
gpu_dst_loss = tf.reduce_mean ( 10*nn.tf_dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3])
gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] )
gpu_src_losses += [gpu_src_loss]
gpu_dst_losses += [gpu_dst_loss]
gpu_G_loss = gpu_src_loss + gpu_dst_loss
gpu_src_dst_loss_gvs += [ nn.tf_gradients ( gpu_G_loss, self.src_dst_trainable_weights ) ]
# Average losses and gradients, and create optimizer update ops
with tf.device (models_opt_device):
pred_src_src = nn.tf_concat(gpu_pred_src_src_list, 0)
pred_dst_dst = nn.tf_concat(gpu_pred_dst_dst_list, 0)
pred_src_dst = nn.tf_concat(gpu_pred_src_dst_list, 0)
pred_src_srcm = nn.tf_concat(gpu_pred_src_srcm_list, 0)
pred_dst_dstm = nn.tf_concat(gpu_pred_dst_dstm_list, 0)
pred_src_dstm = nn.tf_concat(gpu_pred_src_dstm_list, 0)
src_loss = nn.tf_average_tensor_list(gpu_src_losses)
dst_loss = nn.tf_average_tensor_list(gpu_dst_losses)
src_dst_loss_gv = nn.tf_average_gv_list (gpu_src_dst_loss_gvs)
src_dst_loss_gv_op = self.src_dst_opt.get_update_op (src_dst_loss_gv)
# Initializing training and view functions
def src_dst_train(warped_src, target_src, target_srcm, \
warped_dst, target_dst, target_dstm):
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:target_srcm,
self.warped_dst :warped_dst,
self.target_dst :target_dst,
self.target_dstm:target_dstm,
})
s = np.mean(s)
d = np.mean(d)
return s, d
self.src_dst_train = src_dst_train
def AE_view(warped_src, warped_dst):
return nn.tf_sess.run ( [pred_src_src, pred_dst_dst, pred_dst_dstm, pred_src_dst, pred_src_dstm],
feed_dict={self.warped_src:warped_src,
self.warped_dst:warped_dst})
self.AE_view = AE_view
else:
# Initializing merge function
with tf.device( f'/GPU:0' if len(devices) != 0 else f'/CPU:0'):
gpu_dst_code = self.inter(self.encoder(self.warped_dst))
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code)
_, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
def AE_merge( warped_dst):
return nn.tf_sess.run ( [gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm], feed_dict={self.warped_dst:warped_dst})
self.AE_merge = AE_merge
# Loading/initializing all models/optimizers weights
for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"):
if self.pretrain_just_disabled:
do_init = False
if model == self.inter:
do_init = True
else:
do_init = self.is_first_run()
if not do_init:
do_init = not model.load_weights( self.get_strpath_storage_for_file(filename) )
if do_init and self.pretrained_model_path is not None:
pretrained_filepath = self.pretrained_model_path / filename
if pretrained_filepath.exists():
do_init = not model.load_weights(pretrained_filepath)
if do_init:
model.init_weights()
# initializing sample generators
if self.is_training:
t = SampleProcessor.Types
face_type = t.FACE_TYPE_FULL
training_data_src_path = self.training_data_src_path if not self.pretrain else self.get_pretraining_data_path()
training_data_dst_path = self.training_data_dst_path if not self.pretrain else self.get_pretraining_data_path()
cpu_count = min(multiprocessing.cpu_count(), 8)
src_generators_count = cpu_count // 2
dst_generators_count = cpu_count // 2
self.set_training_data_generators ([
SampleGeneratorFace(training_data_src_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=True if self.pretrain else False),
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution':resolution, },
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution, },
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'data_format':nn.data_format, 'resolution': resolution } ],
generators_count=src_generators_count ),
SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=True if self.pretrain else False),
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution':resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'data_format':nn.data_format, 'resolution': resolution} ],
generators_count=dst_generators_count )
])
self.last_samples = None
def on_initialize(self):
nn.initialize()
tf = nn.tf
class EncBlock(nn.ModelBase):
def on_build(self, in_ch, out_ch, level):
self.zero_level = level == 0
self.conv1 = nn.Conv2D(in_ch,
out_ch,
kernel_size=3,
padding='SAME')
self.conv2 = nn.Conv2D(
out_ch,
out_ch,
kernel_size=4 if self.zero_level else 3,
padding='VALID' if self.zero_level else 'SAME')
def forward(self, x):
x = tf.nn.leaky_relu(self.conv1(x), 0.2)
x = tf.nn.leaky_relu(self.conv2(x), 0.2)
if not self.zero_level:
x = nn.max_pool(x)
return x
class DecBlock(nn.ModelBase):
def on_build(self, in_ch, out_ch, level):
self.zero_level = level == 0
self.conv1 = nn.Conv2D(
in_ch,
out_ch,
kernel_size=4 if self.zero_level else 3,
padding=3 if self.zero_level else 'SAME')
self.conv2 = nn.Conv2D(out_ch,
out_ch,
kernel_size=3,
padding='SAME')
def forward(self, x):
if not self.zero_level:
x = nn.upsample2d(x)
x = tf.nn.leaky_relu(self.conv1(x), 0.2)
x = tf.nn.leaky_relu(self.conv2(x), 0.2)
return x
class FromRGB(nn.ModelBase):
def on_build(self, out_ch):
self.conv1 = nn.Conv2D(3,
out_ch,
kernel_size=1,
padding='SAME')
def forward(self, x):
return tf.nn.leaky_relu(self.conv1(x), 0.2)
class ToRGB(nn.ModelBase):
def on_build(self, in_ch):
self.conv = nn.Conv2D(in_ch, 3, kernel_size=1, padding='SAME')
self.convm = nn.Conv2D(in_ch, 1, kernel_size=1, padding='SAME')
def forward(self, x):
return tf.nn.sigmoid(self.conv(x)), tf.nn.sigmoid(
self.convm(x))
class Encoder(nn.ModelBase):
def on_build(self, e_ch, levels):
self.enc_blocks = {}
self.from_rgbs = {}
self.dense_norm = nn.DenseNorm()
in_ch = e_ch
out_ch = in_ch
for level in range(levels, -1, -1):
self.max_ch = out_ch = np.clip(out_ch * 2, 0, 512)
self.enc_blocks[level] = EncBlock(in_ch, out_ch, level)
self.from_rgbs[level] = FromRGB(in_ch)
in_ch = out_ch
def forward(self, inp, stage):
x = inp
for level in range(stage, -1, -1):
if stage in self.enc_blocks:
if level == stage:
x = self.from_rgbs[level](x)
x = self.enc_blocks[level](x)
x = nn.flatten(x)
x = self.dense_norm(x)
x = nn.reshape_4D(x, 1, 1, self.max_ch)
return x
def get_stage_weights(self, stage):
self.get_weights()
weights = []
for level in range(stage, -1, -1):
if stage in self.enc_blocks:
if level == stage:
weights.append(self.from_rgbs[level].get_weights())
weights.append(self.enc_blocks[level].get_weights())
if len(weights) == 0:
return []
elif len(weights) == 1:
return weights[0]
else:
return sum(weights[1:], weights[0])
class Decoder(nn.ModelBase):
def on_build(self, d_ch, total_levels, levels_range):
self.dec_blocks = {}
self.to_rgbs = {}
level_ch = {}
ch = d_ch
for level in range(total_levels, -2, -1):
level_ch[level] = ch
ch = np.clip(ch * 2, 0, 512)
out_ch = level_ch[levels_range[1]]
for level in range(levels_range[1], levels_range[0] - 1, -1):
in_ch = level_ch[level - 1]
self.dec_blocks[level] = DecBlock(in_ch, out_ch, level)
self.to_rgbs[level] = ToRGB(out_ch)
out_ch = in_ch
def forward(self, inp, stage):
x = inp
for level in range(stage + 1):
if level in self.dec_blocks:
x = self.dec_blocks[level](x)
if level == stage:
x = self.to_rgbs[level](x)
return x
def get_stage_weights(self, stage):
# Call internal get_weights in order to initialize inner logic
self.get_weights()
weights = []
for level in range(stage + 1):
if level in self.dec_blocks:
weights.append(self.dec_blocks[level].get_weights())
if level == stage:
weights.append(self.to_rgbs[level].get_weights())
if len(weights) == 0:
return []
elif len(weights) == 1:
return weights[0]
else:
return sum(weights[1:], weights[0])
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
self.stage = stage = self.options['stage']
self.start_stage_iter = self.options.get('start_stage_iter', 0)
self.target_stage_iter = self.options.get('target_stage_iter', 0)
stage_resolutions = [2**(i + 2) for i in range(self.stage_max + 1)]
stage_resolution = stage_resolutions[stage]
ed_dims = 16
self.pretrain = False
self.pretrain_just_disabled = False
masked_training = True
models_opt_on_gpu = len(devices) == 1 and devices[0].total_mem_gb >= 4
models_opt_device = '/GPU:0' if models_opt_on_gpu and self.is_training else '/CPU:0'
optimizer_vars_on_cpu = models_opt_device == '/CPU:0'
input_nc = 3
output_nc = 3
bgr_shape = (stage_resolution, stage_resolution, output_nc)
mask_shape = (stage_resolution, stage_resolution, 1)
self.model_filename_list = []
with tf.device('/CPU:0'):
#Place holders on CPU
self.warped_src = tf.placeholder(tf.float32, (None, ) + bgr_shape)
self.warped_dst = tf.placeholder(tf.float32, (None, ) + bgr_shape)
self.target_src = tf.placeholder(tf.float32, (None, ) + bgr_shape)
self.target_dst = tf.placeholder(tf.float32, (None, ) + bgr_shape)
self.target_srcm = tf.placeholder(tf.float32,
(None, ) + mask_shape)
self.target_dstm = tf.placeholder(tf.float32,
(None, ) + mask_shape)
# Initializing model classes
with tf.device(models_opt_device):
self.encoder = Encoder(e_ch=ed_dims,
levels=self.stage_max,
name='encoder')
self.inter = Decoder(d_ch=ed_dims,
total_levels=self.stage_max,
levels_range=[0, 2],
name='inter')
self.decoder_src = Decoder(d_ch=ed_dims,
total_levels=self.stage_max,
levels_range=[3, self.stage_max],
name='decoder_src')
self.decoder_dst = Decoder(d_ch=ed_dims,
total_levels=self.stage_max,
levels_range=[3, self.stage_max],
name='decoder_dst')
self.model_filename_list += [[self.encoder, 'encoder.npy'],
[self.inter, 'inter.npy'],
[self.decoder_src, 'decoder_src.npy'],
[self.decoder_dst, 'decoder_dst.npy']]
if self.is_training:
self.src_dst_all_weights = self.encoder.get_weights(
) + self.inter.get_weights() + self.decoder_src.get_weights(
) + self.decoder_dst.get_weights()
self.src_dst_trainable_weights = self.encoder.get_stage_weights(stage) + self.inter.get_stage_weights(stage) \
+ self.decoder_src.get_stage_weights(stage) \
+ self.decoder_dst.get_stage_weights(stage)
# Initialize optimizers
self.src_dst_opt = nn.RMSprop(lr=2e-4,
lr_dropout=0.3,
name='src_dst_opt')
self.src_dst_opt.initialize_variables(
self.src_dst_all_weights,
vars_on_cpu=optimizer_vars_on_cpu)
self.model_filename_list += [(self.src_dst_opt,
'src_dst_opt.npy')]
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices))
bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
self.set_batch_size(gpu_count * bs_per_gpu)
# Compute losses per GPU
gpu_pred_src_src_list = []
gpu_pred_dst_dst_list = []
gpu_pred_src_dst_list = []
gpu_pred_src_srcm_list = []
gpu_pred_dst_dstm_list = []
gpu_pred_src_dstm_list = []
gpu_src_losses = []
gpu_dst_losses = []
gpu_src_dst_loss_gvs = []
for gpu_id in range(gpu_count):
with tf.device(
f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0'):
batch_slice = slice(gpu_id * bs_per_gpu,
(gpu_id + 1) * bs_per_gpu)
with tf.device(f'/CPU:0'):
# slice on CPU, otherwise all batch data will be transfered to GPU first
gpu_warped_src = self.warped_src[batch_slice, :, :, :]
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 = self.target_srcm[
batch_slice, :, :, :]
gpu_target_dstm = self.target_dstm[
batch_slice, :, :, :]
# process model tensors
gpu_src_code = self.inter(
self.encoder(gpu_warped_src, stage), stage)
gpu_dst_code = self.inter(
self.encoder(gpu_warped_dst, stage), stage)
gpu_pred_src_src, gpu_pred_src_srcm = self.decoder_src(
gpu_src_code, stage)
gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder_dst(
gpu_dst_code, stage)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(
gpu_dst_code, stage)
import code
code.interact(local=dict(globals(), **locals()))
gpu_pred_src_src_list.append(gpu_pred_src_src)
gpu_pred_dst_dst_list.append(gpu_pred_dst_dst)
gpu_pred_src_dst_list.append(gpu_pred_src_dst)
gpu_pred_src_srcm_list.append(gpu_pred_src_srcm)
gpu_pred_dst_dstm_list.append(gpu_pred_dst_dstm)
gpu_pred_src_dstm_list.append(gpu_pred_src_dstm)
gpu_target_srcm_blur = nn.gaussian_blur(
gpu_target_srcm, max(1, resolution // 32))
gpu_target_dstm_blur = nn.gaussian_blur(
gpu_target_dstm, max(1, resolution // 32))
gpu_target_dst_masked = gpu_target_dst * gpu_target_dstm_blur
gpu_target_dst_anti_masked = gpu_target_dst * (
1.0 - gpu_target_dstm_blur)
gpu_target_srcmasked_opt = gpu_target_src * gpu_target_srcm_blur if masked_training else gpu_target_src
gpu_target_dst_masked_opt = gpu_target_dst_masked if masked_training else gpu_target_dst
gpu_pred_src_src_masked_opt = gpu_pred_src_src * gpu_target_srcm_blur if masked_training else gpu_pred_src_src
gpu_pred_dst_dst_masked_opt = gpu_pred_dst_dst * gpu_target_dstm_blur if masked_training else gpu_pred_dst_dst
gpu_psd_target_dst_masked = gpu_pred_src_dst * gpu_target_dstm_blur
gpu_psd_target_dst_anti_masked = gpu_pred_src_dst * (
1.0 - gpu_target_dstm_blur)
gpu_src_loss = tf.reduce_mean(
10 * nn.dssim(gpu_target_srcmasked_opt,
gpu_pred_src_src_masked_opt,
max_val=1.0,
filter_size=int(resolution / 11.6)),
axis=[1])
gpu_src_loss += tf.reduce_mean(
10 * tf.square(gpu_target_srcmasked_opt -
gpu_pred_src_src_masked_opt),
axis=[1, 2, 3])
gpu_src_loss += tf.reduce_mean(
tf.square(gpu_target_srcm - gpu_pred_src_srcm),
axis=[1, 2, 3])
gpu_dst_loss = tf.reduce_mean(
10 * nn.dssim(gpu_target_dst_masked_opt,
gpu_pred_dst_dst_masked_opt,
max_val=1.0,
filter_size=int(resolution / 11.6)),
axis=[1])
gpu_dst_loss += tf.reduce_mean(
10 * tf.square(gpu_target_dst_masked_opt -
gpu_pred_dst_dst_masked_opt),
axis=[1, 2, 3])
gpu_dst_loss += tf.reduce_mean(
tf.square(gpu_target_dstm - gpu_pred_dst_dstm),
axis=[1, 2, 3])
gpu_src_losses += [gpu_src_loss]
gpu_dst_losses += [gpu_dst_loss]
gpu_src_dst_loss = gpu_src_loss + gpu_dst_loss
gpu_src_dst_loss_gvs += [
nn.gradients(gpu_src_dst_loss,
self.src_dst_trainable_weights)
]
# Average losses and gradients, and create optimizer update ops
with tf.device(models_opt_device):
if gpu_count == 1:
pred_src_src = gpu_pred_src_src_list[0]
pred_dst_dst = gpu_pred_dst_dst_list[0]
pred_src_dst = gpu_pred_src_dst_list[0]
pred_src_srcm = gpu_pred_src_srcm_list[0]
pred_dst_dstm = gpu_pred_dst_dstm_list[0]
pred_src_dstm = gpu_pred_src_dstm_list[0]
src_loss = gpu_src_losses[0]
dst_loss = gpu_dst_losses[0]
src_dst_loss_gv = gpu_src_dst_loss_gvs[0]
else:
pred_src_src = tf.concat(gpu_pred_src_src_list, 0)
pred_dst_dst = tf.concat(gpu_pred_dst_dst_list, 0)
pred_src_dst = tf.concat(gpu_pred_src_dst_list, 0)
pred_src_srcm = tf.concat(gpu_pred_src_srcm_list, 0)
pred_dst_dstm = tf.concat(gpu_pred_dst_dstm_list, 0)
pred_src_dstm = tf.concat(gpu_pred_src_dstm_list, 0)
src_loss = nn.average_tensor_list(gpu_src_losses)
dst_loss = nn.average_tensor_list(gpu_dst_losses)
src_dst_loss_gv = nn.average_gv_list(gpu_src_dst_loss_gvs)
src_dst_loss_gv_op = self.src_dst_opt.get_update_op(
src_dst_loss_gv)
# Initializing training and view functions
def src_dst_train(warped_src, target_src, target_srcm, \
warped_dst, target_dst, target_dstm):
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: target_srcm,
self.warped_dst: warped_dst,
self.target_dst: target_dst,
self.target_dstm: target_dstm,
})
s = np.mean(s)
d = np.mean(d)
return s, d
self.src_dst_train = src_dst_train
def AE_view(warped_src, warped_dst):
return nn.tf_sess.run([
pred_src_src, pred_dst_dst, pred_dst_dstm, pred_src_dst,
pred_src_dstm
],
feed_dict={
self.warped_src: warped_src,
self.warped_dst: warped_dst
})
self.AE_view = AE_view
else:
# Initializing merge function
with tf.device(f'/GPU:0' if len(devices) != 0 else f'/CPU:0'):
gpu_dst_code = self.inter(self.encoder(self.warped_dst))
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(
gpu_dst_code, stage=stage)
_, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code,
stage=stage)
def AE_merge(warped_dst):
return nn.tf_sess.run(
[gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm],
feed_dict={self.warped_dst: warped_dst})
self.AE_merge = AE_merge
# Loading/initializing all models/optimizers weights
for model, filename in io.progress_bar_generator(
self.model_filename_list, "Initializing models"):
do_init = self.is_first_run()
if self.pretrain_just_disabled:
if model == self.inter:
do_init = True
if not do_init:
do_init = not model.load_weights(
self.get_strpath_storage_for_file(filename))
if do_init:
model.init_weights()
# initializing sample generators
if self.is_training:
t = SampleProcessor.Types
face_type = t.FACE_TYPE_FULL
training_data_src_path = self.training_data_src_path if not self.pretrain else self.get_pretraining_data_path(
)
training_data_dst_path = self.training_data_dst_path if not self.pretrain else self.get_pretraining_data_path(
)
cpu_count = multiprocessing.cpu_count()
src_generators_count = cpu_count // 2
dst_generators_count = cpu_count - src_generators_count
self.set_training_data_generators([
SampleGeneratorFace(
training_data_src_path,
debug=self.is_debug(),
batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(
random_flip=True if self.pretrain else False),
output_sample_types=[{
'types':
(t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR),
'resolution':
resolution,
}, {
'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR),
'resolution':
resolution,
}, {
'types': (t.IMG_TRANSFORMED, face_type,
t.MODE_FACE_MASK_ALL_HULL),
'resolution':
resolution
}],
generators_count=src_generators_count),
SampleGeneratorFace(
training_data_dst_path,
debug=self.is_debug(),
batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(
random_flip=True if self.pretrain else False),
output_sample_types=[{
'types':
(t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR),
'resolution':
resolution
}, {
'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR),
'resolution':
resolution
}, {
'types': (t.IMG_TRANSFORMED, face_type,
t.MODE_FACE_MASK_ALL_HULL),
'resolution':
resolution
}],
generators_count=dst_generators_count)
])
self.last_samples = None
def __init__(self, place_model_on_cpu=False):
nn.initialize(data_format="NHWC")
tf = nn.tf
model_path = Path(__file__).parent / "S3FD.npy"
if not model_path.exists():
raise Exception("Unable to load S3FD.npy")
class L2Norm(nn.LayerBase):
def __init__(self, n_channels, **kwargs):
self.n_channels = n_channels
super().__init__(**kwargs)
def build_weights(self):
self.weight = tf.get_variable ("weight", (1, 1, 1, self.n_channels), dtype=nn.floatx, initializer=tf.initializers.ones )
def get_weights(self):
return [self.weight]
def __call__(self, inputs):
x = inputs
x = x / (tf.sqrt( tf.reduce_sum( tf.pow(x, 2), axis=-1, keepdims=True ) ) + 1e-10) * self.weight
return x
class S3FD(nn.ModelBase):
def __init__(self):
super().__init__(name='S3FD')
def on_build(self):
self.minus = tf.constant([104,117,123], dtype=nn.floatx )
self.conv1_1 = nn.Conv2D(3, 64, kernel_size=3, strides=1, padding='SAME')
self.conv1_2 = nn.Conv2D(64, 64, kernel_size=3, strides=1, padding='SAME')
self.conv2_1 = nn.Conv2D(64, 128, kernel_size=3, strides=1, padding='SAME')
self.conv2_2 = nn.Conv2D(128, 128, kernel_size=3, strides=1, padding='SAME')
self.conv3_1 = nn.Conv2D(128, 256, kernel_size=3, strides=1, padding='SAME')
self.conv3_2 = nn.Conv2D(256, 256, kernel_size=3, strides=1, padding='SAME')
self.conv3_3 = nn.Conv2D(256, 256, kernel_size=3, strides=1, padding='SAME')
self.conv4_1 = nn.Conv2D(256, 512, kernel_size=3, strides=1, padding='SAME')
self.conv4_2 = nn.Conv2D(512, 512, kernel_size=3, strides=1, padding='SAME')
self.conv4_3 = nn.Conv2D(512, 512, kernel_size=3, strides=1, padding='SAME')
self.conv5_1 = nn.Conv2D(512, 512, kernel_size=3, strides=1, padding='SAME')
self.conv5_2 = nn.Conv2D(512, 512, kernel_size=3, strides=1, padding='SAME')
self.conv5_3 = nn.Conv2D(512, 512, kernel_size=3, strides=1, padding='SAME')
self.fc6 = nn.Conv2D(512, 1024, kernel_size=3, strides=1, padding=3)
self.fc7 = nn.Conv2D(1024, 1024, kernel_size=1, strides=1, padding='SAME')
self.conv6_1 = nn.Conv2D(1024, 256, kernel_size=1, strides=1, padding='SAME')
self.conv6_2 = nn.Conv2D(256, 512, kernel_size=3, strides=2, padding='SAME')
self.conv7_1 = nn.Conv2D(512, 128, kernel_size=1, strides=1, padding='SAME')
self.conv7_2 = nn.Conv2D(128, 256, kernel_size=3, strides=2, padding='SAME')
self.conv3_3_norm = L2Norm(256)
self.conv4_3_norm = L2Norm(512)
self.conv5_3_norm = L2Norm(512)
self.conv3_3_norm_mbox_conf = nn.Conv2D(256, 4, kernel_size=3, strides=1, padding='SAME')
self.conv3_3_norm_mbox_loc = nn.Conv2D(256, 4, kernel_size=3, strides=1, padding='SAME')
self.conv4_3_norm_mbox_conf = nn.Conv2D(512, 2, kernel_size=3, strides=1, padding='SAME')
self.conv4_3_norm_mbox_loc = nn.Conv2D(512, 4, kernel_size=3, strides=1, padding='SAME')
self.conv5_3_norm_mbox_conf = nn.Conv2D(512, 2, kernel_size=3, strides=1, padding='SAME')
self.conv5_3_norm_mbox_loc = nn.Conv2D(512, 4, kernel_size=3, strides=1, padding='SAME')
self.fc7_mbox_conf = nn.Conv2D(1024, 2, kernel_size=3, strides=1, padding='SAME')
self.fc7_mbox_loc = nn.Conv2D(1024, 4, kernel_size=3, strides=1, padding='SAME')
self.conv6_2_mbox_conf = nn.Conv2D(512, 2, kernel_size=3, strides=1, padding='SAME')
self.conv6_2_mbox_loc = nn.Conv2D(512, 4, kernel_size=3, strides=1, padding='SAME')
self.conv7_2_mbox_conf = nn.Conv2D(256, 2, kernel_size=3, strides=1, padding='SAME')
self.conv7_2_mbox_loc = nn.Conv2D(256, 4, kernel_size=3, strides=1, padding='SAME')
def forward(self, inp):
x, = inp
x = x - self.minus
x = tf.nn.relu(self.conv1_1(x))
x = tf.nn.relu(self.conv1_2(x))
x = tf.nn.max_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
x = tf.nn.relu(self.conv2_1(x))
x = tf.nn.relu(self.conv2_2(x))
x = tf.nn.max_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
x = tf.nn.relu(self.conv3_1(x))
x = tf.nn.relu(self.conv3_2(x))
x = tf.nn.relu(self.conv3_3(x))
f3_3 = x
x = tf.nn.max_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
x = tf.nn.relu(self.conv4_1(x))
x = tf.nn.relu(self.conv4_2(x))
x = tf.nn.relu(self.conv4_3(x))
f4_3 = x
x = tf.nn.max_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
x = tf.nn.relu(self.conv5_1(x))
x = tf.nn.relu(self.conv5_2(x))
x = tf.nn.relu(self.conv5_3(x))
f5_3 = x
x = tf.nn.max_pool(x, [1,2,2,1], [1,2,2,1], "VALID")
x = tf.nn.relu(self.fc6(x))
x = tf.nn.relu(self.fc7(x))
ffc7 = x
x = tf.nn.relu(self.conv6_1(x))
x = tf.nn.relu(self.conv6_2(x))
f6_2 = x
x = tf.nn.relu(self.conv7_1(x))
x = tf.nn.relu(self.conv7_2(x))
f7_2 = x
f3_3 = self.conv3_3_norm(f3_3)
f4_3 = self.conv4_3_norm(f4_3)
f5_3 = self.conv5_3_norm(f5_3)
cls1 = self.conv3_3_norm_mbox_conf(f3_3)
reg1 = self.conv3_3_norm_mbox_loc(f3_3)
cls2 = tf.nn.softmax(self.conv4_3_norm_mbox_conf(f4_3))
reg2 = self.conv4_3_norm_mbox_loc(f4_3)
cls3 = tf.nn.softmax(self.conv5_3_norm_mbox_conf(f5_3))
reg3 = self.conv5_3_norm_mbox_loc(f5_3)
cls4 = tf.nn.softmax(self.fc7_mbox_conf(ffc7))
reg4 = self.fc7_mbox_loc(ffc7)
cls5 = tf.nn.softmax(self.conv6_2_mbox_conf(f6_2))
reg5 = self.conv6_2_mbox_loc(f6_2)
cls6 = tf.nn.softmax(self.conv7_2_mbox_conf(f7_2))
reg6 = self.conv7_2_mbox_loc(f7_2)
# max-out background label
bmax = tf.maximum(tf.maximum(cls1[:,:,:,0:1], cls1[:,:,:,1:2]), cls1[:,:,:,2:3])
cls1 = tf.concat ([bmax, cls1[:,:,:,3:4] ], axis=-1)
cls1 = tf.nn.softmax(cls1)
return [cls1, reg1, cls2, reg2, cls3, reg3, cls4, reg4, cls5, reg5, cls6, reg6]
e = None
if place_model_on_cpu:
e = tf.device("/CPU:0")
if e is not None: e.__enter__()
self.model = S3FD()
self.model.load_weights (model_path)
if e is not None: e.__exit__(None,None,None)
self.model.build_for_run ([ ( tf.float32, nn.get4Dshape (None,None,3) ) ])
def __init__(self,
name,
resolution,
face_type_str=None,
load_weights=True,
weights_file_root=None,
training=False,
place_model_on_cpu=False,
data_format="NHWC"):
nn.initialize(data_format=data_format)
tf = nn.tf
class Ternaus(nn.ModelBase):
def on_build(self, in_ch, base_ch):
self.features_0 = nn.Conv2D(in_ch,
base_ch,
kernel_size=3,
padding='SAME')
self.blurpool_0 = nn.BlurPool(filt_size=3)
self.features_3 = nn.Conv2D(base_ch,
base_ch * 2,
kernel_size=3,
padding='SAME')
self.blurpool_3 = nn.BlurPool(filt_size=3)
self.features_6 = nn.Conv2D(base_ch * 2,
base_ch * 4,
kernel_size=3,
padding='SAME')
self.features_8 = nn.Conv2D(base_ch * 4,
base_ch * 4,
kernel_size=3,
padding='SAME')
self.blurpool_8 = nn.BlurPool(filt_size=3)
self.features_11 = nn.Conv2D(base_ch * 4,
base_ch * 8,
kernel_size=3,
padding='SAME')
self.features_13 = nn.Conv2D(base_ch * 8,
base_ch * 8,
kernel_size=3,
padding='SAME')
self.blurpool_13 = nn.BlurPool(filt_size=3)
self.features_16 = nn.Conv2D(base_ch * 8,
base_ch * 8,
kernel_size=3,
padding='SAME')
self.features_18 = nn.Conv2D(base_ch * 8,
base_ch * 8,
kernel_size=3,
padding='SAME')
self.blurpool_18 = nn.BlurPool(filt_size=3)
self.conv_center = nn.Conv2D(base_ch * 8,
base_ch * 8,
kernel_size=3,
padding='SAME')
self.conv1_up = nn.Conv2DTranspose(base_ch * 8,
base_ch * 4,
kernel_size=3,
padding='SAME')
self.conv1 = nn.Conv2D(base_ch * 12,
base_ch * 8,
kernel_size=3,
padding='SAME')
self.conv2_up = nn.Conv2DTranspose(base_ch * 8,
base_ch * 4,
kernel_size=3,
padding='SAME')
self.conv2 = nn.Conv2D(base_ch * 12,
base_ch * 8,
kernel_size=3,
padding='SAME')
self.conv3_up = nn.Conv2DTranspose(base_ch * 8,
base_ch * 2,
kernel_size=3,
padding='SAME')
self.conv3 = nn.Conv2D(base_ch * 6,
base_ch * 4,
kernel_size=3,
padding='SAME')
self.conv4_up = nn.Conv2DTranspose(base_ch * 4,
base_ch,
kernel_size=3,
padding='SAME')
self.conv4 = nn.Conv2D(base_ch * 3,
base_ch * 2,
kernel_size=3,
padding='SAME')
self.conv5_up = nn.Conv2DTranspose(base_ch * 2,
base_ch // 2,
kernel_size=3,
padding='SAME')
self.conv5 = nn.Conv2D(base_ch // 2 + base_ch,
base_ch,
kernel_size=3,
padding='SAME')
self.out_conv = nn.Conv2D(base_ch,
1,
kernel_size=3,
padding='SAME')
def forward(self, inp):
x, = inp
x = x0 = tf.nn.relu(self.features_0(x))
x = self.blurpool_0(x)
x = x1 = tf.nn.relu(self.features_3(x))
x = self.blurpool_3(x)
x = tf.nn.relu(self.features_6(x))
x = x2 = tf.nn.relu(self.features_8(x))
x = self.blurpool_8(x)
x = tf.nn.relu(self.features_11(x))
x = x3 = tf.nn.relu(self.features_13(x))
x = self.blurpool_13(x)
x = tf.nn.relu(self.features_16(x))
x = x4 = tf.nn.relu(self.features_18(x))
x = self.blurpool_18(x)
x = self.conv_center(x)
x = tf.nn.relu(self.conv1_up(x))
x = tf.concat([x, x4], nn.conv2d_ch_axis)
x = tf.nn.relu(self.conv1(x))
x = tf.nn.relu(self.conv2_up(x))
x = tf.concat([x, x3], nn.conv2d_ch_axis)
x = tf.nn.relu(self.conv2(x))
x = tf.nn.relu(self.conv3_up(x))
x = tf.concat([x, x2], nn.conv2d_ch_axis)
x = tf.nn.relu(self.conv3(x))
x = tf.nn.relu(self.conv4_up(x))
x = tf.concat([x, x1], nn.conv2d_ch_axis)
x = tf.nn.relu(self.conv4(x))
x = tf.nn.relu(self.conv5_up(x))
x = tf.concat([x, x0], nn.conv2d_ch_axis)
x = tf.nn.relu(self.conv5(x))
logits = self.out_conv(x)
return logits, tf.nn.sigmoid(logits)
if weights_file_root is not None:
weights_file_root = Path(weights_file_root)
else:
weights_file_root = Path(__file__).parent
self.weights_file_root = weights_file_root
with tf.device('/CPU:0'):
#Place holders on CPU
self.input_t = tf.placeholder(
nn.tf_floatx, nn.get4Dshape(resolution, resolution, 3))
self.target_t = tf.placeholder(
nn.tf_floatx, nn.get4Dshape(resolution, resolution, 1))
# Initializing model classes
with tf.device('/CPU:0' if place_model_on_cpu else '/GPU:0'):
self.net = Ternaus(3, 64, name='Ternaus')
self.net_weights = self.net.get_weights()
model_name = f'{name}_{resolution}'
if face_type_str is not None:
model_name += f'_{face_type_str}'
self.model_filename_list = [[self.net, f'{model_name}.npy']]
if training:
self.opt = nn.TFRMSpropOptimizer(lr=0.0001, name='opt')
self.opt.initialize_variables(self.net_weights,
vars_on_cpu=place_model_on_cpu)
self.model_filename_list += [[self.opt, f'{model_name}_opt.npy']]
else:
_, pred = self.net([self.input_t])
def net_run(input_np):
return nn.tf_sess.run([pred],
feed_dict={self.input_t: input_np})[0]
self.net_run = net_run
# Loading/initializing all models/optimizers weights
for model, filename in io.progress_bar_generator(
self.model_filename_list, "Initializing models"):
do_init = not load_weights
if not do_init:
do_init = not model.load_weights(
self.weights_file_root / filename)
if do_init:
model.init_weights()
if model == self.net:
try:
with open(
Path(__file__).parent /
'vgg11_enc_weights.npy', 'rb') as f:
d = pickle.loads(f.read())
for i in [0, 3, 6, 8, 11, 13, 16, 18]:
model.get_layer_by_name('features_%d' %
i).set_weights(
d['features.%d' % i])
except:
io.log_err(
"Unable to load VGG11 pretrained weights from vgg11_enc_weights.npy"
)
def sort_by_absdiff(input_path):
io.log_info("Sorting by absolute difference...")
is_sim = io.input_bool("Sort by similar?",
True,
help_message="Otherwise sort by dissimilar.")
from core.leras import nn
device_config = nn.ask_choose_device_idxs(choose_only_one=True,
return_device_config=True)
nn.initialize(device_config=device_config)
tf = nn.tf
image_paths = pathex.get_image_paths(input_path)
image_paths_len = len(image_paths)
batch_size = 1024
batch_size_remain = image_paths_len % batch_size
i_t = tf.placeholder(tf.float32, (None, 256, 256, 3))
j_t = tf.placeholder(tf.float32, (None, 256, 256, 3))
outputs_full = []
outputs_remain = []
for i in range(batch_size):
diff_t = tf.reduce_sum(tf.abs(i_t - j_t[i]), axis=[1, 2, 3])
outputs_full.append(diff_t)
if i < batch_size_remain:
outputs_remain.append(diff_t)
def func_bs_full(i, j):
return nn.tf_sess.run(outputs_full, feed_dict={i_t: i, j_t: j})
def func_bs_remain(i, j):
return nn.tf_sess.run(outputs_remain, feed_dict={i_t: i, j_t: j})
import h5py
db_file_path = Path(tempfile.gettempdir()) / 'sort_cache.hdf5'
db_file = h5py.File(str(db_file_path), "w")
db = db_file.create_dataset("results", (image_paths_len, image_paths_len),
compression="gzip")
pg_len = image_paths_len // batch_size
if batch_size_remain != 0:
pg_len += 1
pg_len = int((pg_len * pg_len - pg_len) / 2 + pg_len)
io.progress_bar("Computing", pg_len)
j = 0
while j < image_paths_len:
j_images = [cv2_imread(x) for x in image_paths[j:j + batch_size]]
j_images_len = len(j_images)
func = func_bs_remain if image_paths_len - j < batch_size else func_bs_full
i = 0
while i < image_paths_len:
if i >= j:
i_images = [
cv2_imread(x) for x in image_paths[i:i + batch_size]
]
i_images_len = len(i_images)
result = func(i_images, j_images)
db[j:j + j_images_len, i:i + i_images_len] = np.array(result)
io.progress_bar_inc(1)
i += batch_size
db_file.flush()
j += batch_size
io.progress_bar_close()
next_id = 0
sorted = [next_id]
for i in io.progress_bar_generator(range(image_paths_len - 1), "Sorting"):
id_ar = np.concatenate([db[:next_id, next_id], db[next_id, next_id:]])
id_ar = np.argsort(id_ar)
next_id = np.setdiff1d(id_ar, sorted, True)[0 if is_sim else -1]
sorted += [next_id]
db_file.close()
db_file_path.unlink()
img_list = [(image_paths[x], ) for x in sorted]
return img_list, []
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
self.model_data_format = "NCHW"
nn.initialize(data_format=self.model_data_format)
tf = nn.tf
self.resolution = resolution = self.options['resolution']
lowest_dense_res = self.lowest_dense_res = resolution // 32
class Downscale(nn.ModelBase):
def __init__(self, in_ch, out_ch, kernel_size=5, *kwargs ):
self.in_ch = in_ch
self.out_ch = out_ch
self.kernel_size = kernel_size
super().__init__(*kwargs)
def on_build(self, *args, **kwargs ):
self.conv1 = nn.Conv2D( self.in_ch, self.out_ch, kernel_size=self.kernel_size, strides=2, padding='SAME')
def forward(self, x):
x = self.conv1(x)
x = tf.nn.leaky_relu(x, 0.1)
return x
def get_out_ch(self):
return self.out_ch
class Upscale(nn.ModelBase):
def on_build(self, in_ch, out_ch, kernel_size=3 ):
self.conv1 = nn.Conv2D( in_ch, out_ch*4, kernel_size=kernel_size, padding='SAME')
def forward(self, x):
x = self.conv1(x)
x = tf.nn.leaky_relu(x, 0.1)
x = nn.depth_to_space(x, 2)
return x
class ResidualBlock(nn.ModelBase):
def on_build(self, ch, kernel_size=3 ):
self.conv1 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
self.conv2 = nn.Conv2D( ch, ch, kernel_size=kernel_size, padding='SAME')
def forward(self, inp):
x = self.conv1(inp)
x = tf.nn.leaky_relu(x, 0.2)
x = self.conv2(x)
x = tf.nn.leaky_relu(inp+x, 0.2)
return x
class Encoder(nn.ModelBase):
def on_build(self, in_ch, e_ch, ae_ch):
self.down1 = Downscale(in_ch, e_ch, kernel_size=5)
self.res1 = ResidualBlock(e_ch)
self.down2 = Downscale(e_ch, e_ch*2, kernel_size=5)
self.down3 = Downscale(e_ch*2, e_ch*4, kernel_size=5)
self.down4 = Downscale(e_ch*4, e_ch*8, kernel_size=5)
self.down5 = Downscale(e_ch*8, e_ch*8, kernel_size=5)
self.res5 = ResidualBlock(e_ch*8)
self.dense1 = nn.Dense( lowest_dense_res*lowest_dense_res*e_ch*8, ae_ch )
def forward(self, inp):
x = inp
x = self.down1(x)
x = self.res1(x)
x = self.down2(x)
x = self.down3(x)
x = self.down4(x)
x = self.down5(x)
x = self.res5(x)
x = nn.flatten(x)
x = nn.pixel_norm(x, axes=-1)
x = self.dense1(x)
return x
class Inter(nn.ModelBase):
def __init__(self, ae_ch, ae_out_ch, **kwargs):
self.ae_ch, self.ae_out_ch = ae_ch, ae_out_ch
super().__init__(**kwargs)
def on_build(self):
ae_ch, ae_out_ch = self.ae_ch, self.ae_out_ch
self.dense2 = nn.Dense( ae_ch, lowest_dense_res * lowest_dense_res * ae_out_ch )
def forward(self, inp):
x = inp
x = self.dense2(x)
x = nn.reshape_4D (x, lowest_dense_res, lowest_dense_res, self.ae_out_ch)
return x
def get_out_ch(self):
return self.ae_out_ch
class Decoder(nn.ModelBase):
def on_build(self, in_ch, d_ch, d_mask_ch ):
self.upscale0 = Upscale(in_ch, d_ch*8, kernel_size=3)
self.upscale1 = Upscale(d_ch*8, d_ch*8, kernel_size=3)
self.upscale2 = Upscale(d_ch*8, d_ch*4, kernel_size=3)
self.upscale3 = Upscale(d_ch*4, d_ch*2, kernel_size=3)
self.res0 = ResidualBlock(d_ch*8, kernel_size=3)
self.res1 = ResidualBlock(d_ch*8, kernel_size=3)
self.res2 = ResidualBlock(d_ch*4, kernel_size=3)
self.res3 = ResidualBlock(d_ch*2, kernel_size=3)
self.upscalem0 = Upscale(in_ch, d_mask_ch*8, kernel_size=3)
self.upscalem1 = Upscale(d_mask_ch*8, d_mask_ch*8, kernel_size=3)
self.upscalem2 = Upscale(d_mask_ch*8, d_mask_ch*4, kernel_size=3)
self.upscalem3 = Upscale(d_mask_ch*4, d_mask_ch*2, kernel_size=3)
self.upscalem4 = Upscale(d_mask_ch*2, d_mask_ch*1, kernel_size=3)
self.out_convm = nn.Conv2D( d_mask_ch*1, 1, kernel_size=1, padding='SAME')
self.out_conv = nn.Conv2D( d_ch*2, 3, kernel_size=1, padding='SAME')
self.out_conv1 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME')
self.out_conv2 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME')
self.out_conv3 = nn.Conv2D( d_ch*2, 3, kernel_size=3, padding='SAME')
def forward(self, inp):
z = inp
x = self.upscale0(z)
x = self.res0(x)
x = self.upscale1(x)
x = self.res1(x)
x = self.upscale2(x)
x = self.res2(x)
x = self.upscale3(x)
x = self.res3(x)
x = tf.nn.sigmoid( nn.depth_to_space(tf.concat( (self.out_conv(x),
self.out_conv1(x),
self.out_conv2(x),
self.out_conv3(x)), nn.conv2d_ch_axis), 2) )
m = self.upscalem0(z)
m = self.upscalem1(m)
m = self.upscalem2(m)
m = self.upscalem3(m)
m = self.upscalem4(m)
m = tf.nn.sigmoid(self.out_convm(m))
return x, m
self.face_type = {'wf' : FaceType.WHOLE_FACE,
'head' : FaceType.HEAD}[ self.options['face_type'] ]
if 'eyes_prio' in self.options:
self.options.pop('eyes_prio')
eyes_mouth_prio = self.options['eyes_mouth_prio']
ae_dims = self.ae_dims = self.options['ae_dims']
e_dims = self.options['e_dims']
d_dims = self.options['d_dims']
d_mask_dims = self.options['d_mask_dims']
morph_factor = self.options['morph_factor']
pretrain = self.pretrain = self.options['pretrain']
if self.pretrain_just_disabled:
self.set_iter(0)
self.gan_power = gan_power = 0.0 if self.pretrain else self.options['gan_power']
random_warp = False if self.pretrain else self.options['random_warp']
random_src_flip = self.random_src_flip if not self.pretrain else True
random_dst_flip = self.random_dst_flip if not self.pretrain else True
if self.pretrain:
self.options_show_override['gan_power'] = 0.0
self.options_show_override['random_warp'] = False
self.options_show_override['lr_dropout'] = 'n'
self.options_show_override['uniform_yaw'] = True
masked_training = self.options['masked_training']
ct_mode = self.options['ct_mode']
if ct_mode == 'none':
ct_mode = None
models_opt_on_gpu = False if len(devices) == 0 else self.options['models_opt_on_gpu']
models_opt_device = nn.tf_default_device_name if models_opt_on_gpu and self.is_training else '/CPU:0'
optimizer_vars_on_cpu = models_opt_device=='/CPU:0'
input_ch=3
bgr_shape = self.bgr_shape = nn.get4Dshape(resolution,resolution,input_ch)
mask_shape = nn.get4Dshape(resolution,resolution,1)
self.model_filename_list = []
with tf.device ('/CPU:0'):
#Place holders on CPU
self.warped_src = tf.placeholder (nn.floatx, bgr_shape, name='warped_src')
self.warped_dst = tf.placeholder (nn.floatx, bgr_shape, name='warped_dst')
self.target_src = tf.placeholder (nn.floatx, bgr_shape, name='target_src')
self.target_dst = tf.placeholder (nn.floatx, bgr_shape, name='target_dst')
self.target_srcm = tf.placeholder (nn.floatx, mask_shape, name='target_srcm')
self.target_srcm_em = tf.placeholder (nn.floatx, mask_shape, name='target_srcm_em')
self.target_dstm = tf.placeholder (nn.floatx, mask_shape, name='target_dstm')
self.target_dstm_em = tf.placeholder (nn.floatx, mask_shape, name='target_dstm_em')
self.morph_value_t = tf.placeholder (nn.floatx, (1,), name='morph_value_t')
# Initializing model classes
with tf.device (models_opt_device):
self.encoder = Encoder(in_ch=input_ch, e_ch=e_dims, ae_ch=ae_dims, name='encoder')
self.inter_src = Inter(ae_ch=ae_dims, ae_out_ch=ae_dims, name='inter_src')
self.inter_dst = Inter(ae_ch=ae_dims, ae_out_ch=ae_dims, name='inter_dst')
self.decoder = Decoder(in_ch=ae_dims, d_ch=d_dims, d_mask_ch=d_mask_dims, name='decoder')
self.model_filename_list += [ [self.encoder, 'encoder.npy'],
[self.inter_src, 'inter_src.npy'],
[self.inter_dst , 'inter_dst.npy'],
[self.decoder , 'decoder.npy'] ]
if self.is_training:
if gan_power != 0:
self.GAN = nn.UNetPatchDiscriminator(patch_size=self.options['gan_patch_size'], in_ch=input_ch, base_ch=self.options['gan_dims'], name="GAN")
self.model_filename_list += [ [self.GAN, 'GAN.npy'] ]
# Initialize optimizers
lr=5e-5
lr_dropout = 0.3 if self.options['lr_dropout'] in ['y','cpu'] and not self.pretrain else 1.0
clipnorm = 1.0 if self.options['clipgrad'] else 0.0
self.all_weights = self.encoder.get_weights() + self.inter_src.get_weights() + self.inter_dst.get_weights() + self.decoder.get_weights()
if pretrain:
self.trainable_weights = self.encoder.get_weights() + self.inter_dst.get_weights() + self.decoder.get_weights()
else:
self.trainable_weights = self.encoder.get_weights() + self.inter_src.get_weights() + self.inter_dst.get_weights() + self.decoder.get_weights()
self.src_dst_opt = nn.AdaBelief(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='src_dst_opt')
self.src_dst_opt.initialize_variables (self.all_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 gan_power != 0:
self.GAN_opt = nn.AdaBelief(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='GAN_opt')
self.GAN_opt.initialize_variables ( self.GAN.get_weights(), vars_on_cpu=optimizer_vars_on_cpu, lr_dropout_on_cpu=self.options['lr_dropout']=='cpu')#+self.D_src_x2.get_weights()
self.model_filename_list += [ (self.GAN_opt, 'GAN_opt.npy') ]
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices) )
bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
self.set_batch_size( gpu_count*bs_per_gpu)
# Compute losses per GPU
gpu_pred_src_src_list = []
gpu_pred_dst_dst_list = []
gpu_pred_src_dst_list = []
gpu_pred_src_srcm_list = []
gpu_pred_dst_dstm_list = []
gpu_pred_src_dstm_list = []
gpu_src_losses = []
gpu_dst_losses = []
gpu_G_loss_gvs = []
gpu_GAN_loss_gvs = []
gpu_D_code_loss_gvs = []
gpu_D_src_dst_loss_gvs = []
for gpu_id in range(gpu_count):
with tf.device( f'/{devices[gpu_id].tf_dev_type}:{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_warped_src = self.warped_src [batch_slice,:,:,:]
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 = self.target_srcm[batch_slice,:,:,:]
gpu_target_srcm_em = self.target_srcm_em[batch_slice,:,:,:]
gpu_target_dstm = self.target_dstm[batch_slice,:,:,:]
gpu_target_dstm_em = self.target_dstm_em[batch_slice,:,:,:]
# process model tensors
gpu_src_code = self.encoder (gpu_warped_src)
gpu_dst_code = self.encoder (gpu_warped_dst)
if pretrain:
gpu_src_inter_src_code = self.inter_src (gpu_src_code)
gpu_dst_inter_dst_code = self.inter_dst (gpu_dst_code)
gpu_src_code = gpu_src_inter_src_code * nn.random_binomial( [bs_per_gpu, gpu_src_inter_src_code.shape.as_list()[1], 1,1] , p=morph_factor)
gpu_dst_code = gpu_src_dst_code = gpu_dst_inter_dst_code * nn.random_binomial( [bs_per_gpu, gpu_dst_inter_dst_code.shape.as_list()[1], 1,1] , p=0.25)
else:
gpu_src_inter_src_code = self.inter_src (gpu_src_code)
gpu_src_inter_dst_code = self.inter_dst (gpu_src_code)
gpu_dst_inter_src_code = self.inter_src (gpu_dst_code)
gpu_dst_inter_dst_code = self.inter_dst (gpu_dst_code)
inter_rnd_binomial = nn.random_binomial( [bs_per_gpu, gpu_src_inter_src_code.shape.as_list()[1], 1,1] , p=morph_factor)
gpu_src_code = gpu_src_inter_src_code * inter_rnd_binomial + gpu_src_inter_dst_code * (1-inter_rnd_binomial)
gpu_dst_code = gpu_dst_inter_dst_code
ae_dims_slice = tf.cast(ae_dims*self.morph_value_t[0], tf.int32)
gpu_src_dst_code = tf.concat( (tf.slice(gpu_dst_inter_src_code, [0,0,0,0], [-1, ae_dims_slice , lowest_dense_res, lowest_dense_res]),
tf.slice(gpu_dst_inter_dst_code, [0,ae_dims_slice,0,0], [-1,ae_dims-ae_dims_slice, lowest_dense_res,lowest_dense_res]) ), 1 )
gpu_pred_src_src, gpu_pred_src_srcm = self.decoder(gpu_src_code)
gpu_pred_dst_dst, gpu_pred_dst_dstm = self.decoder(gpu_dst_code)
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
gpu_pred_src_src_list.append(gpu_pred_src_src)
gpu_pred_dst_dst_list.append(gpu_pred_dst_dst)
gpu_pred_src_dst_list.append(gpu_pred_src_dst)
gpu_pred_src_srcm_list.append(gpu_pred_src_srcm)
gpu_pred_dst_dstm_list.append(gpu_pred_dst_dstm)
gpu_pred_src_dstm_list.append(gpu_pred_src_dstm)
gpu_target_srcm_blur = nn.gaussian_blur(gpu_target_srcm, max(1, resolution // 32) )
gpu_target_srcm_blur = tf.clip_by_value(gpu_target_srcm_blur, 0, 0.5) * 2
gpu_target_dstm_blur = nn.gaussian_blur(gpu_target_dstm, max(1, resolution // 32) )
gpu_target_dstm_blur = tf.clip_by_value(gpu_target_dstm_blur, 0, 0.5) * 2
gpu_target_dst_anti_masked = gpu_target_dst*(1.0-gpu_target_dstm_blur)
gpu_target_src_anti_masked = gpu_target_src*(1.0-gpu_target_srcm_blur)
gpu_target_src_masked_opt = gpu_target_src*gpu_target_srcm_blur if masked_training else gpu_target_src
gpu_target_dst_masked_opt = gpu_target_dst*gpu_target_dstm_blur if masked_training else gpu_target_dst
gpu_pred_src_src_masked_opt = gpu_pred_src_src*gpu_target_srcm_blur if masked_training else gpu_pred_src_src
gpu_pred_src_src_anti_masked = gpu_pred_src_src*(1.0-gpu_target_srcm_blur)
gpu_pred_dst_dst_masked_opt = gpu_pred_dst_dst*gpu_target_dstm_blur if masked_training else gpu_pred_dst_dst
gpu_pred_dst_dst_anti_masked = gpu_pred_dst_dst*(1.0-gpu_target_dstm_blur)
if resolution < 256:
gpu_dst_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
else:
gpu_dst_loss = tf.reduce_mean ( 5*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/11.6) ), axis=[1])
gpu_dst_loss += tf.reduce_mean ( 5*nn.dssim(gpu_target_dst_masked_opt, gpu_pred_dst_dst_masked_opt, max_val=1.0, filter_size=int(resolution/23.2) ), axis=[1])
gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dst_masked_opt- gpu_pred_dst_dst_masked_opt ), axis=[1,2,3])
if eyes_mouth_prio:
gpu_dst_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_dst*gpu_target_dstm_em - gpu_pred_dst_dst*gpu_target_dstm_em ), axis=[1,2,3])
gpu_dst_loss += tf.reduce_mean ( 10*tf.square( gpu_target_dstm - gpu_pred_dst_dstm ),axis=[1,2,3] )
gpu_dst_loss += 0.1*tf.reduce_mean(tf.square(gpu_pred_dst_dst_anti_masked-gpu_target_dst_anti_masked),axis=[1,2,3] )
gpu_dst_losses += [gpu_dst_loss]
if not pretrain:
if resolution < 256:
gpu_src_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
else:
gpu_src_loss = tf.reduce_mean ( 5*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
gpu_src_loss += tf.reduce_mean ( 5*nn.dssim(gpu_target_src_masked_opt, gpu_pred_src_src_masked_opt, max_val=1.0, filter_size=int(resolution/23.2)), axis=[1])
gpu_src_loss += tf.reduce_mean ( 10*tf.square ( gpu_target_src_masked_opt - gpu_pred_src_src_masked_opt ), axis=[1,2,3])
if eyes_mouth_prio:
gpu_src_loss += tf.reduce_mean ( 300*tf.abs ( gpu_target_src*gpu_target_srcm_em - gpu_pred_src_src*gpu_target_srcm_em ), axis=[1,2,3])
gpu_src_loss += tf.reduce_mean ( 10*tf.square( gpu_target_srcm - gpu_pred_src_srcm ),axis=[1,2,3] )
else:
gpu_src_loss = gpu_dst_loss
gpu_src_losses += [gpu_src_loss]
if pretrain:
gpu_G_loss = gpu_dst_loss
else:
gpu_G_loss = gpu_src_loss + gpu_dst_loss
def DLossOnes(logits):
return tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.ones_like(logits), logits=logits), axis=[1,2,3])
def DLossZeros(logits):
return tf.reduce_mean( tf.nn.sigmoid_cross_entropy_with_logits(labels=tf.zeros_like(logits), logits=logits), axis=[1,2,3])
if gan_power != 0:
gpu_pred_src_src_d, gpu_pred_src_src_d2 = self.GAN(gpu_pred_src_src_masked_opt)
gpu_pred_dst_dst_d, gpu_pred_dst_dst_d2 = self.GAN(gpu_pred_dst_dst_masked_opt)
gpu_target_src_d, gpu_target_src_d2 = self.GAN(gpu_target_src_masked_opt)
gpu_target_dst_d, gpu_target_dst_d2 = self.GAN(gpu_target_dst_masked_opt)
gpu_D_src_dst_loss = (DLossOnes (gpu_target_src_d) + DLossOnes (gpu_target_src_d2) + \
DLossZeros(gpu_pred_src_src_d) + DLossZeros(gpu_pred_src_src_d2) + \
DLossOnes (gpu_target_dst_d) + DLossOnes (gpu_target_dst_d2) + \
DLossZeros(gpu_pred_dst_dst_d) + DLossZeros(gpu_pred_dst_dst_d2)
) * ( 1.0 / 8)
gpu_D_src_dst_loss_gvs += [ nn.gradients (gpu_D_src_dst_loss, self.GAN.get_weights() ) ]
gpu_G_loss += (DLossOnes(gpu_pred_src_src_d) + DLossOnes(gpu_pred_src_src_d2) + \
DLossOnes(gpu_pred_dst_dst_d) + DLossOnes(gpu_pred_dst_dst_d2)
) * gan_power
if masked_training:
# Minimal src-src-bg rec with total_variation_mse to suppress random bright dots from gan
gpu_G_loss += 0.000001*nn.total_variation_mse(gpu_pred_src_src)
gpu_G_loss += 0.02*tf.reduce_mean(tf.square(gpu_pred_src_src_anti_masked-gpu_target_src_anti_masked),axis=[1,2,3] )
gpu_G_loss_gvs += [ nn.gradients ( gpu_G_loss, self.trainable_weights ) ]
# Average losses and gradients, and create optimizer update ops
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)
pred_src_srcm = nn.concat(gpu_pred_src_srcm_list, 0)
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))
if gan_power != 0:
src_D_src_dst_loss_gv_op = self.GAN_opt.get_update_op (nn.average_gv_list(gpu_D_src_dst_loss_gvs) )
#GAN_loss_gv_op = self.src_dst_opt.get_update_op (nn.average_gv_list(gpu_GAN_loss_gvs) )
# Initializing training and view functions
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:target_srcm,
self.target_srcm_em:target_srcm_em,
self.warped_dst :warped_dst,
self.target_dst :target_dst,
self.target_dstm:target_dstm,
self.target_dstm_em:target_dstm_em,
})
return s, d
self.src_dst_train = src_dst_train
if gan_power != 0:
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:target_srcm,
self.target_srcm_em:target_srcm_em,
self.warped_dst :warped_dst,
self.target_dst :target_dst,
self.target_dstm:target_dstm,
self.target_dstm_em:target_dstm_em})
self.D_src_dst_train = D_src_dst_train
def AE_view(warped_src, warped_dst, morph_value):
return nn.tf_sess.run ( [pred_src_src, pred_dst_dst, pred_dst_dstm, pred_src_dst, pred_src_dstm],
feed_dict={self.warped_src:warped_src, self.warped_dst:warped_dst, self.morph_value_t:[morph_value] })
self.AE_view = AE_view
else:
#Initializing merge function
with tf.device( nn.tf_default_device_name if len(devices) != 0 else f'/CPU:0'):
gpu_dst_code = self.encoder (self.warped_dst)
gpu_dst_inter_src_code = self.inter_src ( gpu_dst_code)
gpu_dst_inter_dst_code = self.inter_dst ( gpu_dst_code)
ae_dims_slice = tf.cast(ae_dims*self.morph_value_t[0], tf.int32)
gpu_src_dst_code = tf.concat( ( tf.slice(gpu_dst_inter_src_code, [0,0,0,0], [-1, ae_dims_slice , lowest_dense_res, lowest_dense_res]),
tf.slice(gpu_dst_inter_dst_code, [0,ae_dims_slice,0,0], [-1,ae_dims-ae_dims_slice, lowest_dense_res,lowest_dense_res]) ), 1 )
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder(gpu_src_dst_code)
_, gpu_pred_dst_dstm = self.decoder(gpu_dst_inter_dst_code)
def AE_merge(warped_dst, morph_value):
return nn.tf_sess.run ( [gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm], feed_dict={self.warped_dst:warped_dst, self.morph_value_t:[morph_value] })
self.AE_merge = AE_merge
# Loading/initializing all models/optimizers weights
for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"):
if self.pretrain_just_disabled:
do_init = False
if model == self.inter_src or model == self.inter_dst:
do_init = True
else:
do_init = self.is_first_run()
if self.is_training and gan_power != 0 and model == self.GAN:
if self.gan_model_changed:
do_init = True
if not do_init:
do_init = not model.load_weights( self.get_strpath_storage_for_file(filename) )
if do_init:
model.init_weights()
###############
# initializing sample generators
if self.is_training:
training_data_src_path = self.training_data_src_path if not self.pretrain else self.get_pretraining_data_path()
training_data_dst_path = self.training_data_dst_path if not self.pretrain else self.get_pretraining_data_path()
random_ct_samples_path=training_data_dst_path if ct_mode is not None and not self.pretrain else None
cpu_count = min(multiprocessing.cpu_count(), 8)
src_generators_count = cpu_count // 2
dst_generators_count = cpu_count // 2
if ct_mode is not None:
src_generators_count = int(src_generators_count * 1.5)
self.set_training_data_generators ([
SampleGeneratorFace(training_data_src_path, random_ct_samples_path=random_ct_samples_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=random_src_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.EYES_MOUTH, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
],
uniform_yaw_distribution=self.options['uniform_yaw'] or self.pretrain,
generators_count=src_generators_count ),
SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=random_dst_flip),
output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, '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, '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.EYES_MOUTH, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
],
uniform_yaw_distribution=self.options['uniform_yaw'] or self.pretrain,
generators_count=dst_generators_count )
])
self.last_src_samples_loss = []
self.last_dst_samples_loss = []
if self.pretrain_just_disabled:
self.update_sample_for_preview(force_new=True)
def __init__(self, landmarks_3D=False, place_model_on_cpu=False):
model_path = Path(__file__).parent / ("2DFAN.npy" if not landmarks_3D
else "3DFAN.npy")
if not model_path.exists():
raise Exception("Unable to load FANExtractor model")
nn.initialize(data_format="NHWC")
tf = nn.tf
class ConvBlock(nn.ModelBase):
def on_build(self, in_planes, out_planes):
self.in_planes = in_planes
self.out_planes = out_planes
self.bn1 = nn.BatchNorm2D(in_planes)
self.conv1 = nn.Conv2D(in_planes,
out_planes / 2,
kernel_size=3,
strides=1,
padding='SAME',
use_bias=False)
self.bn2 = nn.BatchNorm2D(out_planes // 2)
self.conv2 = nn.Conv2D(out_planes / 2,
out_planes / 4,
kernel_size=3,
strides=1,
padding='SAME',
use_bias=False)
self.bn3 = nn.BatchNorm2D(out_planes // 4)
self.conv3 = nn.Conv2D(out_planes / 4,
out_planes / 4,
kernel_size=3,
strides=1,
padding='SAME',
use_bias=False)
if self.in_planes != self.out_planes:
self.down_bn1 = nn.BatchNorm2D(in_planes)
self.down_conv1 = nn.Conv2D(in_planes,
out_planes,
kernel_size=1,
strides=1,
padding='VALID',
use_bias=False)
else:
self.down_bn1 = None
self.down_conv1 = None
def forward(self, input):
x = input
x = self.bn1(x)
x = tf.nn.relu(x)
x = out1 = self.conv1(x)
x = self.bn2(x)
x = tf.nn.relu(x)
x = out2 = self.conv2(x)
x = self.bn3(x)
x = tf.nn.relu(x)
x = out3 = self.conv3(x)
x = tf.concat([out1, out2, out3], axis=-1)
if self.in_planes != self.out_planes:
downsample = self.down_bn1(input)
downsample = tf.nn.relu(downsample)
downsample = self.down_conv1(downsample)
x = x + downsample
else:
x = x + input
return x
class HourGlass(nn.ModelBase):
def on_build(self, in_planes, depth):
self.b1 = ConvBlock(in_planes, 256)
self.b2 = ConvBlock(in_planes, 256)
if depth > 1:
self.b2_plus = HourGlass(256, depth - 1)
else:
self.b2_plus = ConvBlock(256, 256)
self.b3 = ConvBlock(256, 256)
def forward(self, input):
up1 = self.b1(input)
low1 = tf.nn.avg_pool(input, [1, 2, 2, 1], [1, 2, 2, 1],
'VALID')
low1 = self.b2(low1)
low2 = self.b2_plus(low1)
low3 = self.b3(low2)
up2 = nn.upsample2d(low3)
return up1 + up2
class FAN(nn.ModelBase):
def __init__(self):
super().__init__(name='FAN')
def on_build(self):
self.conv1 = nn.Conv2D(3,
64,
kernel_size=7,
strides=2,
padding='SAME')
self.bn1 = nn.BatchNorm2D(64)
self.conv2 = ConvBlock(64, 128)
self.conv3 = ConvBlock(128, 128)
self.conv4 = ConvBlock(128, 256)
self.m = []
self.top_m = []
self.conv_last = []
self.bn_end = []
self.l = []
self.bl = []
self.al = []
for i in range(4):
self.m += [HourGlass(256, 4)]
self.top_m += [ConvBlock(256, 256)]
self.conv_last += [
nn.Conv2D(256,
256,
kernel_size=1,
strides=1,
padding='VALID')
]
self.bn_end += [nn.BatchNorm2D(256)]
self.l += [
nn.Conv2D(256,
68,
kernel_size=1,
strides=1,
padding='VALID')
]
if i < 4 - 1:
self.bl += [
nn.Conv2D(256,
256,
kernel_size=1,
strides=1,
padding='VALID')
]
self.al += [
nn.Conv2D(68,
256,
kernel_size=1,
strides=1,
padding='VALID')
]
def forward(self, inp):
x, = inp
x = self.conv1(x)
x = self.bn1(x)
x = tf.nn.relu(x)
x = self.conv2(x)
x = tf.nn.avg_pool(x, [1, 2, 2, 1], [1, 2, 2, 1], 'VALID')
x = self.conv3(x)
x = self.conv4(x)
outputs = []
previous = x
for i in range(4):
ll = self.m[i](previous)
ll = self.top_m[i](ll)
ll = self.conv_last[i](ll)
ll = self.bn_end[i](ll)
ll = tf.nn.relu(ll)
tmp_out = self.l[i](ll)
outputs.append(tmp_out)
if i < 4 - 1:
ll = self.bl[i](ll)
previous = previous + ll + self.al[i](tmp_out)
x = outputs[-1]
x = tf.transpose(x, (0, 3, 1, 2))
return x
e = None
if place_model_on_cpu:
e = tf.device("/CPU:0")
if e is not None: e.__enter__()
self.model = FAN()
self.model.load_weights(str(model_path))
if e is not None: e.__exit__(None, None, None)
self.model.build_for_run([(tf.float32, (None, 256, 256, 3))])
def on_initialize(self):
device_config = nn.getCurrentDeviceConfig()
devices = device_config.devices
self.model_data_format = "NHWC"#"NCHW" if len(devices) != 0 and not self.is_debug() else "NHWC"
nn.initialize(data_format=self.model_data_format)
tf = nn.tf
self.resolution = resolution = self.options['resolution']
self.face_type = {'h' : FaceType.HALF,
'mf' : FaceType.MID_FULL,
'f' : FaceType.FULL,
'wf' : FaceType.WHOLE_FACE,
'head' : FaceType.HEAD}[ self.options['face_type'] ]
models_opt_on_gpu = True#False if len(devices) == 0 else self.options['models_opt_on_gpu']
models_opt_device = '/GPU:0' if models_opt_on_gpu and self.is_training else '/CPU:0'
optimizer_vars_on_cpu = models_opt_device=='/CPU:0'
input_ch=3
bgr_shape = nn.get4Dshape(resolution,resolution,input_ch)
mask_shape = nn.get4Dshape(resolution,resolution,1)
self.model_filename_list = []
class BaseModel(nn.ModelBase):
def on_build(self, in_ch, base_ch, out_ch=None):
self.convs = [ nn.Conv2D( in_ch, base_ch, kernel_size=7, strides=1, padding='SAME'),
nn.Conv2D( base_ch, base_ch, kernel_size=3, strides=1, use_bias=False, padding='SAME'),
nn.Conv2D( base_ch, base_ch*2, kernel_size=3, strides=2, use_bias=False, padding='SAME'),
nn.Conv2D( base_ch*2, base_ch*2, kernel_size=3, strides=1, use_bias=False, padding='SAME'),
nn.Conv2D( base_ch*2, base_ch*4, kernel_size=3, strides=2, use_bias=False, padding='SAME'),
nn.Conv2D( base_ch*4, base_ch*4, kernel_size=3, strides=1, use_bias=False, padding='SAME'),
nn.Conv2D( base_ch*4, base_ch*8, kernel_size=3, strides=2, use_bias=False, padding='SAME'),
nn.Conv2D( base_ch*8, base_ch*8, kernel_size=3, strides=1, use_bias=False, padding='SAME')
]
self.frns = [ None,
nn.FRNorm2D(base_ch),
nn.FRNorm2D(base_ch*2),
nn.FRNorm2D(base_ch*2),
nn.FRNorm2D(base_ch*4),
nn.FRNorm2D(base_ch*4),
nn.FRNorm2D(base_ch*8),
nn.FRNorm2D(base_ch*8),
]
self.tlus = [ nn.TLU(base_ch),
nn.TLU(base_ch),
nn.TLU(base_ch*2),
nn.TLU(base_ch*2),
nn.TLU(base_ch*4),
nn.TLU(base_ch*4),
nn.TLU(base_ch*8),
nn.TLU(base_ch*8),
]
if out_ch is not None:
self.out_conv = nn.Conv2D( base_ch*8, out_ch, kernel_size=1, strides=1, use_bias=False, padding='VALID')
else:
self.out_conv = None
def forward(self, inp):
x = inp
for i in range(len(self.convs)):
x = self.convs[i](x)
if self.frns[i] is not None:
x = self.frns[i](x)
x = self.tlus[i](x)
if self.out_conv is not None:
x = self.out_conv(x)
return x
class Regressor(nn.ModelBase):
def on_build(self, lmrks_ch, base_ch, out_ch):
self.convs = [ nn.Conv2D( base_ch*8+lmrks_ch, base_ch*8, kernel_size=3, strides=1, use_bias=False, padding='SAME'),
nn.Conv2D( base_ch*8, base_ch*8*4, kernel_size=3, strides=1, use_bias=False, padding='SAME'),
nn.Conv2D( base_ch*8, base_ch*4, kernel_size=3, strides=1, use_bias=False, padding='SAME'),
nn.Conv2D( base_ch*4, base_ch*4*4, kernel_size=3, strides=1, use_bias=False, padding='SAME'),
nn.Conv2D( base_ch*4, base_ch*2, kernel_size=3, strides=1, use_bias=False, padding='SAME'),
nn.Conv2D( base_ch*2, base_ch*2*4, kernel_size=3, strides=1, use_bias=False, padding='SAME'),
nn.Conv2D( base_ch*2, base_ch, kernel_size=3, strides=1, use_bias=False, padding='SAME'),
]
self.frns = [ nn.FRNorm2D(base_ch*8),
nn.FRNorm2D(base_ch*8*4),
nn.FRNorm2D(base_ch*4),
nn.FRNorm2D(base_ch*4*4),
nn.FRNorm2D(base_ch*2),
nn.FRNorm2D(base_ch*2*4),
nn.FRNorm2D(base_ch),
]
self.tlus = [ nn.TLU(base_ch*8),
nn.TLU(base_ch*8*4),
nn.TLU(base_ch*4),
nn.TLU(base_ch*4*4),
nn.TLU(base_ch*2),
nn.TLU(base_ch*2*4),
nn.TLU(base_ch),
]
self.use_upscale = [ False,
True,
False,
True,
False,
True,
False,
]
self.out_conv = nn.Conv2D( base_ch, out_ch, kernel_size=3, strides=1, padding='SAME')
def forward(self, inp):
x = inp
for i in range(len(self.convs)):
x = self.convs[i](x)
x = self.frns[i](x)
x = self.tlus[i](x)
if self.use_upscale[i]:
x = nn.depth_to_space(x, 2)
x = self.out_conv(x)
x = tf.nn.sigmoid(x)
return x
def get_coord(x, other_axis, axis_size):
# get "x-y" coordinates:
g_c_prob = tf.reduce_mean(x, axis=other_axis) # B,W,NMAP
g_c_prob = tf.nn.softmax(g_c_prob, axis=1) # B,W,NMAP
coord_pt = tf.to_float(tf.linspace(-1.0, 1.0, axis_size)) # W
coord_pt = tf.reshape(coord_pt, [1, axis_size, 1])
g_c = tf.reduce_sum(g_c_prob * coord_pt, axis=1)
return g_c, g_c_prob
def get_gaussian_maps(mu_x, mu_y, width, height, inv_std=10.0, mode='rot'):
"""
Generates [B,SHAPE_H,SHAPE_W,NMAPS] tensor of 2D gaussians,
given the gaussian centers: MU [B, NMAPS, 2] tensor.
STD: is the fixed standard dev.
"""
y = tf.to_float(tf.linspace(-1.0, 1.0, width))
x = tf.to_float(tf.linspace(-1.0, 1.0, height))
if mode in ['rot', 'flat']:
mu_y, mu_x = mu_y[...,None,None], mu_x[...,None,None]
y = tf.reshape(y, [1, 1, width, 1])
x = tf.reshape(x, [1, 1, 1, height])
g_y = tf.square(y - mu_y)
g_x = tf.square(x - mu_x)
dist = (g_y + g_x) * inv_std**2
if mode == 'rot':
g_yx = tf.exp(-dist)
else:
g_yx = tf.exp(-tf.pow(dist + 1e-5, 0.25))
elif mode == 'ankush':
y = tf.reshape(y, [1, 1, width])
x = tf.reshape(x, [1, 1, height])
g_y = tf.exp(-tf.sqrt(1e-4 + tf.abs((mu_y[...,None] - y) * inv_std)))
g_x = tf.exp(-tf.sqrt(1e-4 + tf.abs((mu_x[...,None] - x) * inv_std)))
g_y = tf.expand_dims(g_y, axis=3)
g_x = tf.expand_dims(g_x, axis=2)
g_yx = tf.matmul(g_y, g_x) # [B, NMAPS, H, W]
else:
raise ValueError('Unknown mode: ' + str(mode))
g_yx = tf.transpose(g_yx, perm=[0, 2, 3, 1])
return g_yx
with tf.device ('/CPU:0'):
#Place holders on CPU
self.warped_src = tf.placeholder (nn.floatx, bgr_shape)
self.target_src = tf.placeholder (nn.floatx, bgr_shape)
# Initializing model classes
#model_archi = nn.DeepFakeArchi(resolution, mod='uhd' if 'uhd' in archi else None)
self.landmarks_count = 512
self.n_ch = 32
with tf.device (models_opt_device):
self.detector = BaseModel(3, self.n_ch, out_ch=self.landmarks_count, name='Detector')
self.extractor = BaseModel(3, self.n_ch, name='Extractor')
self.regressor = Regressor(self.landmarks_count, self.n_ch, 3, name='Regressor')
self.model_filename_list += [ [self.detector, 'detector.npy'],
[self.extractor, 'extractor.npy'],
[self.regressor, 'regressor.npy'] ]
if self.is_training:
# Initialize optimizers
lr=5e-5
lr_dropout = 0.3#0.3 if self.options['lr_dropout'] and not self.pretrain else 1.0
clipnorm = 0.0#1.0 if self.options['clipgrad'] else 0.0
self.model_opt = nn.RMSprop(lr=lr, lr_dropout=lr_dropout, clipnorm=clipnorm, name='model_opt')
self.model_filename_list += [ (self.model_opt, 'model_opt.npy') ]
self.model_trainable_weights = self.detector.get_weights() + self.extractor.get_weights() + self.regressor.get_weights()
self.model_opt.initialize_variables (self.model_trainable_weights, vars_on_cpu=optimizer_vars_on_cpu)
if self.is_training:
# Adjust batch size for multiple GPU
gpu_count = max(1, len(devices) )
bs_per_gpu = max(1, self.get_batch_size() // gpu_count)
self.set_batch_size( gpu_count*bs_per_gpu)
# Compute losses per GPU
gpu_src_rec_list = []
gauss_mu_list = []
gpu_src_losses = []
gpu_G_loss_gvs = []
for gpu_id in range(gpu_count):
with tf.device( f'/GPU:{gpu_id}' if len(devices) != 0 else f'/CPU:0' ):
with tf.device(f'/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_warped_src = self.warped_src [batch_slice,:,:,:]
gpu_target_src = self.target_src [batch_slice,:,:,:]
# process model tensors
gpu_src_feat = self.extractor(gpu_warped_src)
gpu_src_heatmaps = self.detector(gpu_target_src)
gauss_y, gauss_y_prob = get_coord(gpu_src_heatmaps, 2, gpu_src_heatmaps.shape.as_list()[1] )
gauss_x, gauss_x_prob = get_coord(gpu_src_heatmaps, 1, gpu_src_heatmaps.shape.as_list()[2] )
gauss_mu = tf.stack ( (gauss_x, gauss_y), -1)
dist_loss = []
for i in range(self.landmarks_count):
t = tf.concat( (gauss_mu[:,0:i], gauss_mu[:,i+1:] ), axis=1 )
diff = t - gauss_mu[:,i:i+1]
dist = tf.sqrt( diff[...,0]**2+diff[...,1]**2 )
dist_loss += [ tf.reduce_mean(2.0 - dist,-1) ]
dist_loss = sum(dist_loss) / self.landmarks_count
#import code
#code.interact(local=dict(globals(), **locals()))
gauss_xy = get_gaussian_maps ( gauss_x, gauss_y, 16, 16 )
gpu_src_rec = self.regressor( tf.concat ( (gpu_src_feat, gauss_xy), -1) )
gpu_src_rec_list.append(gpu_src_rec)
gauss_mu_list.append(gauss_mu)
gpu_src_loss = tf.reduce_mean ( 10*nn.dssim(gpu_target_src, gpu_src_rec, max_val=1.0, filter_size=int(resolution/11.6)), axis=[1])
gpu_src_loss += tf.reduce_mean ( 10*tf.square (gpu_target_src - gpu_src_rec), axis=[1,2,3])
gpu_src_loss += dist_loss
gpu_src_losses += [gpu_src_loss]
gpu_G_loss_gvs += [ nn.gradients ( gpu_src_loss, self.model_trainable_weights ) ]
# Average losses and gradients, and create optimizer update ops
with tf.device (models_opt_device):
src_rec = nn.concat(gpu_src_rec_list, 0)
gauss_mu = nn.concat(gauss_mu_list, 0)
src_loss = tf.concat(gpu_src_losses, 0)
loss_gv_op = self.model_opt.get_update_op (nn.average_gv_list (gpu_G_loss_gvs))
# Initializing training and view functions
def ae_train(warped_src, target_src):
s, _ = nn.tf_sess.run ( [ src_loss, loss_gv_op], feed_dict={self.warped_src:warped_src, self.target_src:target_src})
return s
self.ae_train = ae_train
def AE_view(warped_src, target_src):
return nn.tf_sess.run ( [src_rec, gauss_mu], feed_dict={self.warped_src:warped_src, self.target_src:target_src})
self.AE_view = AE_view
else:
# Initializing merge function
with tf.device( f'/GPU:0' if len(devices) != 0 else f'/CPU:0'):
if 'df' in archi:
gpu_dst_code = self.inter(self.encoder(self.warped_dst))
gpu_pred_src_dst, gpu_pred_src_dstm = self.decoder_src(gpu_dst_code)
_, gpu_pred_dst_dstm = self.decoder_dst(gpu_dst_code)
def AE_merge( warped_dst):
return nn.tf_sess.run ( [gpu_pred_src_dst, gpu_pred_dst_dstm, gpu_pred_src_dstm], feed_dict={self.warped_dst:warped_dst})
self.AE_merge = AE_merge
# Loading/initializing all models/optimizers weights
for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"):
do_init = self.is_first_run()
if not do_init:
do_init = not model.load_weights( self.get_strpath_storage_for_file(filename) )
if do_init:
model.init_weights()
# initializing sample generators
if self.is_training:
training_data_src_path = self.training_data_src_path
cpu_count = min(multiprocessing.cpu_count(), 8)
src_generators_count = cpu_count // 2
self.set_training_data_generators ([
SampleGeneratorFace(training_data_src_path, debug=self.is_debug(), batch_size=self.get_batch_size()*2,
sample_process_options=SampleProcessor.Options(random_flip=False),
output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':True, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, '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, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution},
],
generators_count=src_generators_count ),
])