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