def predictor_func (self, face=None): face = nn.to_data_format(face[None,...], self.model_data_format, "NHWC") bgr, mask_dst_dstm, mask_src_dstm = [ nn.to_data_format(x,"NHWC", self.model_data_format).astype(np.float32) for x in self.AE_merge (face) ] mask = mask_dst_dstm[0] * mask_src_dstm[0] return bgr[0], mask[...,0]
def onGetPreview(self, samples): ( (warped_src, target_src, target_srcm), (warped_dst, target_dst, target_dstm) ) = samples S, D, SS, DD, DDM, SD, SDM = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([target_src,target_dst] + self.AE_view (target_src, target_dst) ) ] DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ] target_srcm, target_dstm = [ nn.to_data_format(x,"NHWC", self.model_data_format) for x in ([target_srcm, target_dstm] )] n_samples = min(4, self.get_batch_size() ) result = [] st = [] for i in range(n_samples): ar = S[i], SS[i], D[i], DD[i], SD[i] st.append ( np.concatenate ( ar, axis=1) ) result += [ ('Quick96', np.concatenate (st, axis=0 )), ] st_m = [] for i in range(n_samples): ar = S[i]*target_srcm[i], SS[i], D[i]*target_dstm[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i]) st_m.append ( np.concatenate ( ar, axis=1) ) result += [ ('Quick96 masked', np.concatenate (st_m, axis=0 )), ] return result
def onGetPreview(self, samples): (( warped_src, target_src, target_srcm_all, ), ( warped_dst, target_dst, target_dstm_all, )) = samples if self.learn_mask: S, D, SS, DD, DDM, SD, SDM = [ np.clip(nn.to_data_format(x, "NHWC", self.model_data_format), 0.0, 1.0) for x in ([target_src, target_dst] + self.AE_view(target_src, target_dst)) ] DDM, SDM, = [np.repeat(x, (3, ), -1) for x in [DDM, SDM]] else: S, D, SS, DD, SD, = [ np.clip(nn.to_data_format(x, "NHWC", self.model_data_format), 0.0, 1.0) for x in ([target_src, target_dst] + self.AE_view(target_src, target_dst)) ] target_srcm_all, target_dstm_all = [ nn.to_data_format(x, "NHWC", self.model_data_format) for x in ([target_srcm_all, target_dstm_all]) ] target_srcm = np.clip(target_srcm_all, 0, 1) target_dstm = np.clip(target_dstm_all, 0, 1) n_samples = min(4, self.get_batch_size(), 800 // self.resolution) result = [] st = [] for i in range(n_samples): ar = D[i], DD[i], SD[i] st.append(np.concatenate(ar, axis=1)) result += [ ('HEAD', np.concatenate(st, axis=0)), ] if self.learn_mask: st_m = [] for i in range(n_samples): ar = D[i] * target_dstm[i], DD[i] * DDM[i], SD[i] * (DDM[i] * SDM[i]) st_m.append(np.concatenate(ar, axis=1)) result += [ ('HEAD masked', np.concatenate(st_m, axis=0)), ] return result
def onGetPreview(self, samples): n_samples = min(4, self.get_batch_size(), 800 // self.resolution ) src_samples, dst_samples = samples image_np, mask_np = src_samples I, M, IM, = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([image_np,mask_np] + self.view (image_np) ) ] M, IM, = [ np.repeat (x, (3,), -1) for x in [M, IM] ] green_bg = np.tile( np.array([0,1,0], dtype=np.float32)[None,None,...], (self.resolution,self.resolution,1) ) result = [] st = [] for i in range(n_samples): ar = I[i]*M[i]+ green_bg*(1-M[i]), IM[i], I[i]*IM[i] + green_bg*(1-IM[i]) st.append ( np.concatenate ( ar, axis=1) ) result += [ ('XSeg training faces', np.concatenate (st, axis=0 )), ] if len(dst_samples) != 0: dst_np, = dst_samples D, DM, = [ np.clip(nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([dst_np] + self.view (dst_np) ) ] DM, = [ np.repeat (x, (3,), -1) for x in [DM] ] st = [] for i in range(n_samples): ar = D[i], DM[i], D[i]*DM[i]+ green_bg*(1-DM[i]) st.append ( np.concatenate ( ar, axis=1) ) result += [ ('XSeg unseen faces', np.concatenate (st, axis=0 )), ] return result
def onGetPreview(self, samples): bs = self.get_batch_size() ( (target_src, check_src), ) = samples warped_src = check_src[0:bs] target_src = check_src[bs:] S, SS, = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([warped_src,target_src]) ] RS, GS = self.AE_view (warped_src, target_src) RS, = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([RS]) ] #import code #code.interact(local=dict(globals(), **locals())) st = [] for i in range(min(4, bs)): #map = np.zeros( (self.resolution,self.resolution,3), dtype=np.float32 ) rs_map = RS[i].copy() for n in range(self.landmarks_count): x,y = GS[i, n] color = colorsys.hsv_to_rgb ( n * (1.0/self.landmarks_count), 1.0, 1.0 ) #cv2.circle (map, ( int( ((x+1)/2.0)*self.resolution), int( ((y+1)/2.0)*self.resolution) ), 1, color, lineType=cv2.LINE_AA ) cv2.circle (rs_map, ( int( ((x+1)/2.0)*self.resolution), int( ((y+1)/2.0)*self.resolution) ), 1, color, lineType=cv2.LINE_AA ) ar = S[i], SS[i], RS[i], rs_map st.append ( np.concatenate ( ar, axis=1) ) return [ ('TEST', np.concatenate (st, axis=0 )), ]
def onGetPreview(self, samples): ( (warped_src, target_src, target_srcm, target_srcm_em), (warped_dst, target_dst, target_dstm, target_dstm_em) ) = samples S, D, SS, DD, DDM_000, _, _ = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([target_src,target_dst] + self.AE_view (target_src, target_dst, 0.0) ) ] _, _, DDM_025, SD_025, SDM_025 = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in self.AE_view (target_src, target_dst, 0.25) ] _, _, DDM_050, SD_050, SDM_050 = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in self.AE_view (target_src, target_dst, 0.50) ] _, _, DDM_065, SD_065, SDM_065 = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in self.AE_view (target_src, target_dst, 0.65) ] _, _, DDM_075, SD_075, SDM_075 = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in self.AE_view (target_src, target_dst, 0.75) ] _, _, DDM_100, SD_100, SDM_100 = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in self.AE_view (target_src, target_dst, 1.00) ] (DDM_000, DDM_025, SDM_025, DDM_050, SDM_050, DDM_065, SDM_065, DDM_075, SDM_075, DDM_100, SDM_100) = [ np.repeat (x, (3,), -1) for x in (DDM_000, DDM_025, SDM_025, DDM_050, SDM_050, DDM_065, SDM_065, DDM_075, SDM_075, DDM_100, SDM_100) ] target_srcm, target_dstm = [ nn.to_data_format(x,"NHWC", self.model_data_format) for x in ([target_srcm, target_dstm] )] n_samples = min(4, self.get_batch_size(), 800 // self.resolution ) result = [] i = np.random.randint(n_samples) st = [ np.concatenate ((S[i], D[i], DD[i]*DDM_000[i]), axis=1) ] st += [ np.concatenate ((SS[i], DD[i], SD_075[i] ), axis=1) ] result += [ ('AMP morph 0.75', np.concatenate (st, axis=0 )), ] st = [ np.concatenate ((DD[i], SD_025[i], SD_050[i]), axis=1) ] st += [ np.concatenate ((SD_065[i], SD_075[i], SD_100[i]), axis=1) ] result += [ ('AMP morph list', np.concatenate (st, axis=0 )), ] st = [ np.concatenate ((DD[i], SD_025[i]*DDM_025[i]*SDM_025[i], SD_050[i]*DDM_050[i]*SDM_050[i]), axis=1) ] st += [ np.concatenate ((SD_065[i]*DDM_065[i]*SDM_065[i], SD_075[i]*DDM_075[i]*SDM_075[i], SD_100[i]*DDM_100[i]*SDM_100[i]), axis=1) ] result += [ ('AMP morph list masked', np.concatenate (st, axis=0 )), ] return result
def onGetPreview(self, samples): ((warped_src, target_src, pyr_src), (warped_dst, target_dst, pyr_dst)) = samples S, D, SS, DD, DS = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([target_src,target_dst] + \ self.AE_view (target_src) + \ self.AE_view (target_dst) + \ self.AE_view2 (target_dst, target_src) ) ] n_samples = min(4, self.get_batch_size()) result = [] st = [] for i in range(n_samples): ar = S[i], SS[i], D[i], DD[i], DS[i] st.append(np.concatenate(ar, axis=1)) result += [ ('TEST', np.concatenate(st, axis=0)), ] return result
def onGetPreview(self, samples): ( (warped_src, target_src, target_srcm_all,), (warped_dst, target_dst, target_dstm_all,) ) = samples S, D, SS, DD, DDM, SD, SDM = [ np.clip( nn.to_data_format(x,"NHWC", self.model_data_format), 0.0, 1.0) for x in ([target_src,target_dst] + self.AE_view (target_src, target_dst) ) ] DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ] target_srcm_all, target_dstm_all = [ nn.to_data_format(x,"NHWC", self.model_data_format) for x in ([target_srcm_all, target_dstm_all] )] target_srcm = np.clip(target_srcm_all, 0, 1) target_dstm = np.clip(target_dstm_all, 0, 1) n_samples = min(4, self.get_batch_size(), 800 // self.resolution ) if self.resolution <= 256: result = [] st = [] for i in range(n_samples): ar = S[i], SS[i], D[i], DD[i], SD[i] st.append ( np.concatenate ( ar, axis=1) ) result += [ ('SAEHD', np.concatenate (st, axis=0 )), ] st_m = [] for i in range(n_samples): SD_mask = DDM[i]*SDM[i] if self.face_type < FaceType.HEAD else SDM[i] ar = S[i]*target_srcm[i], SS[i], D[i]*target_dstm[i], DD[i]*DDM[i], SD[i]*SD_mask st_m.append ( np.concatenate ( ar, axis=1) ) result += [ ('SAEHD masked', np.concatenate (st_m, axis=0 )), ] else: result = [] st = [] for i in range(n_samples): ar = S[i], SS[i] st.append ( np.concatenate ( ar, axis=1) ) result += [ ('SAEHD src-src', np.concatenate (st, axis=0 )), ] st = [] for i in range(n_samples): ar = D[i], DD[i] st.append ( np.concatenate ( ar, axis=1) ) result += [ ('SAEHD dst-dst', np.concatenate (st, axis=0 )), ] st = [] for i in range(n_samples): ar = D[i], SD[i] st.append ( np.concatenate ( ar, axis=1) ) result += [ ('SAEHD pred', np.concatenate (st, axis=0 )), ] st_m = [] for i in range(n_samples): ar = S[i]*target_srcm[i], SS[i] st_m.append ( np.concatenate ( ar, axis=1) ) result += [ ('SAEHD masked src-src', np.concatenate (st_m, axis=0 )), ] st_m = [] for i in range(n_samples): ar = D[i]*target_dstm[i], DD[i]*DDM[i] st_m.append ( np.concatenate ( ar, axis=1) ) result += [ ('SAEHD masked dst-dst', np.concatenate (st_m, axis=0 )), ] st_m = [] for i in range(n_samples): SD_mask = DDM[i]*SDM[i] if self.face_type < FaceType.HEAD else SDM[i] ar = D[i]*target_dstm[i], SD[i]*SD_mask st_m.append ( np.concatenate ( ar, axis=1) ) result += [ ('SAEHD masked pred', np.concatenate (st_m, axis=0 )), ] return result
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 onGetPreview(self, samples): n_samples = min(4 if self.stage_resolution < 256 else 1, self.get_batch_size()) ( (prev_warped_src, warped_src, prev_target_src, target_src, target_srcm), (prev_warped_dst, warped_dst, prev_target_dst, target_dst, target_dstm) ) = \ [ [sample[0:n_samples] for sample in sample_list ] for sample_list in samples ] S, D, SS, DD, DDM, SD, SDM = [ np.clip(nn.to_data_format(x, "NHWC", self.model_data_format), 0.0, 1.0) for x in ([target_src, target_dst] + self.AE_view( prev_target_src, target_src, prev_target_dst, target_dst)) ] DDM, SDM, = [np.repeat(x, (3, ), -1) for x in [DDM, SDM]] target_srcm, target_dstm = [ np.clip(nn.to_data_format(x, "NHWC", self.model_data_format), 0.0, 1.0) for x in ([target_srcm, target_dstm]) ] result = [] if self.stage_resolution < 256: S, D, target_srcm, target_dstm = [[ cv2_resize(y, (self.stage_resolution, self.stage_resolution)) for y in x ] for x in ([S, D, target_srcm, target_dstm])] st = [] for i in range(n_samples): ar = S[i], SS[i], D[i], DD[i], SD[i] st.append(np.concatenate(ar, axis=1)) result += [ ('test', np.concatenate(st, axis=0)), ] st_m = [] for i in range(n_samples): ar = S[i] * target_srcm[i], SS[i], D[i] * target_dstm[i], DD[ i] * DDM[i], SD[i] * (DDM[i] * SDM[i]) st_m.append(np.concatenate(ar, axis=1)) result += [ ('test masked', np.concatenate(st_m, axis=0)), ] else: SS, DD, DDM, SD, SDM = [[ cv2.resize(y, (self.resolution, self.resolution), interpolation=cv2.INTER_NEAREST) for y in x ] for x in ([SS, DD, DDM, SD, SDM])] st = [] for i in range(n_samples): ar = S[i], SS[i] st.append(np.concatenate(ar, axis=1)) result += [ ('SAEHD src-src', np.concatenate(st, axis=0)), ] st = [] for i in range(n_samples): ar = D[i], DD[i] st.append(np.concatenate(ar, axis=1)) result += [ ('SAEHD dst-dst', np.concatenate(st, axis=0)), ] st = [] for i in range(n_samples): ar = D[i], SD[i] st.append(np.concatenate(ar, axis=1)) result += [ ('SAEHD pred', np.concatenate(st, axis=0)), ] st_m = [] for i in range(n_samples): ar = S[i] * target_srcm[i], SS[i] st_m.append(np.concatenate(ar, axis=1)) result += [ ('SAEHD masked src-src', np.concatenate(st_m, axis=0)), ] st_m = [] for i in range(n_samples): ar = D[i] * target_dstm[i], DD[i] * DDM[i] st_m.append(np.concatenate(ar, axis=1)) result += [ ('SAEHD masked dst-dst', np.concatenate(st_m, axis=0)), ] st_m = [] for i in range(n_samples): SD_mask = DDM[i] * SDM[ i] if self.face_type < FaceType.HEAD else SDM[i] ar = D[i] * target_dstm[i], SD[i] * SD_mask st_m.append(np.concatenate(ar, axis=1)) result += [ ('SAEHD masked pred', np.concatenate(st_m, axis=0)), ] """ result = [] st = [] for i in range(n_samples): ar = S[i], SS[i], D[i], DD[i], SD[i] st.append ( np.concatenate ( ar, axis=1) ) result += [ ('test', np.concatenate (st, axis=0 )), ] st_m = [] for i in range(n_samples): ar = S[i]*target_srcm[i], SS[i], D[i]*target_dstm[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i]) st_m.append ( np.concatenate ( ar, axis=1) ) result += [ ('test masked', np.concatenate (st_m, axis=0 )), ] """ return result