Converter: added Apply super resolution? (y/n skip:n) : , Enhance details by applying DCSCN network.

refactorings

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

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
+        if self.super_resolution:
-    def dummy_predict(self):
+            host_proc, dc_upscale = SubprocessFunctionCaller.make_pair( imagelib.DCSCN().upscale )
-        self.predictor_func ( np.zeros ( (self.predictor_input_size,self.predictor_input_size,4), dtype=np.float32 ) )
+            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)
+        output_size = self.predictor_input_size
-        face_output_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, self.output_size, face_type=self.face_type, scale=self.output_face_scale)
+        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 )
+        face_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, output_size, face_type=self.face_type)
-        dst_face_mask_a_0 = cv2.warpAffine( img_face_mask_a, face_mat, (self.output_size, self.output_size), flags=cv2.INTER_LANCZOS4 )
+        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 )
+        if self.super_resolution:
-        prd_face_mask_a_0 = np.clip (predicted_bgra[:,:,3], 0.0, 1.0)
+            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_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, (self.predictor_input_size, self.predictor_input_size))
+            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]
+            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, (self.predictor_input_size, self.predictor_input_size))
+            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 )

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)

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

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
 

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)))
+        rec_src = 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_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):
-
+        x, mx = self.convert ( [ face[np.newaxis,...] ] )
-        face_128_bgr = face[...,0:3]
+        return x[0], mx[0][...,0]
-        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 )
 
     #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)

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]
+        x, mx = self.src_view ( [ face[np.newaxis,...] ] )
-        face_128_mask = np.expand_dims(face[...,3],-1)
+        return x[0], mx[0][...,0]
-
-        x, mx = self.src_view ( [ np.expand_dims(face_128_bgr,0) ] )
-        x, mx = x[0], mx[0]
-
-        return np.concatenate ( (x,mx), -1 )
 
     #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)

models/Model_H64/Model.py

@@ -117,21 +117,14 @@ class Model(ModelBase):
         return [ ('H64', np.concatenate ( st, axis=0 ) ) ]
 
     def predictor_func (self, face):
-
+        x, mx = self.src_view ( [ face[np.newaxis,...] ] )
-        face_64_bgr = face[...,0:3]
+        return x[0], mx[0][...,0]
-        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 )
 
     #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)

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])) )
+        rec_src = self.decoder(Concatenate()([AB, AB]))
-        self.autoencoder_dst = Model([ae_input_layer,mask_layer], self.decoder(Concatenate()([B, 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):
-
+        x, mx = self.convert ( [ face[np.newaxis,...] ] )
-        face_128_bgr = face[...,0:3]
+        return x[0], mx[0][...,0]
-        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 )
 
     #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)

models/Model_SAE/Model.py

@@ -406,13 +406,12 @@ class SAEModel(ModelBase):
         return [ ('SAE', np.concatenate (st, axis=0 )), ]
 
     def predictor_func (self, face):
-
+        if self.options['learn_mask']:
-        prd = [ x[0] for x in self.AE_convert ( [ face[np.newaxis,:,:,0:3] ] ) ]
+            bgr, mask = self.AE_convert ([face[np.newaxis,...]])
-
+            return bgr[0], mask[0][...,0]
-        if not self.options['learn_mask']:
+        else:
-            prd += [ face[...,3:4] ]
+            bgr, = self.AE_convert ([face[np.newaxis,...]])
-
+            return bgr[0]
-        return np.concatenate ( prd, -1 )
 
     #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,

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)