diff --git a/converters/Converter.py b/converters/Converter.py index 38ac748..93db97e 100644 --- a/converters/Converter.py +++ b/converters/Converter.py @@ -19,7 +19,11 @@ class Converter(object): pass #overridable - def convert_face (self, img_bgr, img_face_landmarks, debug): + def on_host_(self): + pass + + #overridable + def cli_convert_face (self, img_bgr, img_face_landmarks, debug): #return float32 image #if debug , return tuple ( images of any size and channels, ...) return image diff --git a/converters/ConverterMasked.py b/converters/ConverterMasked.py index a5fb7eb..ca899ab 100644 --- a/converters/ConverterMasked.py +++ b/converters/ConverterMasked.py @@ -7,6 +7,9 @@ import cv2 import numpy as np import imagelib from interact import interact as io +from joblib import SubprocessFunctionCaller +from utils.pickle_utils import AntiPickler + ''' default_mode = {1:'overlay', 2:'hist-match', @@ -20,7 +23,7 @@ class ConverterMasked(Converter): #override def __init__(self, predictor_func, predictor_input_size=0, - output_size=0, + predictor_masked=True, face_type=FaceType.FULL, default_mode = 4, base_erode_mask_modifier = 0, @@ -30,8 +33,12 @@ class ConverterMasked(Converter): clip_hborder_mask_per = 0): super().__init__(predictor_func, Converter.TYPE_FACE) + + predictor_func_host, predictor_func = SubprocessFunctionCaller.make_pair(predictor_func) + self.predictor_func_host = AntiPickler(predictor_func_host) + self.predictor_func = predictor_func + self.predictor_masked = predictor_masked self.predictor_input_size = predictor_input_size - self.output_size = output_size self.face_type = face_type self.clip_hborder_mask_per = clip_hborder_mask_per @@ -85,6 +92,7 @@ class ConverterMasked(Converter): self.output_face_scale = np.clip ( 1.0 + io.input_int ("Choose output face scale modifier [-50..50] (skip:0) : ", 0)*0.01, 0.5, 1.5) self.color_transfer_mode = io.input_str ("Apply color transfer to predicted face? Choose mode ( rct/lct skip:None ) : ", None, ['rct','lct']) + self.super_resolution = io.input_bool("Apply super resolution? (y/n skip:n) : ", False, help_message="Enhance details by applying DCSCN network.") if self.mode != 'raw': self.final_image_color_degrade_power = np.clip ( io.input_int ("Degrade color power of final image [0..100] (skip:0) : ", 0), 0, 100) @@ -93,9 +101,19 @@ class ConverterMasked(Converter): io.log_info ("") self.over_res = 4 if self.suppress_seamless_jitter else 1 - #override - def dummy_predict(self): - self.predictor_func ( np.zeros ( (self.predictor_input_size,self.predictor_input_size,4), dtype=np.float32 ) ) + if self.super_resolution: + host_proc, dc_upscale = SubprocessFunctionCaller.make_pair( imagelib.DCSCN().upscale ) + self.dc_host = AntiPickler(host_proc) + self.dc_upscale = dc_upscale + else: + self.dc_host = None + + #overridable + def on_host_tick(self): + self.predictor_func_host.obj.process_messages() + + if self.dc_host is not None: + self.dc_host.obj.process_messages() #overridable def on_cli_initialize(self): @@ -103,7 +121,7 @@ class ConverterMasked(Converter): self.fan_seg = FANSegmentator(256, FaceType.toString(FaceType.FULL) ) #override - def convert_face (self, img_bgr, img_face_landmarks, debug): + def cli_convert_face (self, img_bgr, img_face_landmarks, debug): if self.over_res != 1: img_bgr = cv2.resize ( img_bgr, ( img_bgr.shape[1]*self.over_res, img_bgr.shape[0]*self.over_res ) ) img_face_landmarks = img_face_landmarks*self.over_res @@ -115,32 +133,52 @@ class ConverterMasked(Converter): img_face_mask_a = LandmarksProcessor.get_image_hull_mask (img_bgr.shape, img_face_landmarks) - face_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, self.output_size, face_type=self.face_type) - face_output_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, self.output_size, face_type=self.face_type, scale=self.output_face_scale) + output_size = self.predictor_input_size + if self.super_resolution: + output_size *= 2 - dst_face_bgr = cv2.warpAffine( img_bgr , face_mat, (self.output_size, self.output_size), flags=cv2.INTER_LANCZOS4 ) - dst_face_mask_a_0 = cv2.warpAffine( img_face_mask_a, face_mat, (self.output_size, self.output_size), flags=cv2.INTER_LANCZOS4 ) + face_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, output_size, face_type=self.face_type) + face_output_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, output_size, face_type=self.face_type, scale=self.output_face_scale) + + dst_face_bgr = cv2.warpAffine( img_bgr , face_mat, (output_size, output_size), flags=cv2.INTER_LANCZOS4 ) + dst_face_mask_a_0 = cv2.warpAffine( img_face_mask_a, face_mat, (output_size, output_size), flags=cv2.INTER_LANCZOS4 ) predictor_input_bgr = cv2.resize (dst_face_bgr, (self.predictor_input_size,self.predictor_input_size)) - predictor_input_mask_a_0 = cv2.resize (dst_face_mask_a_0, (self.predictor_input_size,self.predictor_input_size)) - predictor_input_mask_a = np.expand_dims (predictor_input_mask_a_0, -1) - predicted_bgra = self.predictor_func ( np.concatenate( (predictor_input_bgr, predictor_input_mask_a), -1) ) + if self.predictor_masked: + prd_face_bgr, prd_face_mask_a_0 = self.predictor_func (predictor_input_bgr) + prd_face_bgr = np.clip (prd_face_bgr, 0, 1.0 ) + prd_face_mask_a_0 = np.clip (prd_face_mask_a_0, 0.0, 1.0) + else: + predicted = self.predictor_func (predictor_input_bgr) + prd_face_bgr = np.clip (predicted, 0, 1.0 ) + prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (self.predictor_input_size,self.predictor_input_size)) - prd_face_bgr = np.clip (predicted_bgra[:,:,0:3], 0, 1.0 ) - prd_face_mask_a_0 = np.clip (predicted_bgra[:,:,3], 0.0, 1.0) + if self.super_resolution: + if debug: + tmp = cv2.resize (prd_face_bgr, (output_size,output_size), cv2.INTER_CUBIC) + debugs += [ np.clip( cv2.warpAffine( tmp, face_output_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + + prd_face_bgr = self.dc_upscale(prd_face_bgr) + if debug: + debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + + if self.predictor_masked: + prd_face_mask_a_0 = cv2.resize (prd_face_mask_a_0, (output_size, output_size), cv2.INTER_CUBIC) + else: + prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (output_size, output_size), cv2.INTER_CUBIC) if self.mask_mode == 2: #dst - prd_face_mask_a_0 = predictor_input_mask_a_0 + prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC) elif self.mask_mode == 3: #FAN-prd prd_face_bgr_256 = cv2.resize (prd_face_bgr, (256,256) ) - prd_face_bgr_256_mask = self.fan_seg.extract_from_bgr( np.expand_dims(prd_face_bgr_256,0) ) [0] - prd_face_mask_a_0 = cv2.resize (prd_face_bgr_256_mask, (self.predictor_input_size, self.predictor_input_size)) + prd_face_bgr_256_mask = self.fan_seg.extract_from_bgr( prd_face_bgr_256[np.newaxis,...] ) [0] + prd_face_mask_a_0 = cv2.resize (prd_face_bgr_256_mask, (output_size,output_size), cv2.INTER_CUBIC) elif self.mask_mode == 4: #FAN-dst face_256_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, 256, face_type=FaceType.FULL) dst_face_256_bgr = cv2.warpAffine(img_bgr, face_256_mat, (256, 256), flags=cv2.INTER_LANCZOS4 ) - dst_face_256_mask = self.fan_seg.extract_from_bgr( np.expand_dims(dst_face_256_bgr,0) ) [0] - prd_face_mask_a_0 = cv2.resize (dst_face_256_mask, (self.predictor_input_size, self.predictor_input_size)) + dst_face_256_mask = self.fan_seg.extract_from_bgr( dst_face_256_bgr[np.newaxis,...] ) [0] + prd_face_mask_a_0 = cv2.resize (dst_face_256_mask, (output_size,output_size), cv2.INTER_CUBIC) prd_face_mask_a_0[ prd_face_mask_a_0 < 0.001 ] = 0.0 @@ -333,7 +371,7 @@ class ConverterMasked(Converter): out_img = np.clip( img_bgr*(1-img_mask_blurry_aaa) + (out_img*img_mask_blurry_aaa) , 0, 1.0 ) if self.mode == 'seamless-hist-match': - out_face_bgr = cv2.warpAffine( out_img, face_mat, (self.output_size, self.output_size) ) + out_face_bgr = cv2.warpAffine( out_img, face_mat, (output_size, output_size) ) new_out_face_bgr = imagelib.color_hist_match(out_face_bgr, dst_face_bgr, self.hist_match_threshold) new_out = cv2.warpAffine( new_out_face_bgr, face_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) out_img = np.clip( img_bgr*(1-img_mask_blurry_aaa) + (new_out*img_mask_blurry_aaa) , 0, 1.0 ) diff --git a/imagelib/DCSCN.h5 b/imagelib/DCSCN.h5 new file mode 100644 index 0000000..d9c8ac7 Binary files /dev/null and b/imagelib/DCSCN.h5 differ diff --git a/imagelib/DCSCN.py b/imagelib/DCSCN.py new file mode 100644 index 0000000..332e70a --- /dev/null +++ b/imagelib/DCSCN.py @@ -0,0 +1,164 @@ +import numpy as np +import cv2 +from pathlib import Path +from nnlib import nnlib +from interact import interact as io + +class DCSCN(): + def __init__(self): + exec( nnlib.import_all(), locals(), globals() ) + + inp_x = KL.Input([None, None, 1]) + inp_x2 = KL.Input([None, None, 1]) + + x = inp_x + layers_count = 12 + layers = [] + for i in range(1,layers_count+1): + if i == 1: + output_feature_num = 196 + else: + x1 = (i-1) / float(layers_count - 1) + y1 = x1 ** (1.0 / 1.5) + output_feature_num = int((196 - 48) * (1 - y1) + 48) + x = Conv2D(output_feature_num, kernel_size=3, strides=1, padding='same', name='CNN%d' % (i) ) (x) + x = PReLU(shared_axes=[1,2], name='CNN%d_prelu' % (i) ) (x) + layers.append(x) + + x_concat = KL.Concatenate()(layers) + + A1 = Conv2D(64, kernel_size=1, strides=1, padding='same', name='A1' ) (x_concat) + A1 = PReLU(shared_axes=[1,2], name='A1_prelu') (A1) + + B1 = Conv2D(32, kernel_size=1, strides=1, padding='same', name='B1' ) (x_concat) + B1 = PReLU(shared_axes=[1,2], name='B1_prelu') (B1) + + B2 = Conv2D(32, kernel_size=3, strides=1, padding='same', name='B2' ) (B1) + B2 = PReLU(shared_axes=[1,2], name='B2_prelu') (B2) + + x = KL.Concatenate()([B2,A1]) + x = Conv2D(96*4, kernel_size=3, strides=1, padding='same', name='Up_PS' )(x) + x = PixelShuffler()(x) + x = Conv2D(1, kernel_size=3, strides=1, padding='same', name='R_CNN1', use_bias=False )(x) + x = KL.Add()([x, inp_x2]) + self.model = keras.models.Model ([inp_x, inp_x2], [x]) + self.model.load_weights ( Path(__file__).parent / 'DCSCN.h5' ) + + def upscale(self, img, is_bgr=True, is_float=True): + if is_bgr: + img = img[...,::-1] + + if is_float: + img = np.clip (img*255, 0, 255) + + img_shape_len = len(img.shape) + h, w = img.shape[:2] + ch = img.shape[2] if len(img.shape) >= 3 else 1 + + nh, nw = h*2, w*2 + + img_x = self.convert_rgb_to_y(img) + + img_bx = cv2.resize(img_x, (nh, nw), cv2.INTER_CUBIC) + + ensemble = 8 + + output = np.zeros([nh,nw,1], dtype=np.float32) + + for i in range(ensemble): + x = np.reshape( self.flip(img_x, i), (1,h,w,1) ) + bx = np.reshape( self.flip(img_bx, i), (1,nh,nw,1) ) + y = self.model.predict([x,bx])[0] + y = self.flip(y, i, invert=True) + output += y + + output /= ensemble + + bimg = cv2.resize(img, (nh, nw), cv2.INTER_CUBIC) + bimg_ycbcr = self.convert_rgb_to_ycbcr(bimg) + + if ch > 1: + output = self.convert_y_and_cbcr_to_rgb(output, bimg_ycbcr[:, :, 1:3]) + + if is_float: + output = np.clip (output/255.0, 0, 1.0) + + if is_bgr: + output = output[...,::-1] + + return output + + def convert_rgb_to_y(self, image): + if len(image.shape) <= 2 or image.shape[2] == 1: + return image + + xform = np.array([[65.738 / 256.0, 129.057 / 256.0, 25.064 / 256.0]], dtype=np.float32) + y_image = image.dot(xform.T) + 16.0 + + return y_image + + + def convert_rgb_to_ycbcr(self, image): + if len(image.shape) <= 2 or image.shape[2] == 1: + return image + + xform = np.array( + [[65.738 / 256.0, 129.057 / 256.0, 25.064 / 256.0], + [- 37.945 / 256.0, - 74.494 / 256.0, 112.439 / 256.0], + [112.439 / 256.0, - 94.154 / 256.0, - 18.285 / 256.0]], dtype=np.float32) + + ycbcr_image = image.dot(xform.T) + ycbcr_image[:, :, 0] += 16.0 + ycbcr_image[:, :, [1, 2]] += 128.0 + + return ycbcr_image + + def convert_ycbcr_to_rgb(self,ycbcr_image): + rgb_image = np.zeros([ycbcr_image.shape[0], ycbcr_image.shape[1], 3], dtype=np.float32) + + rgb_image[:, :, 0] = ycbcr_image[:, :, 0] - 16.0 + rgb_image[:, :, [1, 2]] = ycbcr_image[:, :, [1, 2]] - 128.0 + xform = np.array( + [[298.082 / 256.0, 0, 408.583 / 256.0], + [298.082 / 256.0, -100.291 / 256.0, -208.120 / 256.0], + [298.082 / 256.0, 516.412 / 256.0, 0]], dtype=np.float32) + rgb_image = rgb_image.dot(xform.T) + + return rgb_image + + def convert_y_and_cbcr_to_rgb(self,y_image, cbcr_image): + if len(y_image.shape) <= 2: + y_image = y_image.reshape[y_image.shape[0], y_image.shape[1], 1] + + if len(y_image.shape) == 3 and y_image.shape[2] == 3: + y_image = y_image[:, :, 0:1] + + ycbcr_image = np.zeros([y_image.shape[0], y_image.shape[1], 3], dtype=np.float32) + ycbcr_image[:, :, 0] = y_image[:, :, 0] + ycbcr_image[:, :, 1:3] = cbcr_image[:, :, 0:2] + + return self.convert_ycbcr_to_rgb(ycbcr_image) + + def flip(self, image, flip_type, invert=False): + if flip_type == 0: + return image + elif flip_type == 1: + return np.flipud(image) + elif flip_type == 2: + return np.fliplr(image) + elif flip_type == 3: + return np.flipud(np.fliplr(image)) + elif flip_type == 4: + return np.rot90(image, 1 if invert is False else -1) + elif flip_type == 5: + return np.rot90(image, -1 if invert is False else 1) + elif flip_type == 6: + if invert is False: + return np.flipud(np.rot90(image)) + else: + return np.rot90(np.flipud(image), -1) + elif flip_type == 7: + if invert is False: + return np.flipud(np.rot90(image, -1)) + else: + return np.rot90(np.flipud(image), 1) diff --git a/imagelib/__init__.py b/imagelib/__init__.py index 2f674e1..64e5a90 100644 --- a/imagelib/__init__.py +++ b/imagelib/__init__.py @@ -17,3 +17,5 @@ from .reduce_colors import reduce_colors from .color_transfer import color_hist_match from .color_transfer import reinhard_color_transfer from .color_transfer import linear_color_transfer + +from .DCSCN import DCSCN \ No newline at end of file diff --git a/mainscripts/Converter.py b/mainscripts/Converter.py index 7d035db..6e12bcf 100644 --- a/mainscripts/Converter.py +++ b/mainscripts/Converter.py @@ -108,9 +108,9 @@ class ConvertSubprocessor(Subprocessor): self.log_info ( '\nConverting face_num [%d] in file [%s]' % (face_num, filename_path) ) if self.debug: - debug_images += self.converter.convert_face(image, image_landmarks, self.debug) + debug_images += self.converter.cli_convert_face(image, image_landmarks, self.debug) else: - image = self.converter.convert_face(image, image_landmarks, self.debug) + image = self.converter.cli_convert_face(image, image_landmarks, self.debug) except Exception as e: e_str = traceback.format_exc() @@ -140,9 +140,6 @@ class ConvertSubprocessor(Subprocessor): super().__init__('Converter', ConvertSubprocessor.Cli, 86400 if debug == True else 60) self.converter = converter - self.host_processor, self.cli_func = SubprocessFunctionCaller.make_pair ( self.converter.predictor_func ) - self.process_converter = self.converter.copy_and_set_predictor(self.cli_func) - self.input_data = self.input_path_image_paths = input_path_image_paths self.output_path = output_path self.alignments = alignments @@ -158,7 +155,7 @@ class ConvertSubprocessor(Subprocessor): for i in r: yield 'CPU%d' % (i), {}, {'device_idx': i, 'device_name': 'CPU%d' % (i), - 'converter' : self.process_converter, + 'converter' : self.converter, 'output_dir' : str(self.output_path), 'alignments' : self.alignments, 'debug': self.debug, @@ -202,7 +199,7 @@ class ConvertSubprocessor(Subprocessor): #override def on_tick(self): - self.host_processor.process_messages() + self.converter.on_host_tick() #override def get_result(self): @@ -235,7 +232,6 @@ def main (args, device_args): import models model = models.import_model( args['model_name'] )(model_path, device_args=device_args) converter = model.get_converter() - converter.dummy_predict() alignments = None diff --git a/models/Model_DF/Model.py b/models/Model_DF/Model.py index 5cdd723..e6f2cc4 100644 --- a/models/Model_DF/Model.py +++ b/models/Model_DF/Model.py @@ -33,12 +33,16 @@ class Model(ModelBase): ] self.load_weights_safe(weights_to_load) - self.autoencoder_src = Model([ae_input_layer,mask_layer], self.decoder_src(self.encoder(ae_input_layer))) - self.autoencoder_dst = Model([ae_input_layer,mask_layer], self.decoder_dst(self.encoder(ae_input_layer))) + rec_src = self.decoder_src(self.encoder(ae_input_layer)) + rec_dst = self.decoder_dst(self.encoder(ae_input_layer)) + self.autoencoder_src = Model([ae_input_layer,mask_layer], rec_src) + self.autoencoder_dst = Model([ae_input_layer,mask_layer], rec_dst) self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] ) self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] ) + self.convert = K.function([ae_input_layer], rec_src) + if self.is_training_mode: f = SampleProcessor.TypeFlags self.set_training_data_generators ([ @@ -103,21 +107,14 @@ class Model(ModelBase): return [ ('DF', np.concatenate ( st, axis=0 ) ) ] def predictor_func (self, face): - - face_128_bgr = face[...,0:3] - face_128_mask = np.expand_dims(face[...,3],-1) - - x, mx = self.autoencoder_src.predict ( [ np.expand_dims(face_128_bgr,0), np.expand_dims(face_128_mask,0) ] ) - x, mx = x[0], mx[0] - - return np.concatenate ( (x,mx), -1 ) + x, mx = self.convert ( [ face[np.newaxis,...] ] ) + return x[0], mx[0][...,0] #override def get_converter(self): from converters import ConverterMasked return ConverterMasked(self.predictor_func, predictor_input_size=128, - output_size=128, face_type=FaceType.FULL, base_erode_mask_modifier=30, base_blur_mask_modifier=0) diff --git a/models/Model_H128/Model.py b/models/Model_H128/Model.py index 78d810c..9a9c444 100644 --- a/models/Model_H128/Model.py +++ b/models/Model_H128/Model.py @@ -116,20 +116,14 @@ class Model(ModelBase): return [ ('H128', np.concatenate ( st, axis=0 ) ) ] def predictor_func (self, face): - face_128_bgr = face[...,0:3] - face_128_mask = np.expand_dims(face[...,3],-1) - - x, mx = self.src_view ( [ np.expand_dims(face_128_bgr,0) ] ) - x, mx = x[0], mx[0] - - return np.concatenate ( (x,mx), -1 ) + x, mx = self.src_view ( [ face[np.newaxis,...] ] ) + return x[0], mx[0][...,0] #override def get_converter(self): from converters import ConverterMasked return ConverterMasked(self.predictor_func, predictor_input_size=128, - output_size=128, face_type=FaceType.HALF, base_erode_mask_modifier=100, base_blur_mask_modifier=100) diff --git a/models/Model_H64/Model.py b/models/Model_H64/Model.py index 337ade1..2ed75ff 100644 --- a/models/Model_H64/Model.py +++ b/models/Model_H64/Model.py @@ -117,21 +117,14 @@ class Model(ModelBase): return [ ('H64', np.concatenate ( st, axis=0 ) ) ] def predictor_func (self, face): - - face_64_bgr = face[...,0:3] - face_64_mask = np.expand_dims(face[...,3],-1) - - x, mx = self.src_view ( [ np.expand_dims(face_64_bgr,0) ] ) - x, mx = x[0], mx[0] - - return np.concatenate ( (x,mx), -1 ) + x, mx = self.src_view ( [ face[np.newaxis,...] ] ) + return x[0], mx[0][...,0] #override def get_converter(self): from converters import ConverterMasked return ConverterMasked(self.predictor_func, predictor_input_size=64, - output_size=64, face_type=FaceType.HALF, base_erode_mask_modifier=100, base_blur_mask_modifier=100) diff --git a/models/Model_LIAEF128/Model.py b/models/Model_LIAEF128/Model.py index 66b2754..53119a0 100644 --- a/models/Model_LIAEF128/Model.py +++ b/models/Model_LIAEF128/Model.py @@ -37,12 +37,17 @@ class Model(ModelBase): code = self.encoder(ae_input_layer) AB = self.inter_AB(code) B = self.inter_B(code) - self.autoencoder_src = Model([ae_input_layer,mask_layer], self.decoder(Concatenate()([AB, AB])) ) - self.autoencoder_dst = Model([ae_input_layer,mask_layer], self.decoder(Concatenate()([B, AB])) ) + rec_src = self.decoder(Concatenate()([AB, AB])) + rec_dst = self.decoder(Concatenate()([B, AB])) + self.autoencoder_src = Model([ae_input_layer,mask_layer], rec_src ) + self.autoencoder_dst = Model([ae_input_layer,mask_layer], rec_dst ) self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] ) self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMSEMaskLoss(mask_layer, is_mse=self.options['pixel_loss']), 'mse'] ) + self.convert = K.function([ae_input_layer],rec_src) + + if self.is_training_mode: f = SampleProcessor.TypeFlags self.set_training_data_generators ([ @@ -111,21 +116,14 @@ class Model(ModelBase): return [ ('LIAEF128', np.concatenate ( st, axis=0 ) ) ] def predictor_func (self, face): - - face_128_bgr = face[...,0:3] - face_128_mask = np.expand_dims(face[...,3],-1) - - x, mx = self.autoencoder_src.predict ( [ np.expand_dims(face_128_bgr,0), np.expand_dims(face_128_mask,0) ] ) - x, mx = x[0], mx[0] - - return np.concatenate ( (x,mx), -1 ) + x, mx = self.convert ( [ face[np.newaxis,...] ] ) + return x[0], mx[0][...,0] #override def get_converter(self): from converters import ConverterMasked return ConverterMasked(self.predictor_func, predictor_input_size=128, - output_size=128, face_type=FaceType.FULL, base_erode_mask_modifier=30, base_blur_mask_modifier=0) diff --git a/models/Model_SAE/Model.py b/models/Model_SAE/Model.py index 9ce0294..9647ba1 100644 --- a/models/Model_SAE/Model.py +++ b/models/Model_SAE/Model.py @@ -406,13 +406,12 @@ class SAEModel(ModelBase): return [ ('SAE', np.concatenate (st, axis=0 )), ] def predictor_func (self, face): - - prd = [ x[0] for x in self.AE_convert ( [ face[np.newaxis,:,:,0:3] ] ) ] - - if not self.options['learn_mask']: - prd += [ face[...,3:4] ] - - return np.concatenate ( prd, -1 ) + if self.options['learn_mask']: + bgr, mask = self.AE_convert ([face[np.newaxis,...]]) + return bgr[0], mask[0][...,0] + else: + bgr, = self.AE_convert ([face[np.newaxis,...]]) + return bgr[0] #override def get_converter(self): @@ -428,7 +427,7 @@ class SAEModel(ModelBase): from converters import ConverterMasked return ConverterMasked(self.predictor_func, predictor_input_size=self.options['resolution'], - output_size=self.options['resolution'], + predictor_masked=self.options['learn_mask'], face_type=face_type, default_mode = 1 if self.options['face_style_power'] or self.options['bg_style_power'] else 4, base_erode_mask_modifier=base_erode_mask_modifier, diff --git a/utils/pickle_utils.py b/utils/pickle_utils.py new file mode 100644 index 0000000..37c4c72 --- /dev/null +++ b/utils/pickle_utils.py @@ -0,0 +1,9 @@ +class AntiPickler(): + def __init__(self, obj): + self.obj = obj + + def __getstate__(self): + return dict() + + def __setstate__(self, d): + self.__dict__.update(d) \ No newline at end of file