From b333fcea4b74436bd17388d5fe2cad7bcfcf1163 Mon Sep 17 00:00:00 2001 From: iperov Date: Tue, 23 Mar 2021 15:00:24 +0400 Subject: [PATCH 1/7] dump_ckpt --- main.py | 2 ++ mainscripts/Trainer.py | 11 +++++-- models/Model_SAEHD/Model.py | 62 +++++++++++++++++++++++++++++++------ 3 files changed, 62 insertions(+), 13 deletions(-) diff --git a/main.py b/main.py index 2ba085f..cad3e92 100644 --- a/main.py +++ b/main.py @@ -127,6 +127,7 @@ if __name__ == "__main__": 'silent_start' : arguments.silent_start, 'execute_programs' : [ [int(x[0]), x[1] ] for x in arguments.execute_program ], 'debug' : arguments.debug, + 'dump_ckpt' : arguments.dump_ckpt, } from mainscripts import Trainer Trainer.main(**kwargs) @@ -144,6 +145,7 @@ if __name__ == "__main__": p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Train on CPU.") p.add_argument('--force-gpu-idxs', dest="force_gpu_idxs", default=None, help="Force to choose GPU indexes separated by comma.") p.add_argument('--silent-start', action="store_true", dest="silent_start", default=False, help="Silent start. Automatically chooses Best GPU and last used model.") + p.add_argument('--dump-ckpt', action="store_true", dest="dump_ckpt", default=False, help="Dump the model to ckpt format.") p.add_argument('--execute-program', dest="execute_program", default=[], action='append', nargs='+') diff --git a/mainscripts/Trainer.py b/mainscripts/Trainer.py index 7d73e2f..4afc218 100644 --- a/mainscripts/Trainer.py +++ b/mainscripts/Trainer.py @@ -27,6 +27,7 @@ def trainerThread (s2c, c2s, e, silent_start=False, execute_programs = None, debug=False, + dump_ckpt=False, **kwargs): while True: try: @@ -44,7 +45,7 @@ def trainerThread (s2c, c2s, e, saved_models_path.mkdir(exist_ok=True, parents=True) model = models.import_model(model_class_name)( - is_training=True, + is_training=not dump_ckpt, saved_models_path=saved_models_path, training_data_src_path=training_data_src_path, training_data_dst_path=training_data_dst_path, @@ -55,9 +56,13 @@ def trainerThread (s2c, c2s, e, force_gpu_idxs=force_gpu_idxs, cpu_only=cpu_only, silent_start=silent_start, - debug=debug, - ) + debug=debug) + if dump_ckpt: + e.set() + model.dump_ckpt() + break + is_reached_goal = model.is_reached_iter_goal() shared_state = { 'after_save' : False } diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 380a053..0ef99a6 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -204,6 +204,8 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... archi_type, archi_opts = archi_split elif len(archi_split) == 1: archi_type, archi_opts = archi_split[0], None + + self.archi_type = archi_type ae_dims = self.options['ae_dims'] e_dims = self.options['e_dims'] @@ -236,22 +238,22 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... optimizer_vars_on_cpu = models_opt_device=='/CPU:0' input_ch=3 - bgr_shape = nn.get4Dshape(resolution,resolution,input_ch) + bgr_shape = self.bgr_shape = nn.get4Dshape(resolution,resolution,input_ch) mask_shape = nn.get4Dshape(resolution,resolution,1) self.model_filename_list = [] with tf.device ('/CPU:0'): #Place holders on CPU - self.warped_src = tf.placeholder (nn.floatx, bgr_shape) - self.warped_dst = tf.placeholder (nn.floatx, bgr_shape) + self.warped_src = tf.placeholder (nn.floatx, bgr_shape, name='warped_src') + self.warped_dst = tf.placeholder (nn.floatx, bgr_shape, name='warped_dst') - self.target_src = tf.placeholder (nn.floatx, bgr_shape) - self.target_dst = tf.placeholder (nn.floatx, bgr_shape) + self.target_src = tf.placeholder (nn.floatx, bgr_shape, name='target_src') + self.target_dst = tf.placeholder (nn.floatx, bgr_shape, name='target_dst') - self.target_srcm = tf.placeholder (nn.floatx, mask_shape) - self.target_srcm_em = tf.placeholder (nn.floatx, mask_shape) - self.target_dstm = tf.placeholder (nn.floatx, mask_shape) - self.target_dstm_em = tf.placeholder (nn.floatx, mask_shape) + self.target_srcm = tf.placeholder (nn.floatx, mask_shape, name='target_srcm') + self.target_srcm_em = tf.placeholder (nn.floatx, mask_shape, name='target_srcm_em') + self.target_dstm = tf.placeholder (nn.floatx, mask_shape, name='target_dstm') + self.target_dstm_em = tf.placeholder (nn.floatx, mask_shape, name='target_dstm_em') # Initializing model classes model_archi = nn.DeepFakeArchi(resolution, opts=archi_opts) @@ -609,7 +611,10 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... 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() @@ -650,7 +655,44 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... if self.pretrain_just_disabled: self.update_sample_for_preview(force_new=True) + + def dump_ckpt(self): + tf = nn.tf + + + with tf.device ('/CPU:0'): + warped_dst = tf.placeholder (nn.floatx, (None, self.resolution, self.resolution, 3), name='in_face') + warped_dst = tf.transpose(warped_dst, (0,3,1,2)) + + + if 'df' in self.archi_type: + gpu_dst_code = self.inter(self.encoder(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 self.archi_type: + gpu_dst_code = self.encoder (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) + + gpu_pred_src_dst = tf.transpose(gpu_pred_src_dst, (0,2,3,1)) + gpu_pred_dst_dstm = tf.transpose(gpu_pred_dst_dstm, (0,2,3,1)) + gpu_pred_src_dstm = tf.transpose(gpu_pred_src_dstm, (0,2,3,1)) + + + saver = tf.train.Saver() + tf.identity(gpu_pred_dst_dstm, name='out_face_mask') + tf.identity(gpu_pred_src_dst, name='out_celeb_face') + tf.identity(gpu_pred_src_dstm, name='out_celeb_face_mask') + + saver.save(nn.tf_sess, self.get_strpath_storage_for_file('.ckpt') ) + + #override def get_model_filename_list(self): return self.model_filename_list From e47b602ec8acaf7ad175afe4dbf3be3ec9192015 Mon Sep 17 00:00:00 2001 From: iperov Date: Tue, 23 Mar 2021 15:13:10 +0400 Subject: [PATCH 2/7] SAEHD: random flip replaced with random SRC flip(default False) and random DST flip(default True) --- models/ModelBase.py | 12 +++++++++++- models/Model_SAEHD/Model.py | 7 ++++--- 2 files changed, 15 insertions(+), 4 deletions(-) diff --git a/models/ModelBase.py b/models/ModelBase.py index 3cb88c5..ce219ab 100644 --- a/models/ModelBase.py +++ b/models/ModelBase.py @@ -185,7 +185,9 @@ class ModelBase(object): self.write_preview_history = self.options.get('write_preview_history', False) self.target_iter = self.options.get('target_iter',0) self.random_flip = self.options.get('random_flip',True) - + self.random_src_flip = self.options.get('random_src_flip', False) + self.random_dst_flip = self.options.get('random_dst_flip', True) + self.on_initialize() self.options['batch_size'] = self.batch_size @@ -297,6 +299,14 @@ class ModelBase(object): def ask_random_flip(self): default_random_flip = self.load_or_def_option('random_flip', True) self.options['random_flip'] = io.input_bool("Flip faces randomly", default_random_flip, help_message="Predicted face will look more naturally without this option, but src faceset should cover all face directions as dst faceset.") + + def ask_random_src_flip(self): + default_random_src_flip = self.load_or_def_option('random_src_flip', False) + self.options['random_src_flip'] = io.input_bool("Flip SRC faces randomly", default_random_src_flip, help_message="Random horizontal flip SRC faceset. Covers more angles, but the face may look less naturally.") + + def ask_random_dst_flip(self): + default_random_dst_flip = self.load_or_def_option('random_dst_flip', True) + self.options['random_dst_flip'] = io.input_bool("Flip DST faces randomly", default_random_dst_flip, help_message="Random horizontal flip DST faceset. Makes generalization of src->dst better, if src random flip is not enabled.") def ask_batch_size(self, suggest_batch_size=None, range=None): default_batch_size = self.load_or_def_option('batch_size', suggest_batch_size or self.batch_size) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 0ef99a6..2772870 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -65,7 +65,8 @@ class SAEHDModel(ModelBase): self.ask_autobackup_hour() self.ask_write_preview_history() self.ask_target_iter() - self.ask_random_flip() + self.ask_random_src_flip() + self.ask_random_dst_flip() self.ask_batch_size(suggest_batch_size) if self.is_first_run(): @@ -630,7 +631,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... 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), + sample_process_options=SampleProcessor.Options(random_flip=self.random_src_flip), output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': ct_mode, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': ct_mode, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, @@ -640,7 +641,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... 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), + sample_process_options=SampleProcessor.Options(random_flip=self.random_dst_flip), output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, From f387179cbaa11ef607d4b92662454c1fc0458499 Mon Sep 17 00:00:00 2001 From: iperov Date: Tue, 23 Mar 2021 16:01:24 +0400 Subject: [PATCH 3/7] Added faceset resize tool via MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 4.2) data_src util faceset resize.bat 5.2) data_dst util faceset resize.bat Resize faceset to match model resolution to reduce CPU load during training. Don’t forget to keep original faceset. --- core/imagelib/SegIEPolys.py | 8 +++++++- main.py | 10 ++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/core/imagelib/SegIEPolys.py b/core/imagelib/SegIEPolys.py index e658711..1a4c3d2 100644 --- a/core/imagelib/SegIEPolys.py +++ b/core/imagelib/SegIEPolys.py @@ -77,6 +77,8 @@ class SegIEPoly(): self.pts = np.array(pts) self.n_max = self.n = len(pts) + def mult_points(self, val): + self.pts *= val @@ -136,7 +138,11 @@ class SegIEPolys(): def dump(self): return {'polys' : [ poly.dump() for poly in self.polys ] } - + + def mult_points(self, val): + for poly in self.polys: + poly.mult_points(val) + @staticmethod def load(data=None): ie_polys = SegIEPolys() diff --git a/main.py b/main.py index cad3e92..d9fd4fd 100644 --- a/main.py +++ b/main.py @@ -255,6 +255,16 @@ if __name__ == "__main__": p.set_defaults(func=process_faceset_enhancer) + + p = facesettool_parser.add_parser ("resize", help="Resize DFL faceset.") + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.") + + def process_faceset_resizer(arguments): + osex.set_process_lowest_prio() + from mainscripts import FacesetResizer + FacesetResizer.process_folder ( Path(arguments.input_dir) ) + p.set_defaults(func=process_faceset_resizer) + def process_dev_test(arguments): osex.set_process_lowest_prio() from mainscripts import dev_misc From 8d46cd94fdcb0b962ec7cb081636163de4ece11c Mon Sep 17 00:00:00 2001 From: iperov Date: Tue, 23 Mar 2021 16:07:50 +0400 Subject: [PATCH 4/7] upd pretrain option help --- models/Model_SAEHD/Model.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/models/Model_SAEHD/Model.py b/models/Model_SAEHD/Model.py index 2772870..79329ff 100644 --- a/models/Model_SAEHD/Model.py +++ b/models/Model_SAEHD/Model.py @@ -170,7 +170,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... self.options['ct_mode'] = io.input_str (f"Color transfer for src faceset", default_ct_mode, ['none','rct','lct','mkl','idt','sot'], help_message="Change color distribution of src samples close to dst samples. Try all modes to find the best.") self.options['clipgrad'] = io.input_bool ("Enable gradient clipping", default_clipgrad, help_message="Gradient clipping reduces chance of model collapse, sacrificing speed of training.") - self.options['pretrain'] = io.input_bool ("Enable pretraining mode", default_pretrain, help_message="Pretrain the model with large amount of various faces. After that, model can be used to train the fakes more quickly.") + self.options['pretrain'] = io.input_bool ("Enable pretraining mode", default_pretrain, help_message="Pretrain the model with large amount of various faces. After that, model can be used to train the fakes more quickly. Forces random_warp=Y, random_flips=Y, gan_power=0.0, lr_dropout=N, styles=0.0, uniform_yaw=Y") if self.options['pretrain'] and self.get_pretraining_data_path() is None: raise Exception("pretraining_data_path is not defined") @@ -220,7 +220,9 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... self.gan_power = gan_power = 0.0 if self.pretrain else self.options['gan_power'] random_warp = False if self.pretrain else self.options['random_warp'] - + random_src_flip = self.random_src_flip if not self.pretrain else True + random_dst_flip = self.random_dst_flip if not self.pretrain else True + if self.pretrain: self.options_show_override['gan_power'] = 0.0 self.options_show_override['random_warp'] = False @@ -631,7 +633,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... self.set_training_data_generators ([ SampleGeneratorFace(training_data_src_path, random_ct_samples_path=random_ct_samples_path, debug=self.is_debug(), batch_size=self.get_batch_size(), - sample_process_options=SampleProcessor.Options(random_flip=self.random_src_flip), + sample_process_options=SampleProcessor.Options(random_flip=random_src_flip), output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': ct_mode, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'ct_mode': ct_mode, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, @@ -641,7 +643,7 @@ Examples: df, liae, df-d, df-ud, liae-ud, ... generators_count=src_generators_count ), SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(), - sample_process_options=SampleProcessor.Options(random_flip=self.random_dst_flip), + sample_process_options=SampleProcessor.Options(random_flip=random_dst_flip), output_sample_types = [ {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':random_warp, 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, {'sample_type': SampleProcessor.SampleType.FACE_IMAGE,'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.BGR, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, {'sample_type': SampleProcessor.SampleType.FACE_MASK, 'warp':False , 'transform':True, 'channel_type' : SampleProcessor.ChannelType.G, 'face_mask_type' : SampleProcessor.FaceMaskType.FULL_FACE, 'face_type':self.face_type, 'data_format':nn.data_format, 'resolution': resolution}, From 0ad94211017aeb754c9b971368f51307934db47a Mon Sep 17 00:00:00 2001 From: iperov Date: Tue, 23 Mar 2021 17:22:42 +0400 Subject: [PATCH 5/7] upd windows magnet link --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index 42b5383..614d798 100644 --- a/README.md +++ b/README.md @@ -194,7 +194,7 @@ Unfortunately, there is no "make everything ok" button in DeepFaceLab. You shoul -Windows (magnet link) +Windows (magnet link) Last release. Use torrent client to download. From b7245d888b4ba2b5b5d5a786403a5e9bd23b18b5 Mon Sep 17 00:00:00 2001 From: iperov Date: Wed, 24 Mar 2021 13:00:47 +0400 Subject: [PATCH 6/7] missing file --- mainscripts/FacesetResizer.py | 147 ++++++++++++++++++++++++++++++++++ 1 file changed, 147 insertions(+) create mode 100644 mainscripts/FacesetResizer.py diff --git a/mainscripts/FacesetResizer.py b/mainscripts/FacesetResizer.py new file mode 100644 index 0000000..f2ec391 --- /dev/null +++ b/mainscripts/FacesetResizer.py @@ -0,0 +1,147 @@ +import multiprocessing +import shutil + +from core import pathex +from core.cv2ex import * +from core.interact import interact as io +from core.joblib import Subprocessor +from DFLIMG import * +from facelib import LandmarksProcessor, FaceType + + +class FacesetResizerSubprocessor(Subprocessor): + + #override + def __init__(self, image_paths, output_dirpath, image_size): + self.image_paths = image_paths + self.output_dirpath = output_dirpath + self.image_size = image_size + self.result = [] + + super().__init__('FacesetResizer', FacesetResizerSubprocessor.Cli, 600) + + #override + def on_clients_initialized(self): + io.progress_bar (None, len (self.image_paths)) + + #override + def on_clients_finalized(self): + io.progress_bar_close() + + #override + def process_info_generator(self): + base_dict = {'output_dirpath':self.output_dirpath, 'image_size':self.image_size} + + for device_idx in range( min(8, multiprocessing.cpu_count()) ): + client_dict = base_dict.copy() + device_name = f'CPU #{device_idx}' + client_dict['device_name'] = device_name + yield device_name, {}, client_dict + + #override + def get_data(self, host_dict): + if len (self.image_paths) > 0: + return self.image_paths.pop(0) + + #override + def on_data_return (self, host_dict, data): + self.image_paths.insert(0, data) + + #override + def on_result (self, host_dict, data, result): + io.progress_bar_inc(1) + if result[0] == 1: + self.result +=[ (result[1], result[2]) ] + + #override + def get_result(self): + return self.result + + class Cli(Subprocessor.Cli): + + #override + def on_initialize(self, client_dict): + self.output_dirpath = client_dict['output_dirpath'] + self.image_size = client_dict['image_size'] + self.log_info (f"Running on { client_dict['device_name'] }") + + #override + def process_data(self, filepath): + try: + dflimg = DFLIMG.load (filepath) + if dflimg is None or not dflimg.has_data(): + self.log_err (f"{filepath.name} is not a dfl image file") + else: + dfl_dict = dflimg.get_dict() + + img = cv2_imread(filepath) + h,w = img.shape[:2] + if h != w: + raise Exception(f'w != h in {filepath}') + + image_size = self.image_size + scale = w / image_size + + img = cv2.resize(img, (image_size, image_size), interpolation=cv2.INTER_LANCZOS4) + + output_filepath = self.output_dirpath / filepath.name + cv2_imwrite ( str(output_filepath), img, [int(cv2.IMWRITE_JPEG_QUALITY), 100] ) + + dflimg = DFLIMG.load (output_filepath) + dflimg.set_dict(dfl_dict) + + + lmrks = dflimg.get_landmarks() + lmrks /= scale + dflimg.set_landmarks(lmrks) + + seg_ie_polys = dflimg.get_seg_ie_polys() + seg_ie_polys.mult_points( 1.0 / scale) + dflimg.set_seg_ie_polys(seg_ie_polys) + + mat = dflimg.get_image_to_face_mat() + if mat is not None: + face_type = FaceType.fromString ( dflimg.get_face_type() ) + mat = LandmarksProcessor.get_transform_mat ( dflimg.get_source_landmarks(), image_size, face_type ) + dflimg.set_image_to_face_mat(mat) + dflimg.save() + + return (1, filepath, output_filepath) + except: + self.log_err (f"Exception occured while processing file {filepath}. Error: {traceback.format_exc()}") + + return (0, filepath, None) + +def process_folder ( dirpath): + + image_size = io.input_int(f"New image size", 512, valid_range=[256,2048]) + + + output_dirpath = dirpath.parent / (dirpath.name + '_resized') + output_dirpath.mkdir (exist_ok=True, parents=True) + + dirpath_parts = '/'.join( dirpath.parts[-2:]) + output_dirpath_parts = '/'.join( output_dirpath.parts[-2:] ) + io.log_info (f"Resizing faceset in {dirpath_parts}") + io.log_info ( f"Processing to {output_dirpath_parts}") + + output_images_paths = pathex.get_image_paths(output_dirpath) + if len(output_images_paths) > 0: + for filename in output_images_paths: + Path(filename).unlink() + + image_paths = [Path(x) for x in pathex.get_image_paths( dirpath )] + result = FacesetResizerSubprocessor ( image_paths, output_dirpath, image_size).run() + + is_merge = io.input_bool (f"\r\nMerge {output_dirpath_parts} to {dirpath_parts} ?", True) + if is_merge: + io.log_info (f"Copying processed files to {dirpath_parts}") + + for (filepath, output_filepath) in result: + try: + shutil.copy (output_filepath, filepath) + except: + pass + + io.log_info (f"Removing {output_dirpath_parts}") + shutil.rmtree(output_dirpath) From 1652bffeb0bb242dc06516880f754abfea1e207e Mon Sep 17 00:00:00 2001 From: iperov Date: Wed, 24 Mar 2021 13:42:57 +0400 Subject: [PATCH 7/7] change in xseg apply and fetch --- mainscripts/XSegUtil.py | 65 ++++++++++++++++++++++++++++++++--------- 1 file changed, 51 insertions(+), 14 deletions(-) diff --git a/mainscripts/XSegUtil.py b/mainscripts/XSegUtil.py index 9f9f893..b95bb2c 100644 --- a/mainscripts/XSegUtil.py +++ b/mainscripts/XSegUtil.py @@ -10,7 +10,7 @@ from core.cv2ex import * from core.interact import interact as io from core.leras import nn from DFLIMG import * -from facelib import XSegNet +from facelib import XSegNet, LandmarksProcessor, FaceType def apply_xseg(input_path, model_path): @@ -19,7 +19,16 @@ def apply_xseg(input_path, model_path): if not model_path.exists(): raise ValueError(f'{model_path} not found. Please ensure it exists.') - + + face_type = io.input_str ("XSeg model face type", 'same', ['h','mf','f','wf','head','same'], help_message="Specify face type of trained XSeg model. For example if XSeg model trained as WF, but faceset is HEAD, specify WF to apply xseg only on WF part of HEAD. Default is 'same'").lower() + if face_type == 'same': + face_type = None + else: + face_type = {'h' : FaceType.HALF, + 'mf' : FaceType.MID_FULL, + 'f' : FaceType.FULL, + 'wf' : FaceType.WHOLE_FACE, + 'head' : FaceType.HEAD}[face_type] io.log_info(f'Applying trained XSeg model to {input_path.name}/ folder.') device_config = nn.DeviceConfig.ask_choose_device(choose_only_one=True) @@ -30,7 +39,7 @@ def apply_xseg(input_path, model_path): weights_file_root=model_path, data_format=nn.data_format, raise_on_no_model_files=True) - res = xseg.get_resolution() + xseg_res = xseg.get_resolution() images_paths = pathex.get_image_paths(input_path, return_Path_class=True) @@ -42,15 +51,36 @@ def apply_xseg(input_path, model_path): img = cv2_imread(filepath).astype(np.float32) / 255.0 h,w,c = img.shape - if w != res: - img = cv2.resize( img, (res,res), interpolation=cv2.INTER_CUBIC ) - if len(img.shape) == 2: - img = img[...,None] - - mask = xseg.extract(img) - mask[mask < 0.5]=0 - mask[mask >= 0.5]=1 + img_face_type = FaceType.fromString( dflimg.get_face_type() ) + if face_type is not None and img_face_type != face_type: + lmrks = dflimg.get_source_landmarks() + + fmat = LandmarksProcessor.get_transform_mat(lmrks, w, face_type) + imat = LandmarksProcessor.get_transform_mat(lmrks, w, img_face_type) + + g_p = LandmarksProcessor.transform_points (np.float32([(0,0),(w,0),(0,w) ]), fmat, True) + g_p2 = LandmarksProcessor.transform_points (g_p, imat) + + mat = cv2.getAffineTransform( g_p2, np.float32([(0,0),(w,0),(0,w) ]) ) + + img = cv2.warpAffine(img, mat, (w, w), cv2.INTER_LANCZOS4) + img = cv2.resize(img, (xseg_res, xseg_res), interpolation=cv2.INTER_LANCZOS4) + else: + if w != xseg_res: + img = cv2.resize( img, (xseg_res,xseg_res), interpolation=cv2.INTER_LANCZOS4 ) + + if len(img.shape) == 2: + img = img[...,None] + + mask = xseg.extract(img) + + if face_type is not None and img_face_type != face_type: + mask = cv2.resize(mask, (w, w), interpolation=cv2.INTER_LANCZOS4) + mask = cv2.warpAffine( mask, mat, (w,w), np.zeros( (h,w,c), dtype=np.float), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4) + mask = cv2.resize(mask, (xseg_res, xseg_res), interpolation=cv2.INTER_LANCZOS4) + mask[mask < 0.5]=0 + mask[mask >= 0.5]=1 dflimg.set_xseg_mask(mask) dflimg.save() @@ -67,7 +97,8 @@ def fetch_xseg(input_path): images_paths = pathex.get_image_paths(input_path, return_Path_class=True) - files_copied = 0 + + files_copied = [] for filepath in io.progress_bar_generator(images_paths, "Processing"): dflimg = DFLIMG.load(filepath) if dflimg is None or not dflimg.has_data(): @@ -77,10 +108,16 @@ def fetch_xseg(input_path): ie_polys = dflimg.get_seg_ie_polys() if ie_polys.has_polys(): - files_copied += 1 + files_copied.append(filepath) shutil.copy ( str(filepath), str(output_path / filepath.name) ) - io.log_info(f'Files copied: {files_copied}') + io.log_info(f'Files copied: {len(files_copied)}') + + is_delete = io.input_bool (f"\r\nDelete original files?", True) + if is_delete: + for filepath in files_copied: + Path(filepath).unlink() + def remove_xseg(input_path): if not input_path.exists():