|
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 2ba085f..d9fd4fd 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='+')
@@ -253,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
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)
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/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():
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 cfa3d75..6df40c1 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():
@@ -169,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")
@@ -204,6 +205,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']
@@ -217,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
@@ -236,22 +241,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 +614,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()
@@ -625,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_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},
@@ -635,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_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},
@@ -650,7 +658,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
|