def onInitialize(self): exec(nnlib.import_all(), locals(), globals()) self.set_vram_batch_requirements( {4:64} ) self.resolution = 128 self.face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF self.pose_est = PoseEstimator(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: t = SampleProcessor.Types face_type = t.FACE_TYPE_FULL if self.options['face_type'] == 'f' else t.FACE_TYPE_HALF self.set_training_data_generators ([ SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, generators_count=4, sample_process_options=SampleProcessor.Options( rotation_range=[0,0] ), #random_flip=True, output_sample_types=[ {'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR_SHUFFLE), 'resolution':self.resolution, 'motion_blur':(25, 1) }, {'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR_SHUFFLE), 'resolution':self.resolution }, {'types': (t.IMG_PITCH_YAW_ROLL,)} ]), SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, generators_count=4, sample_process_options=SampleProcessor.Options( rotation_range=[0,0] ), #random_flip=True, output_sample_types=[ {'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'resolution':self.resolution }, {'types': (t.IMG_PITCH_YAW_ROLL,)} ]) ])
class Model(ModelBase): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs, ask_write_preview_history=False, ask_target_iter=False, ask_sort_by_yaw=False, ask_random_flip=False, ask_src_scale_mod=False) #override def onInitializeOptions(self, is_first_run, ask_override): default_face_type = 'f' if is_first_run: self.options['face_type'] = io.input_str( "Half or Full face? (h/f, ?:help skip:f) : ", default_face_type, ['h', 'f'], help_message= "Half face has better resolution, but covers less area of cheeks." ).lower() else: self.options['face_type'] = self.options.get( 'face_type', default_face_type) #override def onInitialize(self): exec(nnlib.import_all(), locals(), globals()) self.set_vram_batch_requirements({4: 32}) self.resolution = 128 self.face_type = FaceType.FULL if self.options[ 'face_type'] == 'f' else FaceType.HALF self.pose_est = PoseEstimator( 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: t = SampleProcessor.Types face_type = t.FACE_TYPE_FULL if self.options[ 'face_type'] == 'f' else t.FACE_TYPE_HALF self.set_training_data_generators([ SampleGeneratorFace( self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, generators_count=4, sample_process_options=SampleProcessor.Options( rotation_range=[0, 0]), #random_flip=True, output_sample_types=[{ 'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR_SHUFFLE), 'resolution': self.resolution, 'motion_blur': (25, 1) }, { 'types': (t.IMG_PITCH_YAW_ROLL, ) }]), SampleGeneratorFace( self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, generators_count=4, sample_process_options=SampleProcessor.Options( rotation_range=[0, 0]), #random_flip=True, output_sample_types=[{ 'types': (t.IMG_TRANSFORMED, face_type, t.MODE_BGR_SHUFFLE), 'resolution': self.resolution }, { 'types': (t.IMG_PITCH_YAW_ROLL, ) }]) ]) #override def onSave(self): self.pose_est.save_weights() #override def onTrainOneIter(self, generators_samples, generators_list): target_src, pitch_yaw_roll = generators_samples[0] pitch_loss, yaw_loss, roll_loss = self.pose_est.train_on_batch( target_src, pitch_yaw_roll) return (('pitch_loss', pitch_loss), ('yaw_loss', yaw_loss), ('roll_loss', roll_loss)) #override def onGetPreview(self, generators_samples): test_src = generators_samples[0][0][0:4] #first 4 samples test_pyr_src = generators_samples[0][1][0:4] test_dst = generators_samples[1][0][0:4] test_pyr_dst = generators_samples[1][1][0:4] h, w, c = self.resolution, self.resolution, 3 h_line = 13 result = [] for name, img, pyr in [ ['training data', test_src, test_pyr_src], \ ['evaluating data',test_dst, test_pyr_dst] ]: pyr_pred = self.pose_est.extract(img) hor_imgs = [] for i in range(len(img)): img_info = np.ones((h, w, c)) * 0.1 lines = ["%s" % (str(pyr[i])), "%s" % (str(pyr_pred[i]))] lines_count = len(lines) for ln in range(lines_count): img_info[ ln*h_line:(ln+1)*h_line, 0:w] += \ imagelib.get_text_image ( (h_line,w,c), lines[ln], color=[0.8]*c ) hor_imgs.append( np.concatenate((img[i, :, :, 0:3], img_info), axis=1)) result += [(name, np.concatenate(hor_imgs, axis=0))] return result
def onInitialize(self, batch_size=-1, **in_options): exec(nnlib.code_import_all, locals(), globals()) self.set_vram_batch_requirements({4: 4}) resolution = self.options['resolution'] bgr_shape = (resolution, resolution, 3) mask_shape = (resolution, resolution, 1) bgrm_shape = (resolution, resolution, 4) ngf = 64 ndf = 64 lambda_A = 100 lambda_B = 100 use_batch_norm = True #created_batch_size > 1 poseest = self.poseest = PoseEstimator( resolution, FaceType.toString(FaceType.FULL)) self.enc = modelify(AVATARModel.DFEncFlow())([ Input(bgr_shape), Input((poseest.class_nums[0], )), Input((poseest.class_nums[0], )), Input((poseest.class_nums[0], )) ]) dec_Inputs = [Input(K.int_shape(x)[1:]) for x in self.enc.outputs] self.decA = modelify(AVATARModel.DFDecFlow(bgr_shape[2]))(dec_Inputs) self.decB = modelify(AVATARModel.DFDecFlow(bgr_shape[2]))(dec_Inputs) def GA(x): return self.decA(self.enc(x)) self.GA = GA def GB(x): return self.decB(self.enc(x)) self.GB = GB #self.GA = modelify(AVATARModel.ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=True))( Input(bgr_shape) ) #self.GB = modelify(AVATARModel.ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=True))( Input(bgr_shape) ) #self.GA = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape)) #self.GB = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape)) self.DA = modelify( AVATARModel.NLayerDiscriminator(use_batch_norm, ndf=ndf))(Input(bgr_shape)) self.DB = modelify( AVATARModel.NLayerDiscriminator(use_batch_norm, ndf=ndf))(Input(bgr_shape)) if not self.is_first_run(): weights_to_load = [ # (self.GA, 'GA.h5'), # (self.GB, 'GB.h5'), (self.enc, 'enc.h5'), (self.decA, 'decA.h5'), (self.decB, 'decB.h5'), (self.DA, 'DA.h5'), (self.DB, 'DB.h5'), ] self.load_weights_safe(weights_to_load) real_A0 = Input(bgr_shape) real_A0m = Input(mask_shape) real_B0 = Input(bgr_shape) real_B0m = Input(mask_shape) real_A0p, real_A0y, real_A0r = poseest.flow(real_A0) real_B0p, real_B0y, real_B0r = poseest.flow(real_B0) DA_ones = K.ones_like(K.shape(self.DA.outputs[0])) DA_zeros = K.zeros_like(K.shape(self.DA.outputs[0])) DB_ones = K.ones_like(K.shape(self.DB.outputs[0])) DB_zeros = K.zeros_like(K.shape(self.DB.outputs[0])) def DLoss(labels, logits): return K.mean(K.binary_crossentropy(labels, logits)) def CycleLOSS(t1, t2): return dssim(kernel_size=int(resolution / 11.6), max_value=2.0)(t1 + 1, t2 + 1) return K.mean(K.abs(t1 - t2)) fake_B0 = self.GA([real_A0, real_B0p, real_B0y, real_B0r]) fake_A0 = self.GB([real_B0, real_A0p, real_A0y, real_A0r]) fake_B0p, fake_B0y, fake_B0r = poseest.flow(fake_B0) fake_A0p, fake_A0y, fake_A0r = poseest.flow(fake_A0) real_A0_d = self.DA(real_A0) real_A0_d_ones = K.ones_like(real_A0_d) fake_A0_d = self.DA(fake_A0) fake_A0_d_ones = K.ones_like(fake_A0_d) fake_A0_d_zeros = K.zeros_like(fake_A0_d) real_B0_d = self.DB(real_B0) real_B0_d_ones = K.ones_like(real_B0_d) fake_B0_d = self.DB(fake_B0) fake_B0_d_ones = K.ones_like(fake_B0_d) fake_B0_d_zeros = K.zeros_like(fake_B0_d) rec_A0 = self.GB([fake_B0, real_A0p, real_A0y, real_A0r]) rec_B0 = self.GA([fake_A0, real_B0p, real_B0y, real_B0r]) #import code #code.interact(local=dict(globals(), **locals())) loss_GA = DLoss(fake_B0_d_ones, fake_B0_d ) + \ lambda_A * 0.1 * K.mean( K.square(fake_B0p-real_A0p) + K.square(fake_B0y-real_A0y) + K.square(fake_B0r-real_A0r) ) + \ lambda_A * (CycleLOSS(rec_B0, real_B0) ) weights_GA = self.enc.trainable_weights + self.decA.trainable_weights # + #self.GA.trainable_weights loss_GB = DLoss(fake_A0_d_ones, fake_A0_d ) + \ lambda_B * 0.1 * K.mean( K.square(fake_A0p-real_B0p) + K.square(fake_A0y-real_B0y) + K.square(fake_A0r-real_B0r) ) + \ lambda_B * (CycleLOSS(rec_A0, real_A0) ) weights_GB = self.enc.trainable_weights + self.decB.trainable_weights # + #self.GB.trainable_weights def opt(): return Adam(lr=2e-5, beta_1=0.5, beta_2=0.999, tf_cpu_mode=2) #, clipnorm=1) self.GA_train = K.function([real_A0, real_A0m, real_B0, real_B0m], [loss_GA], opt().get_updates(loss_GA, weights_GA)) self.GB_train = K.function([real_A0, real_A0m, real_B0, real_B0m], [loss_GB], opt().get_updates(loss_GB, weights_GB)) ########### loss_D_A = ( DLoss(real_A0_d_ones, real_A0_d ) + \ DLoss(fake_A0_d_zeros, fake_A0_d ) ) * 0.5 self.DA_train = K.function( [real_A0, real_A0m, real_B0, real_B0m], [loss_D_A], opt().get_updates(loss_D_A, self.DA.trainable_weights)) ############ loss_D_B = ( DLoss(real_B0_d_ones, real_B0_d ) + \ DLoss(fake_B0_d_zeros, fake_B0_d ) ) * 0.5 self.DB_train = K.function( [real_A0, real_A0m, real_B0, real_B0m], [loss_D_B], opt().get_updates(loss_D_B, self.DB.trainable_weights)) ############ self.G_view = K.function([real_A0, real_A0m, real_B0, real_B0m], [fake_A0, rec_A0, fake_B0, rec_B0]) if self.is_training_mode: t = SampleProcessor.Types face_type = t.FACE_TYPE_FULL output_sample_types = [ { 'types': (t.IMG_SOURCE, face_type, t.MODE_BGR), 'resolution': resolution }, { 'types': (t.IMG_SOURCE, face_type, t.MODE_M, t.FACE_MASK_FULL), 'resolution': resolution }, ] 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), output_sample_types=output_sample_types), 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), output_sample_types=output_sample_types) ]) else: self.G_convert = K.function([real_A0, real_B0m], [fake_B0])