def DFEncFlow(padding='zero', **kwargs):
exec(nnlib.import_all(), locals(), globals())
use_bias = False
def XNormalization(x):
return BatchNormalization(axis=-1)(x)
XConv2D = partial(Conv2D, padding=padding, use_bias=use_bias)
def Act(lrelu_alpha=0.1):
return LeakyReLU(alpha=lrelu_alpha)
def downscale(dim, **kwargs):
def func(x):
return Act()(XNormalization(
XConv2D(dim, kernel_size=5, strides=2)(x)))
return func
def upscale(dim, **kwargs):
def func(x):
return SubpixelUpscaler()(Act()(XNormalization(
XConv2D(dim * 4, kernel_size=3, strides=1)(x))))
return func
upscale = partial(upscale)
downscale = partial(downscale)
def func(input):
x, emb_pyr = input
b, h, w, c = K.int_shape(x)
lowest_dense_res = w // 16
dims = 64
x = downscale(dims)(x)
x = downscale(dims * 2)(x)
x = downscale(dims * 4)(x)
x = downscale(dims * 8)(x)
x = Dense(256)(Flatten()(x))
x = Concatenate()([x, emb_pyr])
x = Dense(lowest_dense_res * lowest_dense_res * 256)(x)
x = Reshape((lowest_dense_res, lowest_dense_res, 256))(x)
x = upscale(256)(x)
return x
return func
def LIAEInterFlow(resolution, ae_dims=256, **kwargs):
exec(nnlib.import_all(), locals(), globals())
upscale = partial(SAEModel.upscale, **kwargs)
lowest_dense_res = resolution // 16
def func(input):
x = input[0]
x = Dense(ae_dims)(x)
x = Dense(lowest_dense_res * lowest_dense_res * ae_dims * 2)(x)
x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims * 2))(x)
x = upscale(ae_dims * 2)(x)
return x
return func
def NLayerDiscriminator(use_batch_norm, ndf=64, n_layers=3):
exec(nnlib.import_all(), locals(), globals())
if not use_batch_norm:
use_bias = True
def XNormalization(x):
return InstanceNormalization(axis=-1)(x)
else:
use_bias = False
def XNormalization(x):
return BatchNormalization(axis=-1)(x)
XConv2D = partial(Conv2D, use_bias=use_bias)
def func(input):
b, h, w, c = K.int_shape(input)
x = input
f = ndf
x = ZeroPadding2D((1, 1))(x)
x = XConv2D(f, 4, strides=2, padding='valid', use_bias=True)(x)
f = min(ndf * 8, f * 2)
x = LeakyReLU(0.2)(x)
for i in range(n_layers):
x = ZeroPadding2D((1, 1))(x)
x = XConv2D(f, 4, strides=2, padding='valid')(x)
f = min(ndf * 8, f * 2)
x = XNormalization(x)
x = Dropout(0.5)(x)
x = LeakyReLU(0.2)(x)
x = ZeroPadding2D((1, 1))(x)
x = XConv2D(f, 4, strides=1, padding='valid')(x)
x = XNormalization(x)
x = LeakyReLU(0.2)(x)
x = ZeroPadding2D((1, 1))(x)
return XConv2D(1,
4,
strides=1,
padding='valid',
use_bias=True,
activation='sigmoid')(x) #
return func
def DecFlow(resolution):
exec(nnlib.import_all(), locals(), globals())
XConv2D = partial(Conv2D, padding='zero')
def Act(lrelu_alpha=0.1):
return LeakyReLU(alpha=lrelu_alpha)
def downscale(dim, **kwargs):
def func(x):
return MaxPooling2D()(Act()(XConv2D(dim,
kernel_size=5,
strides=1)(x)))
return func
def upscale(dim, **kwargs):
def func(x):
return SubpixelUpscaler()(Act()(XConv2D(dim * 4,
kernel_size=3,
strides=1)(x)))
return func
def to_bgr(output_nc, **kwargs):
def func(x):
return XConv2D(output_nc, kernel_size=5,
activation='sigmoid')(x)
return func
upscale = partial(upscale)
downscale = partial(downscale)
lowest_dense_res = resolution // 16
def func(input):
x = input
x = Dense(lowest_dense_res * lowest_dense_res * 256,
use_bias=False)(x)
x = Reshape((lowest_dense_res, lowest_dense_res, 256))(x)
x = upscale(512)(x)
x = upscale(256)(x)
x = upscale(128)(x)
x = upscale(64)(x)
bgr = to_bgr(3)(x)
return [bgr]
return func
def LIAEEncFlow(resolution, adapt_k_size, light_enc, ed_ch_dims=42):
exec(nnlib.import_all(), locals(), globals())
k_size = resolution // 16 + 1 if adapt_k_size else 5
strides = resolution // 32 if adapt_k_size else 2
def downscale(dim):
def func(x):
return LeakyReLU(0.1)(Conv2D(dim,
k_size,
strides=strides,
padding='same')(x))
return func
def downscale_sep(dim):
def func(x):
return LeakyReLU(0.1)(SeparableConv2D(dim,
k_size,
strides=strides,
padding='same')(x))
return func
def upscale(dim):
def func(x):
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(
dim * 4, 3, strides=1, padding='same')(x)))
return func
def func(input):
ed_dims = K.int_shape(input)[-1] * ed_ch_dims
x = input
x = downscale(ed_dims)(x)
if not light_enc:
x = downscale(ed_dims * 2)(x)
x = downscale(ed_dims * 4)(x)
x = downscale(ed_dims * 8)(x)
else:
x = downscale_sep(ed_dims * 2)(x)
x = downscale(ed_dims * 4)(x)
x = downscale_sep(ed_dims * 8)(x)
x = Flatten()(x)
return x
return func
def VGEncFlow(resolution, light_enc, ae_dims=512, ed_ch_dims=42):
exec (nnlib.import_all(), locals(), globals())
upscale = SAEModel.upscale
downscale = SAEModel.downscale
downscale_sep = SAEModel.downscale_sep
ResidualBlock = SAEModel.ResidualBlock
lowest_dense_res = resolution // 16
def func(input):
x = input
ed_dims = K.int_shape(input)[-1]*ed_ch_dims
while np.modf(ed_dims / 4)[0] != 0.0:
ed_dims -= 1
in_conv_filters = ed_dims# if resolution <= 128 else ed_dims + (resolution//128)*ed_ch_dims
x = tmp_x = Conv2D (in_conv_filters, kernel_size=5, strides=2, padding='same') (x)
for _ in range ( 8 if light_enc else 16 ):
x = ResidualBlock(ed_dims)(x)
x = Add()([x, tmp_x])
x = downscale(ed_dims)(x)
x = SubpixelUpscaler()(x)
x = downscale(ed_dims)(x)
x = SubpixelUpscaler()(x)
x = downscale(ed_dims)(x)
if light_enc:
x = downscale_sep (ed_dims*2)(x)
else:
x = downscale (ed_dims*2)(x)
x = downscale(ed_dims*4)(x)
if light_enc:
x = downscale_sep (ed_dims*8)(x)
else:
x = downscale (ed_dims*8)(x)
x = Dense(ae_dims)(Flatten()(x))
x = Dense(lowest_dense_res * lowest_dense_res * ae_dims)(x)
x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims))(x)
x = upscale(ae_dims)(x)
return x
return func
def onInitialize(self):
exec(nnlib.import_all(), locals(), globals())
self.set_vram_batch_requirements( {4.5:4} )
ae_input_layer = Input(shape=(128, 128, 3))
mask_layer = Input(shape=(128, 128, 1)) #same as output
self.encoder, self.decoder, self.inter_B, self.inter_AB = self.Build(ae_input_layer)
if not self.is_first_run():
weights_to_load = [ [self.encoder, 'encoder.h5'],
[self.decoder, 'decoder.h5'],
[self.inter_B, 'inter_B.h5'],
[self.inter_AB, 'inter_AB.h5']
]
self.load_weights_safe(weights_to_load)
code = self.encoder(ae_input_layer)
AB = self.inter_AB(code)
B = self.inter_B(code)
rec_src = self.decoder(Concatenate()([AB, AB]))
rec_dst = self.decoder(Concatenate()([B, AB]))
self.autoencoder_src = Model([ae_input_layer,mask_layer], rec_src )
self.autoencoder_dst = Model([ae_input_layer,mask_layer], rec_dst )
self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] )
self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] )
self.convert = K.function([ae_input_layer],rec_src)
if self.is_training_mode:
f = SampleProcessor.TypeFlags
self.set_training_data_generators ([
SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None,
debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ),
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] ),
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_BGR, 128],
[f.TRANSFORMED | f.FACE_ALIGN_FULL | f.MODE_M | f.FACE_MASK_FULL, 128] ] )
])
def InterFlow(dims=256, lowest_dense_res=8):
exec (nnlib.import_all(), locals(), globals())
def upscale (dim):
def func(x):
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
return func
def func(input):
x = input[0]
x = Dense(dims)(x)
x = Dense(lowest_dense_res * lowest_dense_res * dims*2)(x)
x = Reshape((lowest_dense_res, lowest_dense_res, dims*2))(x)
x = upscale(dims*2)(x)
return x
return func
def onClientInitialize(self, client_dict):
self.safe_print('Running on %s.' % (client_dict['device_name']))
self.type = client_dict['type']
self.image_size = client_dict['image_size']
self.face_type = client_dict['face_type']
self.device_idx = client_dict['device_idx']
self.cpu_only = client_dict['device_type'] == 'CPU'
self.output_path = Path(
client_dict['output_dir']) if 'output_dir' in client_dict.keys(
) else None
self.debug = client_dict['debug']
self.detector = client_dict['detector']
self.e = None
device_config = nnlib.DeviceConfig(cpu_only=self.cpu_only,
force_gpu_idx=self.device_idx,
allow_growth=True)
if self.type == 'rects':
if self.detector is not None:
if self.detector == 'mt':
nnlib.import_all(device_config)
self.e = facelib.MTCExtractor()
elif self.detector == 'dlib':
nnlib.import_dlib(device_config)
self.e = facelib.DLIBExtractor(nnlib.dlib)
self.e.__enter__()
elif self.type == 'landmarks':
nnlib.import_all(device_config)
self.e = facelib.LandmarksExtractor(nnlib.keras)
self.e.__enter__()
elif self.type == 'final':
pass
return None
def UNetTemporalPredictor(output_nc,
use_batch_norm,
ngf=64,
use_dropout=False):
exec(nnlib.import_all(), locals(), globals())
def func(inputs):
past_2_image_tensor, past_1_image_tensor = inputs
x = Concatenate(axis=-1)(
[past_2_image_tensor, past_1_image_tensor])
x = UNet(3, use_batch_norm, ngf=ngf, use_dropout=use_dropout)(x)
return x
return func
def Dec64Flow(output_nc=3, **kwargs):
exec(nnlib.import_all(), locals(), globals())
ResidualBlock = CONVModel.ResidualBlock
upscale = CONVModel.upscale
to_bgr = CONVModel.to_bgr
def func(input):
x = input[0]
x = upscale(512)(x)
x = upscale(256)(x)
x = upscale(128)(x)
return to_bgr(output_nc, activation="tanh")(x)
return func
def DecFlow(output_nc=3, **kwargs):
exec(nnlib.import_all(), locals(), globals())
ResidualBlock = AVATARModel.ResidualBlock
upscale = AVATARModel.upscale
to_bgr = AVATARModel.to_bgr
def func(input):
x = input[0]
x = upscale(512)(x)
x = upscale(256)(x)
x = upscale(128)(x)
return to_bgr(output_nc)(x)
return func
def BuildModels(resolution, class_nums, ae_dims=128):
exec(nnlib.import_all(), locals(), globals())
x = inp = Input((resolution, resolution, 3))
x = PoseEstimator.EncFlow(ae_dims)(x)
encoder = Model(inp, x)
x = inp = Input(K.int_shape(encoder.outputs[0][1:]))
x = PoseEstimator.DecFlow(resolution, ae_dims)(x)
decoder = Model(inp, x)
x = inp = Input(K.int_shape(encoder.outputs[0][1:]))
x = PoseEstimator.LatentFlow(class_nums=class_nums)(x)
model_l = Model(inp, x)
return encoder, decoder, model_l
def onInitialize(self):
exec(nnlib.import_all(), locals(), globals())
self.set_vram_batch_requirements({1.5: 4})
self.resolution = 256
self.face_type = FaceType.FULL
self.fan_seg = FANSegmentator(
self.resolution,
FaceType.toString(self.face_type),
load_weights=not self.is_first_run(),
weights_file_root=self.get_model_root_path())
if self.is_training_mode:
f = SampleProcessor.TypeFlags
f_type = f.FACE_ALIGN_FULL #if self.face_type == FaceType.FULL else f.FACE_ALIGN_HALF
self.set_training_data_generators([
SampleGeneratorFace(
self.training_data_src_path,
debug=self.is_debug(),
batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(
random_flip=self.random_flip,
normalize_tanh=True,
scale_range=np.array([-0.05, 0.05]) +
self.src_scale_mod / 100.0),
output_sample_types=[
[f.TRANSFORMED | f_type | f.MODE_BGR, self.resolution],
[
f.TRANSFORMED | f_type | f.MODE_M
| f.FACE_MASK_FULL, self.resolution
]
]),
SampleGeneratorFace(
self.training_data_dst_path,
debug=self.is_debug(),
batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(
random_flip=self.random_flip,
normalize_tanh=True,
scale_range=np.array([-0.05, 0.05]) +
self.src_scale_mod / 100.0),
output_sample_types=[[
f.TRANSFORMED | f_type | f.MODE_BGR, self.resolution
]])
])
def onInitialize(self):
exec(nnlib.import_all(), locals(), globals())
self.set_vram_batch_requirements({1.5: 4})
self.resolution = 256
self.face_type = FaceType.FULL if self.options[
'face_type'] == 'f' else FaceType.HALF
self.fan_seg = FANSegmentator(
self.resolution,
FaceType.toString(self.face_type),
load_weights=not self.is_first_run(),
weights_file_root=self.get_model_root_path(),
training=True)
if self.is_training_mode:
f = SampleProcessor.TypeFlags
face_type = f.FACE_TYPE_FULL if self.options[
'face_type'] == 'f' else f.FACE_TYPE_HALF
self.set_training_data_generators([
SampleGeneratorFace(
self.training_data_src_path,
debug=self.is_debug(),
batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(
random_flip=True, motion_blur=[25, 1]),
output_sample_types=[[
f.WARPED_TRANSFORMED | face_type | f.MODE_BGR_SHUFFLE
| f.OPT_APPLY_MOTION_BLUR, self.resolution
],
[
f.WARPED_TRANSFORMED | face_type
| f.MODE_M | f.FACE_MASK_FULL,
self.resolution
]]),
SampleGeneratorFace(
self.training_data_dst_path,
debug=self.is_debug(),
batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(
random_flip=True),
output_sample_types=[[
f.TRANSFORMED | face_type | f.MODE_BGR_SHUFFLE,
self.resolution
]])
])
def EncFlow():
exec(nnlib.import_all(), locals(), globals())
XConv2D = partial(Conv2D, padding='zero')
def Act(lrelu_alpha=0.1):
return LeakyReLU(alpha=lrelu_alpha)
def downscale(dim, **kwargs):
def func(x):
return MaxPooling2D()(Act()(XConv2D(dim,
kernel_size=5,
strides=1)(x)))
return func
def upscale(dim, **kwargs):
def func(x):
return SubpixelUpscaler()(Act()(XConv2D(dim * 4,
kernel_size=3,
strides=1)(x)))
return func
def to_bgr(output_nc, **kwargs):
def func(x):
return XConv2D(output_nc, kernel_size=5,
activation='sigmoid')(x)
return func
upscale = partial(upscale)
downscale = partial(downscale)
ae_dims = 512
def func(input):
x = input
x = downscale(64)(x)
x = downscale(128)(x)
x = downscale(256)(x)
x = downscale(512)(x)
x = Dense(ae_dims, name="latent", use_bias=False)(Flatten()(x))
x = Lambda(lambda x: x + 0.1 * K.random_normal(K.shape(x), 0, 1),
output_shape=(None, ae_dims))(x)
return x
return func
def DecFlow(resolution, ae_dims):
exec(nnlib.import_all(), locals(), globals())
def upscale(dim, strides=2, **kwargs):
def func(x):
return ReLU()((Conv2DTranspose(dim,
kernel_size=3,
strides=strides,
padding='same')(x)))
return func
def to_bgr(output_nc, **kwargs):
def func(x):
return Conv2D(output_nc,
kernel_size=5,
padding='same',
activation='sigmoid')(x)
return func
upscale = partial(upscale)
lowest_dense_res = resolution // 16
def func(input):
x = input
x = Dense(256)(x)
x = ReLU()(x)
x = Dense(256)(x)
x = ReLU()(x)
x = Dense((lowest_dense_res * lowest_dense_res * 256))(x)
x = ReLU()(x)
x = Reshape((lowest_dense_res, lowest_dense_res, 256))(x)
x = upscale(512)(x)
x = upscale(256)(x)
x = upscale(128)(x)
x = upscale(64)(x)
x = to_bgr(3)(x)
return x
return func
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.device_idx = client_dict['device_idx']
self.cpu_only = client_dict['device_type'] == 'CPU'
self.final_output_path = Path(
client_dict['final_output_dir']
) if 'final_output_dir' in client_dict.keys() else None
self.debug_dir = client_dict['debug_dir']
self.cached_image = (None, None)
self.e = None
device_config = nnlib.DeviceConfig(cpu_only=self.cpu_only,
force_gpu_idx=self.device_idx,
allow_growth=True)
self.device_vram = device_config.gpu_vram_gb[0]
intro_str = 'Running on %s.' % (client_dict['device_name'])
if not self.cpu_only and self.device_vram <= 2:
intro_str += " Recommended to close all programs using this device."
self.log_info(intro_str)
if 'rects' in self.type:
if self.type == 'rects-mt':
nnlib.import_all(device_config)
self.e = facelib.MTCExtractor()
elif self.type == 'rects-dlib':
nnlib.import_dlib(device_config)
self.e = facelib.DLIBExtractor(nnlib.dlib)
elif self.type == 'rects-s3fd':
nnlib.import_all(device_config)
self.e = facelib.S3FDExtractor()
else:
raise ValueError("Wrong type.")
if self.e is not None:
self.e.__enter__()
elif self.type == 'landmarks':
nnlib.import_all(device_config)
self.e = facelib.LandmarksExtractor(nnlib.keras)
self.e.__enter__()
if self.device_vram >= 2:
self.second_pass_e = facelib.S3FDExtractor()
self.second_pass_e.__enter__()
else:
self.second_pass_e = None
elif self.type == 'final':
pass
def LIAEDecFlow(output_nc,
ch_dims,
multiscale_count=1,
add_residual_blocks=False,
padding='zero',
norm='',
**kwargs):
exec(nnlib.import_all(), locals(), globals())
upscale = SAEModel.upscale_pre(**kwargs)
to_bgr = SAEModel.to_bgr_pre(**kwargs)
dims = output_nc * ch_dims
ResidualBlock = SAEModel.ResidualBlock_pre(**kwargs)
def func(input):
x = input[0]
outputs = []
x1 = upscale(dims * 8)(x)
if add_residual_blocks:
x1 = ResidualBlock(dims * 8)(x1)
x1 = ResidualBlock(dims * 8)(x1)
if multiscale_count >= 3:
outputs += [to_bgr(output_nc)(x1)]
x2 = upscale(dims * 4)(x1)
if add_residual_blocks:
x2 = ResidualBlock(dims * 4)(x2)
x2 = ResidualBlock(dims * 4)(x2)
if multiscale_count >= 2:
outputs += [to_bgr(output_nc)(x2)]
x3 = upscale(dims * 2)(x2)
if add_residual_blocks:
x3 = ResidualBlock(dims * 2)(x3)
x3 = ResidualBlock(dims * 2)(x3)
outputs += [to_bgr(output_nc)(x3)]
return outputs
return func
def DFDecFlow(output_nc,dims,activation='tanh'):
exec (nnlib.import_all(), locals(), globals())
def upscale (dim):
def func(x):
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
return func
def func(input):
x = input[0]
x = upscale(dims)(x)
x = upscale(dims//2)(x)
x = upscale(dims//4)(x)
x = Conv2D(output_nc, kernel_size=5, padding='same', activation=activation)(x)
return x
return func
def DecoderFlow(ups, n_res_blks=2, mlp_nf=256, mlp_blks=2, subpixel_decoder=False ):
exec (nnlib.import_all(), locals(), globals())
def ResBlock(dim):
def func(input):
inp, mlp = input
x = inp
x = Conv2D(dim, 3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
x = FUNITAdain()([x,mlp])
x = ReLU()(x)
x = Conv2D(dim, 3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
x = FUNITAdain()([x,mlp])
return Add()([x,inp])
return func
def func(inputs):
x , class_code = inputs
nf = K.int_shape(x)[-1]
### MLP block inside decoder
mlp = class_code
for i in range(mlp_blks):
mlp = Dense(mlp_nf, activation='relu')(mlp)
for i in range(n_res_blks):
x = ResBlock(nf)( [x,mlp] )
for i in range(ups):
if subpixel_decoder:
x = Conv2D (4* (nf // 2**(i+1)), kernel_size=3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
x = SubpixelUpscaler()(x)
else:
x = UpSampling2D()(x)
x = Conv2D (nf // 2**(i+1), kernel_size=5, strides=1, padding='valid')(ZeroPadding2D(2)(x))
x = InstanceNormalization()(x)
x = ReLU()(x)
rgb = Conv2D (3, kernel_size=7, strides=1, padding='valid', activation='tanh')(ZeroPadding2D(3)(x))
return rgb
return func
def LIAEEncFlow(resolution, ch_dims, **kwargs):
exec (nnlib.import_all(), locals(), globals())
upscale = SAEModel.upscale_pre(**kwargs)
downscale = SAEModel.downscale_pre(**kwargs)
def func(input):
dims = K.int_shape(input)[-1]*ch_dims
x = input
x = downscale(dims)(x)
x = downscale(dims*2)(x)
x = downscale(dims*4)(x)
x = downscale(dims*8)(x)
x = Flatten()(x)
return x
return func
def onInitialize(self):
exec(nnlib.import_all(), locals(), globals())
self.set_vram_batch_requirements( {1.5:4} )
bgr_shape, mask_shape, self.encoder, self.decoder_src, self.decoder_dst = self.Build(self.options['lighter_ae'])
if not self.is_first_run():
weights_to_load = [ [self.encoder , 'encoder.h5'],
[self.decoder_src, 'decoder_src.h5'],
[self.decoder_dst, 'decoder_dst.h5']
]
self.load_weights_safe(weights_to_load)
input_src_bgr = Input(bgr_shape)
input_src_mask = Input(mask_shape)
input_dst_bgr = Input(bgr_shape)
input_dst_mask = Input(mask_shape)
rec_src_bgr, rec_src_mask = self.decoder_src( self.encoder(input_src_bgr) )
rec_dst_bgr, rec_dst_mask = self.decoder_dst( self.encoder(input_dst_bgr) )
self.ae = Model([input_src_bgr,input_src_mask,input_dst_bgr,input_dst_mask], [rec_src_bgr, rec_src_mask, rec_dst_bgr, rec_dst_mask] )
self.ae.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[ DSSIMMSEMaskLoss(input_src_mask, is_mse=self.options['pixel_loss']), 'mae', DSSIMMSEMaskLoss(input_dst_mask, is_mse=self.options['pixel_loss']), 'mae' ] )
self.src_view = K.function([input_src_bgr],[rec_src_bgr, rec_src_mask])
self.dst_view = K.function([input_dst_bgr],[rec_dst_bgr, rec_dst_mask])
if self.is_training_mode:
f = SampleProcessor.TypeFlags
self.set_training_data_generators ([
SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None,
debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ),
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 64] ] ),
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 64],
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 64] ] )
])
def onInitialize(self, **in_options):
exec(nnlib.import_all(), locals(), globals())
self.set_vram_batch_requirements( {2.5:2,3:2,4:2,4:4,5:8,6:12,7:16,8:16,9:24,10:24,11:32,12:32,13:48} )
bgr_shape, mask_shape, self.encoder, self.decoder_src, self.decoder_dst = self.Build(self.created_vram_gb)
if not self.is_first_run():
self.encoder.load_weights (self.get_strpath_storage_for_file(self.encoderH5))
self.decoder_src.load_weights (self.get_strpath_storage_for_file(self.decoder_srcH5))
self.decoder_dst.load_weights (self.get_strpath_storage_for_file(self.decoder_dstH5))
input_src_bgr = Input(bgr_shape)
input_src_mask = Input(mask_shape)
input_dst_bgr = Input(bgr_shape)
input_dst_mask = Input(mask_shape)
rec_src_bgr, rec_src_mask = self.decoder_src( self.encoder(input_src_bgr) )
rec_dst_bgr, rec_dst_mask = self.decoder_dst( self.encoder(input_dst_bgr) )
self.ae = Model([input_src_bgr,input_src_mask,input_dst_bgr,input_dst_mask], [rec_src_bgr, rec_src_mask, rec_dst_bgr, rec_dst_mask] )
if self.is_training_mode:
self.ae, = self.to_multi_gpu_model_if_possible ( [self.ae,] )
self.ae.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999),
loss=[ DSSIMMaskLoss([input_src_mask]), 'mae', DSSIMMaskLoss([input_dst_mask]), 'mae' ] )
self.src_view = K.function([input_src_bgr],[rec_src_bgr, rec_src_mask])
self.dst_view = K.function([input_dst_bgr],[rec_dst_bgr, rec_dst_mask])
if self.is_training_mode:
f = SampleProcessor.TypeFlags
self.set_training_data_generators ([
SampleGeneratorFace(self.training_data_src_path, sort_by_yaw_target_samples_path=self.training_data_dst_path if self.sort_by_yaw else None,
debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ),
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128],
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128],
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 128] ] ),
SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
output_sample_types=[ [f.WARPED_TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128],
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_BGR, 128],
[f.TRANSFORMED | f.FACE_ALIGN_HALF | f.MODE_M | f.FACE_MASK_FULL, 128] ] )
])
def on_initialize(self, client_dict):
self.log_info('Running on %s.' % (client_dict['device_name']))
self.type = client_dict['type']
self.image_size = client_dict['image_size']
self.face_type = client_dict['face_type']
self.device_idx = client_dict['device_idx']
self.cpu_only = client_dict['device_type'] == 'CPU'
self.output_path = Path(
client_dict['output_dir']) if 'output_dir' in client_dict.keys(
) else None
self.debug_dir = client_dict['debug_dir']
self.detector = client_dict['detector']
self.accurate_landmarks_extractor = client_dict[
'accurate_landmarks_extractor']
self.cached_image = (None, None)
self.e = None
device_config = nnlib.DeviceConfig(cpu_only=self.cpu_only,
force_gpu_idx=self.device_idx,
allow_growth=True)
if self.type == 'rects':
if self.detector is not None:
if self.detector == 'mt':
nnlib.import_all(device_config)
self.e = facelib.MTCExtractor()
elif self.detector == 'dlib':
nnlib.import_dlib(device_config)
self.e = facelib.DLIBExtractor(nnlib.dlib)
elif self.detector == 's3fd':
nnlib.import_all(device_config)
self.e = facelib.S3FDExtractor()
else:
raise ValueError("Wrong detector type.")
if self.e is not None:
self.e.__enter__()
elif self.type == 'landmarks':
nnlib.import_all(device_config)
self.e = facelib.LandmarksExtractor(nnlib.keras)
self.e.__enter__()
if self.accurate_landmarks_extractor and device_config.gpu_vram_gb[
0] >= 2:
self.second_pass_e = facelib.S3FDExtractor()
self.second_pass_e.__enter__()
else:
self.second_pass_e = None
elif self.type == 'final':
pass
def PatchDiscriminator(ndf=64):
exec (nnlib.import_all(), locals(), globals())
#use_bias = True
#def XNormalization(x):
# return InstanceNormalization (axis=-1)(x)
use_bias = False
def XNormalization(x):
return BatchNormalization (axis=-1)(x)
XConv2D = partial(Conv2D, use_bias=use_bias)
def func(input):
b,h,w,c = K.int_shape(input)
x = input
x = ZeroPadding2D((1,1))(x)
x = XConv2D( ndf, 4, strides=2, padding='valid', use_bias=True)(x)
x = LeakyReLU(0.2)(x)
x = ZeroPadding2D((1,1))(x)
x = XConv2D( ndf*2, 4, strides=2, padding='valid')(x)
x = XNormalization(x)
x = LeakyReLU(0.2)(x)
x = ZeroPadding2D((1,1))(x)
x = XConv2D( ndf*4, 4, strides=2, padding='valid')(x)
x = XNormalization(x)
x = LeakyReLU(0.2)(x)
x = ZeroPadding2D((1,1))(x)
x = XConv2D( ndf*8, 4, strides=2, padding='valid')(x)
x = XNormalization(x)
x = LeakyReLU(0.2)(x)
x = ZeroPadding2D((1,1))(x)
x = XConv2D( ndf*8, 4, strides=2, padding='valid')(x)
x = XNormalization(x)
x = LeakyReLU(0.2)(x)
x = ZeroPadding2D((1,1))(x)
return XConv2D( 1, 4, strides=1, padding='valid', use_bias=True, activation='sigmoid')(x)#
return func
def ResNet(output_nc, ngf=64, n_blocks=6, use_dropout=False):
exec (nnlib.import_all(), locals(), globals())
def func(input):
def ResnetBlock(dim, use_dropout=False):
def func(input):
x = input
x = Conv2D(dim, 3, strides=1, padding='same')(x)
x = InstanceNormalization (axis=-1)(x)
x = ReLU()(x)
if use_dropout:
x = Dropout(0.5)(x)
x = Conv2D(dim, 3, strides=1, padding='same')(x)
x = InstanceNormalization (axis=-1)(x)
x = ReLU()(x)
return Add()([x,input])
return func
x = input
x = ReLU()(InstanceNormalization (axis=-1)(Conv2D(ngf, 7, strides=1, padding='same')(x)))
x = ReLU()(InstanceNormalization (axis=-1)(Conv2D(ngf*2, 3, strides=2, padding='same')(x)))
x = ReLU()(InstanceNormalization (axis=-1)(Conv2D(ngf*4, 3, strides=2, padding='same')(x)))
x = ReLU()(InstanceNormalization (axis=-1)(Conv2D(ngf*4, 3, strides=2, padding='same')(x)))
for i in range(n_blocks):
x = ResnetBlock(ngf*4, use_dropout=use_dropout)(x)
x = ReLU()(InstanceNormalization (axis=-1)(Conv2DTranspose(ngf*4, 3, strides=2, padding='same')(x)))
x = ReLU()(InstanceNormalization (axis=-1)(Conv2DTranspose(ngf*2, 3, strides=2, padding='same')(x)))
x = ReLU()(InstanceNormalization (axis=-1)(Conv2DTranspose(ngf , 3, strides=2, padding='same')(x)))
x = Conv2D(output_nc, 7, strides=1, activation='sigmoid', padding='same')(x)
return x
return func
def __init__(self,
resolution,
face_type_str,
load_weights=True,
weights_file_root=None):
exec(nnlib.import_all(), locals(), globals())
self.model = FANSegmentator.BuildModel(resolution, ngf=32)
if weights_file_root:
weights_file_root = Path(weights_file_root)
else:
weights_file_root = Path(__file__).parent
self.weights_path = weights_file_root / ('FANSeg_%d_%s.h5' %
(resolution, face_type_str))
if load_weights:
self.model.load_weights(str(self.weights_path))
def DiscriminatorFlow(nf, n_res_blks, num_classes ):
exec (nnlib.import_all(), locals(), globals())
n_layers = n_res_blks // 2
def ActFirstResBlock(fout):
def func(x):
fin = K.int_shape(x)[-1]
fhid = min(fin, fout)
if fin != fout:
x_s = Conv2D (fout, kernel_size=1, strides=1, padding='valid', use_bias=False)(x)
else:
x_s = x
x = LeakyReLU(0.2)(x)
x = Conv2D (fhid, kernel_size=3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
x = LeakyReLU(0.2)(x)
x = Conv2D (fout, kernel_size=3, strides=1, padding='valid')(ZeroPadding2D(1)(x))
return Add()([x_s, x])
return func
def func( x ):
l_nf = nf
x = Conv2D (l_nf, kernel_size=7, strides=1, padding='valid')(ZeroPadding2D(3)(x))
for i in range(n_layers-1):
l_nf_out = min( l_nf*2, 1024 )
x = ActFirstResBlock(l_nf)(x)
x = ActFirstResBlock(l_nf_out)(x)
x = AveragePooling2D( pool_size=3, strides=2, padding='valid' )(ZeroPadding2D(1)(x))
l_nf = min( l_nf*2, 1024 )
l_nf_out = min( l_nf*2, 1024 )
x = ActFirstResBlock(l_nf)(x)
feat = x = ActFirstResBlock(l_nf_out)(x)
x = LeakyReLU(0.2)(x)
x = Conv2D (num_classes, kernel_size=1, strides=1, padding='valid')(x)
return x, feat
return func
def PYLatentFlow():
exec(nnlib.import_all(), locals(), globals())
k_size = 5
strides = 2
def func(input):
x = input
x = Dense(512, kernel_regularizer='l2')(x)
x = Dense(256, kernel_regularizer='l2')(x)
x = Dense(128, kernel_regularizer='l2')(x)
x = Dense(64, kernel_regularizer='l2')(x)
x = Dense(32, kernel_regularizer='l2')(x)
x = Dense(16, kernel_regularizer='l2')(x)
x = Dense(8, kernel_regularizer='l2')(x)
x = Dense(4, kernel_regularizer='l2')(x)
x = Dense(2, activation='tanh')(x)
return x
return func
