diff --git a/converters/Converter.py b/converters/Converter.py index ef0b819..b4e4213 100644 --- a/converters/Converter.py +++ b/converters/Converter.py @@ -5,8 +5,9 @@ You can implement your own Converter, check example ConverterMasked.py class Converter(object): TYPE_FACE = 0 #calls convert_face - TYPE_IMAGE = 1 #calls convert_image without landmarks - TYPE_IMAGE_WITH_LANDMARKS = 2 #calls convert_image with landmarks + TYPE_FACE_AVATAR = 1 #calls convert_face with avatar_operator_face + TYPE_IMAGE = 2 #calls convert_image without landmarks + TYPE_IMAGE_WITH_LANDMARKS = 3 #calls convert_image with landmarks #overridable def __init__(self, predictor_func, type): @@ -23,13 +24,13 @@ class Converter(object): pass #overridable - def cli_convert_face (self, img_bgr, img_face_landmarks, debug): + def cli_convert_face (self, img_bgr, img_face_landmarks, debug, avaperator_face_bgr=None, **kwargs): #return float32 image #if debug , return tuple ( images of any size and channels, ...) return image - + #overridable - def convert_image (self, img_bgr, img_landmarks, debug): + def cli_convert_image (self, img_bgr, img_landmarks, debug): #img_landmarks not None, if input image is png with embedded data #return float32 image #if debug , return tuple ( images of any size and channels, ...) diff --git a/converters/ConverterAvatar.py b/converters/ConverterAvatar.py new file mode 100644 index 0000000..4a9e9b5 --- /dev/null +++ b/converters/ConverterAvatar.py @@ -0,0 +1,70 @@ +import time + +import cv2 +import numpy as np + +from facelib import FaceType, LandmarksProcessor +from joblib import SubprocessFunctionCaller +from utils.pickle_utils import AntiPickler + +from .Converter import Converter + +class ConverterAvatar(Converter): + + #override + def __init__(self, predictor_func, + predictor_input_size=0): + + super().__init__(predictor_func, Converter.TYPE_FACE_AVATAR) + + self.predictor_input_size = predictor_input_size + + #dummy predict and sleep, tensorflow caching kernels. If remove it, conversion speed will be x2 slower + predictor_func ( np.zeros ( (predictor_input_size,predictor_input_size,3), dtype=np.float32 ), + np.zeros ( (predictor_input_size,predictor_input_size,1), dtype=np.float32 ) ) + time.sleep(2) + + predictor_func_host, predictor_func = SubprocessFunctionCaller.make_pair(predictor_func) + self.predictor_func_host = AntiPickler(predictor_func_host) + self.predictor_func = predictor_func + + #overridable + def on_host_tick(self): + self.predictor_func_host.obj.process_messages() + + #override + def cli_convert_face (self, img_bgr, img_face_landmarks, debug, avaperator_face_bgr=None, **kwargs): + if debug: + debugs = [img_bgr.copy()] + + img_size = img_bgr.shape[1], img_bgr.shape[0] + + img_face_mask_a = LandmarksProcessor.get_image_hull_mask (img_bgr.shape, img_face_landmarks) + img_face_mask_aaa = np.repeat(img_face_mask_a, 3, -1) + + output_size = self.predictor_input_size + face_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, output_size, face_type=FaceType.FULL) + + dst_face_mask_a_0 = cv2.warpAffine( img_face_mask_a, face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC ) + + predictor_input_dst_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (self.predictor_input_size,self.predictor_input_size), cv2.INTER_CUBIC ) + prd_inp_dst_face_mask_a = predictor_input_dst_face_mask_a_0[...,np.newaxis] + + prd_inp_avaperator_face_bgr = cv2.resize (avaperator_face_bgr, (self.predictor_input_size,self.predictor_input_size), cv2.INTER_CUBIC ) + + prd_face_bgr = self.predictor_func ( prd_inp_avaperator_face_bgr, prd_inp_dst_face_mask_a ) + + out_img = img_bgr.copy() + out_img = cv2.warpAffine( prd_face_bgr, face_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) + out_img = np.clip(out_img, 0.0, 1.0) + + if debug: + debugs += [out_img.copy()] + + out_img = np.clip( img_bgr*(1-img_face_mask_aaa) + (out_img*img_face_mask_aaa) , 0, 1.0 ) + + if debug: + debugs += [out_img.copy()] + + + return debugs if debug else out_img diff --git a/converters/ConverterImage.py b/converters/ConverterImage.py index a3dcc22..58b1faa 100644 --- a/converters/ConverterImage.py +++ b/converters/ConverterImage.py @@ -1,40 +1,50 @@ -from .Converter import Converter -from facelib import LandmarksProcessor -from facelib import FaceType +import time import cv2 import numpy as np -''' -predictor_func: - input: [predictor_input_size, predictor_input_size, BGR] - output: [predictor_input_size, predictor_input_size, BGR] -''' +from facelib import FaceType, LandmarksProcessor +from joblib import SubprocessFunctionCaller +from utils.pickle_utils import AntiPickler + +from .Converter import Converter class ConverterImage(Converter): #override def __init__(self, predictor_func, - predictor_input_size=0, - output_size=0): + predictor_input_size=0): super().__init__(predictor_func, Converter.TYPE_IMAGE) self.predictor_input_size = predictor_input_size - self.output_size = output_size + + #dummy predict and sleep, tensorflow caching kernels. If remove it, conversion speed will be x2 slower + predictor_func ( np.zeros ( (predictor_input_size,predictor_input_size,3), dtype=np.float32 ) ) + time.sleep(2) - #override - def dummy_predict(self): - self.predictor_func ( np.zeros ( (self.predictor_input_size, self.predictor_input_size,3), dtype=np.float32) ) + predictor_func_host, predictor_func = SubprocessFunctionCaller.make_pair(predictor_func) + self.predictor_func_host = AntiPickler(predictor_func_host) + self.predictor_func = predictor_func + #overridable + def on_host_tick(self): + self.predictor_func_host.obj.process_messages() + #override - def convert_image (self, img_bgr, img_landmarks, debug): + def cli_convert_image (self, img_bgr, img_landmarks, debug): img_size = img_bgr.shape[1], img_bgr.shape[0] predictor_input_bgr = cv2.resize ( img_bgr, (self.predictor_input_size, self.predictor_input_size), cv2.INTER_LANCZOS4 ) - predicted_bgr = self.predictor_func ( predictor_input_bgr ) + + if debug: + debugs = [predictor_input_bgr] + + output = self.predictor_func ( predictor_input_bgr ) - output = cv2.resize ( predicted_bgr, (self.output_size, self.output_size), cv2.INTER_LANCZOS4 ) if debug: return (predictor_input_bgr,output,) - return output + if debug: + debugs += [out_img.copy()] + + return debugs if debug else output diff --git a/converters/ConverterMasked.py b/converters/ConverterMasked.py index 1b7e3cc..fd853e9 100644 --- a/converters/ConverterMasked.py +++ b/converters/ConverterMasked.py @@ -30,7 +30,8 @@ class ConverterMasked(Converter): base_blur_mask_modifier = 0, default_erode_mask_modifier = 0, default_blur_mask_modifier = 0, - clip_hborder_mask_per = 0): + clip_hborder_mask_per = 0, + force_mask_mode=-1): super().__init__(predictor_func, Converter.TYPE_FACE) @@ -76,10 +77,13 @@ class ConverterMasked(Converter): if self.mode == 'hist-match' or self.mode == 'hist-match-bw' or self.mode == 'seamless-hist-match': self.hist_match_threshold = np.clip ( io.input_int("Hist match threshold [0..255] (skip:255) : ", 255), 0, 255) - if face_type == FaceType.FULL: - self.mask_mode = np.clip ( io.input_int ("Mask mode: (1) learned, (2) dst, (3) FAN-prd, (4) FAN-dst , (5) FAN-prd*FAN-dst (6) learned*FAN-prd*FAN-dst (?) help. Default - %d : " % (1) , 1, help_message="If you learned mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images. 'FAN-prd' - using super smooth mask by pretrained FAN-model from predicted face. 'FAN-dst' - using super smooth mask by pretrained FAN-model from dst face. 'FAN-prd*FAN-dst' or 'learned*FAN-prd*FAN-dst' - using multiplied masks."), 1, 6 ) + if force_mask_mode != -1: + self.mask_mode = force_mask_mode else: - self.mask_mode = np.clip ( io.input_int ("Mask mode: (1) learned, (2) dst . Default - %d : " % (1) , 1), 1, 2 ) + if face_type == FaceType.FULL: + self.mask_mode = np.clip ( io.input_int ("Mask mode: (1) learned, (2) dst, (3) FAN-prd, (4) FAN-dst , (5) FAN-prd*FAN-dst (6) learned*FAN-prd*FAN-dst (?) help. Default - %d : " % (1) , 1, help_message="If you learned mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images. 'FAN-prd' - using super smooth mask by pretrained FAN-model from predicted face. 'FAN-dst' - using super smooth mask by pretrained FAN-model from dst face. 'FAN-prd*FAN-dst' or 'learned*FAN-prd*FAN-dst' - using multiplied masks."), 1, 6 ) + else: + self.mask_mode = np.clip ( io.input_int ("Mask mode: (1) learned, (2) dst . Default - %d : " % (1) , 1), 1, 2 ) if self.mask_mode >= 3 and self.mask_mode <= 6: self.fan_seg = None @@ -118,10 +122,10 @@ class ConverterMasked(Converter): #overridable def on_cli_initialize(self): if (self.mask_mode >= 3 and self.mask_mode <= 6) and self.fan_seg == None: - self.fan_seg = FANSegmentator(256, FaceType.toString(FaceType.FULL) ) + self.fan_seg = FANSegmentator(256, FaceType.toString( self.face_type ) ) #override - def cli_convert_face (self, img_bgr, img_face_landmarks, debug): + def cli_convert_face (self, img_bgr, img_face_landmarks, debug, **kwargs): if debug: debugs = [img_bgr.copy()] @@ -171,13 +175,13 @@ class ConverterMasked(Converter): if self.mask_mode == 3 or self.mask_mode == 5 or self.mask_mode == 6: prd_face_bgr_256 = cv2.resize (prd_face_bgr, (256,256) ) - prd_face_bgr_256_mask = self.fan_seg.extract_from_bgr( prd_face_bgr_256[np.newaxis,...] ) [0] + prd_face_bgr_256_mask = self.fan_seg.extract( prd_face_bgr_256 ) FAN_prd_face_mask_a_0 = cv2.resize (prd_face_bgr_256_mask, (output_size,output_size), cv2.INTER_CUBIC) if self.mask_mode == 4 or self.mask_mode == 5 or self.mask_mode == 6: 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( dst_face_256_bgr[np.newaxis,...] ) [0] + dst_face_256_mask = self.fan_seg.extract( dst_face_256_bgr ) FAN_dst_face_mask_a_0 = cv2.resize (dst_face_256_mask, (output_size,output_size), cv2.INTER_CUBIC) if self.mask_mode == 3: #FAN-prd diff --git a/converters/__init__.py b/converters/__init__.py index 1471297..ef544cb 100644 --- a/converters/__init__.py +++ b/converters/__init__.py @@ -1,3 +1,4 @@ from .Converter import Converter from .ConverterMasked import ConverterMasked from .ConverterImage import ConverterImage +from .ConverterAvatar import ConverterAvatar diff --git a/doc/manual_en_google_translated.docx b/doc/manual_en_google_translated.docx index b1b9834..cd6c76d 100644 Binary files a/doc/manual_en_google_translated.docx and b/doc/manual_en_google_translated.docx differ diff --git a/doc/manual_en_google_translated.pdf b/doc/manual_en_google_translated.pdf index 4634b9c..828349d 100644 Binary files a/doc/manual_en_google_translated.pdf and b/doc/manual_en_google_translated.pdf differ diff --git a/doc/manual_ru.pdf b/doc/manual_ru.pdf index e55606a..6712983 100644 Binary files a/doc/manual_ru.pdf and b/doc/manual_ru.pdf differ diff --git a/doc/manual_ru_source.docx b/doc/manual_ru_source.docx index a4d0951..d733115 100644 Binary files a/doc/manual_ru_source.docx and b/doc/manual_ru_source.docx differ diff --git a/facelib/FANSeg_256_full_face.h5 b/facelib/FANSeg_256_full_face.h5 index 495764b..a9b2e59 100644 Binary files a/facelib/FANSeg_256_full_face.h5 and b/facelib/FANSeg_256_full_face.h5 differ diff --git a/facelib/FANSegmentator.py b/facelib/FANSegmentator.py index 09e14af..47a4fa1 100644 --- a/facelib/FANSegmentator.py +++ b/facelib/FANSegmentator.py @@ -1,15 +1,30 @@ -import numpy as np import os -import cv2 +import pickle +from functools import partial from pathlib import Path -from nnlib import nnlib + +import cv2 +import numpy as np + from interact import interact as io +from nnlib import nnlib + +""" +FANSegmentator is designed to segment faces aligned by 2DFAN-4 landmarks extractor. + +Dataset used to train located in official DFL mega.nz folder +https://mega.nz/#F!b9MzCK4B!zEAG9txu7uaRUjXz9PtBqg + +using https://github.com/ternaus/TernausNet +TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentation +""" class FANSegmentator(object): + VERSION = 1 def __init__ (self, resolution, face_type_str, load_weights=True, weights_file_root=None, training=False): exec( nnlib.import_all(), locals(), globals() ) - self.model = FANSegmentator.BuildModel(resolution, ngf=32) + self.model = FANSegmentator.BuildModel(resolution, ngf=64) if weights_file_root: weights_file_root = Path(weights_file_root) @@ -22,14 +37,21 @@ class FANSegmentator(object): self.model.load_weights (str(self.weights_path)) else: if training: - conv_weights_list = [] - for layer in self.model.layers: - if type(layer) == keras.layers.Conv2D: - conv_weights_list += [layer.weights[0]] # Conv2D kernel_weights - CAInitializerMP(conv_weights_list) - if training: - self.model.compile(loss='mse', optimizer=Adam(tf_cpu_mode=2)) + try: + with open( Path(__file__).parent / 'vgg11_enc_weights.npy', 'rb' ) as f: + d = pickle.loads (f.read()) + for i in [0,3,6,8,11,13,16,18]: + s = 'features.%d' % i + + self.model.get_layer (s).set_weights ( d[s] ) + except: + io.log_err("Unable to load VGG11 pretrained weights from vgg11_enc_weights.npy") + + if training: + #self.model.compile(loss='mse', optimizer=Adam(tf_cpu_mode=2)) + self.model.compile(loss='binary_crossentropy', optimizer=Adam(tf_cpu_mode=2), metrics=['accuracy']) + def __enter__(self): return self @@ -42,66 +64,76 @@ class FANSegmentator(object): def train_on_batch(self, inp, outp): return self.model.train_on_batch(inp, outp) - def extract_from_bgr (self, input_image): - return np.clip ( (self.model.predict(input_image) + 1) / 2.0, 0, 1.0 ) + def extract (self, input_image, is_input_tanh=False): + input_shape_len = len(input_image.shape) + if input_shape_len == 3: + input_image = input_image[np.newaxis,...] + + result = np.clip ( self.model.predict( [input_image] ), 0, 1.0 ) + result[result < 0.1] = 0 #get rid of noise + + if input_shape_len == 3: + result = result[0] + + return result @staticmethod - def BuildModel ( resolution, ngf=64): + def BuildModel ( resolution, ngf=64, norm='', act='lrelu'): exec( nnlib.import_all(), locals(), globals() ) inp = Input ( (resolution,resolution,3) ) x = inp - x = FANSegmentator.EncFlow(ngf=ngf)(x) - x = FANSegmentator.DecFlow(ngf=ngf)(x) + x = FANSegmentator.Flow(ngf=ngf, norm=norm, act=act)(x) model = Model(inp,x) return model @staticmethod - def EncFlow(ngf=64, num_downs=4): + def Flow(ngf=64, num_downs=4, norm='', act='lrelu'): exec( nnlib.import_all(), locals(), globals() ) - use_bias = True - def XNormalization(x): - return InstanceNormalization (axis=3, gamma_initializer=RandomNormal(1., 0.02))(x) - - def downscale (dim): - def func(x): - return LeakyReLU(0.1)(XNormalization(Conv2D(dim, kernel_size=5, strides=2, padding='same', kernel_initializer=RandomNormal(0, 0.02))(x))) - return func - def func(input): x = input - result = [] - for i in range(num_downs): - x = downscale ( min(ngf*(2**i), ngf*8) )(x) - result += [x] + x0 = x = Conv2D(ngf, kernel_size=3, strides=1, padding='same', activation='relu', name='features.0')(x) + x = MaxPooling2D()(x) + + x1 = x = Conv2D(ngf*2, kernel_size=3, strides=1, padding='same', activation='relu', name='features.3')(x) + x = MaxPooling2D()(x) + + x = Conv2D(ngf*4, kernel_size=3, strides=1, padding='same', activation='relu', name='features.6')(x) + x2 = x = Conv2D(ngf*4, kernel_size=3, strides=1, padding='same', activation='relu', name='features.8')(x) + x = MaxPooling2D()(x) + + x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.11')(x) + x3 = x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.13')(x) + x = MaxPooling2D()(x) + + x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.16')(x) + x4 = x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.18')(x) + x = MaxPooling2D()(x) + + x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same')(x) + + x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu') (x) + x = Concatenate(axis=3)([ x, x4]) + x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu') (x) + + x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu') (x) + x = Concatenate(axis=3)([ x, x3]) + x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu') (x) + + x = Conv2DTranspose (ngf*2, 3, strides=2, padding='same', activation='relu') (x) + x = Concatenate(axis=3)([ x, x2]) + x = Conv2D (ngf*4, 3, strides=1, padding='same', activation='relu') (x) + + x = Conv2DTranspose (ngf, 3, strides=2, padding='same', activation='relu') (x) + x = Concatenate(axis=3)([ x, x1]) + x = Conv2D (ngf*2, 3, strides=1, padding='same', activation='relu') (x) + + x = Conv2DTranspose (ngf // 2, 3, strides=2, padding='same', activation='relu') (x) + x = Concatenate(axis=3)([ x, x0]) + x = Conv2D (ngf, 3, strides=1, padding='same', activation='relu') (x) + + return Conv2D(1, 3, strides=1, padding='same', activation='sigmoid')(x) + - return result return func - - @staticmethod - def DecFlow(output_nc=1, ngf=64, activation='tanh'): - exec (nnlib.import_all(), locals(), globals()) - - use_bias = True - def XNormalization(x): - return InstanceNormalization (axis=3, gamma_initializer=RandomNormal(1., 0.02))(x) - - def Conv2D (filters, kernel_size, strides=(1, 1), padding='valid', data_format=None, dilation_rate=(1, 1), activation=None, use_bias=use_bias, kernel_initializer=RandomNormal(0, 0.02), bias_initializer='zeros', kernel_regularizer=None, bias_regularizer=None, activity_regularizer=None, kernel_constraint=None, bias_constraint=None): - return keras.layers.Conv2D( filters=filters, kernel_size=kernel_size, strides=strides, padding=padding, data_format=data_format, dilation_rate=dilation_rate, activation=activation, use_bias=use_bias, kernel_initializer=kernel_initializer, bias_initializer=bias_initializer, kernel_regularizer=kernel_regularizer, bias_regularizer=bias_regularizer, activity_regularizer=activity_regularizer, kernel_constraint=kernel_constraint, bias_constraint=bias_constraint ) - - def upscale (dim): - def func(x): - return SubpixelUpscaler()( LeakyReLU(0.1)(XNormalization(Conv2D(dim, kernel_size=3, strides=1, padding='same', kernel_initializer=RandomNormal(0, 0.02))(x)))) - return func - - def func(input): - input_len = len(input) - x = input[input_len-1] - for i in range(input_len-1, -1, -1): - x = upscale( min(ngf* (2**i) *4, ngf*8 *4 ) )(x) - if i != 0: - x = Concatenate(axis=3)([ input[i-1] , x]) - - return Conv2D(output_nc, 3, 1, 'same', activation=activation)(x) - return func \ No newline at end of file diff --git a/facelib/vgg11_enc_weights.npy b/facelib/vgg11_enc_weights.npy new file mode 100644 index 0000000..ea9df4e Binary files /dev/null and b/facelib/vgg11_enc_weights.npy differ diff --git a/joblib/SubprocessFunctionCaller.py b/joblib/SubprocessFunctionCaller.py index 7b3a528..4ea3101 100644 --- a/joblib/SubprocessFunctionCaller.py +++ b/joblib/SubprocessFunctionCaller.py @@ -8,9 +8,9 @@ class SubprocessFunctionCaller(object): self.c2s = c2s self.lock = lock - def __call__(self, value): + def __call__(self, *args, **kwargs): self.lock.acquire() - self.c2s.put (value) + self.c2s.put ( {'args':args, 'kwargs':kwargs} ) while True: if not self.s2c.empty(): obj = self.s2c.get() @@ -27,7 +27,7 @@ class SubprocessFunctionCaller(object): def process_messages(self): while not self.c2s.empty(): obj = self.c2s.get() - result = self.func (obj) + result = self.func ( *obj['args'], **obj['kwargs'] ) self.s2c.put (result) @staticmethod diff --git a/main.py b/main.py index 5f210d6..6e86e3b 100644 --- a/main.py +++ b/main.py @@ -49,6 +49,23 @@ if __name__ == "__main__": p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU. Forces to use MT extractor.") p.set_defaults (func=process_extract) + """ + def process_extract_fanseg(arguments): + os_utils.set_process_lowest_prio() + from mainscripts import Extractor + Extractor.extract_fanseg( arguments.input_dir, + device_args={'cpu_only' : arguments.cpu_only, + 'multi_gpu' : arguments.multi_gpu, + } + ) + + p = subparsers.add_parser( "extract_fanseg", help="Extract fanseg mask from faces.") + p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") + p.add_argument('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="Enables multi GPU.") + p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU.") + p.set_defaults (func=process_extract_fanseg) + """ + def process_sort(arguments): os_utils.set_process_lowest_prio() from mainscripts import Sorter @@ -71,12 +88,17 @@ if __name__ == "__main__": if arguments.recover_original_aligned_filename: Util.recover_original_aligned_filename (input_path=arguments.input_dir) - + + #if arguments.remove_fanseg: + # Util.remove_fanseg_folder (input_path=arguments.input_dir) + p = subparsers.add_parser( "util", help="Utilities.") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") p.add_argument('--convert-png-to-jpg', action="store_true", dest="convert_png_to_jpg", default=False, help="Convert DeepFaceLAB PNG files to JPEG.") p.add_argument('--add-landmarks-debug-images', action="store_true", dest="add_landmarks_debug_images", default=False, help="Add landmarks debug image for aligned faces.") p.add_argument('--recover-original-aligned-filename', action="store_true", dest="recover_original_aligned_filename", default=False, help="Recover original aligned filename.") + #p.add_argument('--remove-fanseg', action="store_true", dest="remove_fanseg", default=False, help="Remove fanseg mask from aligned faces.") + p.set_defaults (func=process_util) def process_train(arguments): @@ -112,6 +134,7 @@ if __name__ == "__main__": args = {'input_dir' : arguments.input_dir, 'output_dir' : arguments.output_dir, 'aligned_dir' : arguments.aligned_dir, + 'avaperator_aligned_dir' : arguments.avaperator_aligned_dir, 'model_dir' : arguments.model_dir, 'model_name' : arguments.model_name, 'debug' : arguments.debug, @@ -125,7 +148,8 @@ if __name__ == "__main__": p = subparsers.add_parser( "convert", help="Converter") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the converted files will be stored.") - p.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", help="Aligned directory. This is where the extracted of dst faces stored. Not used in AVATAR model.") + p.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", help="Aligned directory. This is where the extracted of dst faces stored.") + p.add_argument('--avaperator-aligned-dir', action=fixPathAction, dest="avaperator_aligned_dir", help="Only for AVATAR model. Directory of aligned avatar operator faces.") p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.") p.add_argument('--model', required=True, dest="model_name", choices=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Type of model") p.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug converter.") diff --git a/mainscripts/Converter.py b/mainscripts/Converter.py index 51c14c1..5ce9b7a 100644 --- a/mainscripts/Converter.py +++ b/mainscripts/Converter.py @@ -1,19 +1,23 @@ import sys +import multiprocessing +import operator import os +import shutil +import time import traceback from pathlib import Path -from utils import Path_utils + import cv2 -from utils.DFLPNG import DFLPNG -from utils.DFLJPG import DFLJPG -from utils.cv2_utils import * -import shutil import numpy as np -import time -import multiprocessing + from converters import Converter -from joblib import Subprocessor, SubprocessFunctionCaller from interact import interact as io +from joblib import SubprocessFunctionCaller, Subprocessor +from utils import Path_utils +from utils.cv2_utils import * +from utils.DFLJPG import DFLJPG +from utils.DFLPNG import DFLPNG +from imagelib import normalize_channels class ConvertSubprocessor(Subprocessor): class Cli(Subprocessor.Cli): @@ -26,6 +30,7 @@ class ConvertSubprocessor(Subprocessor): self.converter = client_dict['converter'] self.output_path = Path(client_dict['output_dir']) if 'output_dir' in client_dict.keys() else None self.alignments = client_dict['alignments'] + self.avatar_image_paths = client_dict['avatar_image_paths'] self.debug = client_dict['debug'] #transfer and set stdin in order to work code.interact in debug subprocess @@ -45,13 +50,15 @@ class ConvertSubprocessor(Subprocessor): #override def process_data(self, data): - filename_path = Path(data) + idx, filename = data + filename_path = Path(filename) files_processed = 1 faces_processed = 0 output_filename_path = self.output_path / (filename_path.stem + '.png') - if self.converter.type == Converter.TYPE_FACE and filename_path.stem not in self.alignments.keys(): + if (self.converter.type == Converter.TYPE_FACE or self.converter.type == Converter.TYPE_FACE_AVATAR ) \ + and filename_path.stem not in self.alignments.keys(): if not self.debug: self.log_info ( 'no faces found for %s, copying without faces' % (filename_path.name) ) @@ -62,19 +69,18 @@ class ConvertSubprocessor(Subprocessor): cv2_imwrite ( str(output_filename_path), image ) else: image = (cv2_imread(str(filename_path)) / 255.0).astype(np.float32) - h,w,c = image.shape - if c > 3: - image = image[...,0:3] - + image = normalize_channels (image, 3) + if self.converter.type == Converter.TYPE_IMAGE: - image = self.converter.convert_image(image, None, self.debug) + image = self.converter.cli_convert_image(image, None, self.debug) + if self.debug: - raise NotImplementedError - #for img in image: - # io.show_image ('Debug convert', img ) - # cv2.waitKey(0) + return (1, image) + faces_processed = 1 + elif self.converter.type == Converter.TYPE_IMAGE_WITH_LANDMARKS: + #currently unused if filename_path.suffix == '.png': dflimg = DFLPNG.load( str(filename_path) ) elif filename_path.suffix == '.jpg': @@ -85,7 +91,7 @@ class ConvertSubprocessor(Subprocessor): if dflimg is not None: image_landmarks = dflimg.get_landmarks() - image = self.converter.convert_image(image, image_landmarks, self.debug) + image = self.converter.convert_image(image, image_landmarks, self.debug) if self.debug: raise NotImplementedError @@ -96,7 +102,13 @@ class ConvertSubprocessor(Subprocessor): else: self.log_err ("%s is not a dfl image file" % (filename_path.name) ) - elif self.converter.type == Converter.TYPE_FACE: + elif self.converter.type == Converter.TYPE_FACE or self.converter.type == Converter.TYPE_FACE_AVATAR: + + ava_face = None + if self.converter.type == Converter.TYPE_FACE_AVATAR: + ava_filename_path = self.avatar_image_paths[idx] + ava_face = (cv2_imread(str(ava_filename_path)) / 255.0).astype(np.float32) + ava_face = normalize_channels (ava_face, 3) faces = self.alignments[filename_path.stem] if self.debug: @@ -108,9 +120,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.cli_convert_face(image, image_landmarks, self.debug) + debug_images += self.converter.cli_convert_face(image, image_landmarks, self.debug, avaperator_face_bgr=ava_face) else: - image = self.converter.cli_convert_face(image, image_landmarks, self.debug) + image = self.converter.cli_convert_face(image, image_landmarks, self.debug, avaperator_face_bgr=ava_face) except Exception as e: e_str = traceback.format_exc() @@ -133,16 +145,19 @@ class ConvertSubprocessor(Subprocessor): #overridable def get_data_name (self, data): #return string identificator of your data - return data + idx, filename = data + return filename #override - def __init__(self, converter, input_path_image_paths, output_path, alignments, debug = False): + def __init__(self, converter, input_path_image_paths, output_path, alignments, avatar_image_paths=None, debug = False): super().__init__('Converter', ConvertSubprocessor.Cli, 86400 if debug == True else 60) self.converter = converter self.input_data = self.input_path_image_paths = input_path_image_paths + self.input_data_idxs = [ *range(len(self.input_data)) ] self.output_path = output_path self.alignments = alignments + self.avatar_image_paths = avatar_image_paths self.debug = debug self.files_processed = 0 @@ -158,6 +173,7 @@ class ConvertSubprocessor(Subprocessor): 'converter' : self.converter, 'output_dir' : str(self.output_path), 'alignments' : self.alignments, + 'avatar_image_paths' : self.avatar_image_paths, 'debug': self.debug, 'stdin_fd': sys.stdin.fileno() if self.debug else None } @@ -167,7 +183,7 @@ class ConvertSubprocessor(Subprocessor): if self.debug: io.named_window ("Debug convert") - io.progress_bar ("Converting", len (self.input_data) ) + io.progress_bar ("Converting", len (self.input_data_idxs) ) #overridable optional def on_clients_finalized(self): @@ -178,13 +194,15 @@ class ConvertSubprocessor(Subprocessor): #override def get_data(self, host_dict): - if len (self.input_data) > 0: - return self.input_data.pop(0) + if len (self.input_data_idxs) > 0: + idx = self.input_data_idxs.pop(0) + return (idx, self.input_data[idx]) return None #override def on_data_return (self, host_dict, data): - self.input_data.insert(0, data) + idx, filename = data + self.input_data_idxs.insert(0, idx) #override def on_result (self, host_dict, data, result): @@ -209,7 +227,8 @@ def main (args, device_args): io.log_info ("Running converter.\r\n") aligned_dir = args.get('aligned_dir', None) - + avaperator_aligned_dir = args.get('avaperator_aligned_dir', None) + try: input_path = Path(args['input_dir']) output_path = Path(args['output_dir']) @@ -233,9 +252,10 @@ def main (args, device_args): model = models.import_model( args['model_name'] )(model_path, device_args=device_args) converter = model.get_converter() + input_path_image_paths = Path_utils.get_image_paths(input_path) alignments = None - - if converter.type == Converter.TYPE_FACE: + avatar_image_paths = None + if converter.type == Converter.TYPE_FACE or converter.type == Converter.TYPE_FACE_AVATAR: if aligned_dir is None: io.log_err('Aligned directory not found. Please ensure it exists.') return @@ -267,12 +287,45 @@ def main (args, device_args): alignments[ source_filename_stem ] = [] alignments[ source_filename_stem ].append (dflimg.get_source_landmarks()) + + + if converter.type == Converter.TYPE_FACE_AVATAR: + if avaperator_aligned_dir is None: + io.log_err('Avatar operator aligned directory not found. Please ensure it exists.') + return + avaperator_aligned_path = Path(avaperator_aligned_dir) + if not avaperator_aligned_path.exists(): + io.log_err('Avatar operator aligned directory not found. Please ensure it exists.') + return + + avatar_image_paths = [] + for filename in io.progress_bar_generator( Path_utils.get_image_paths(avaperator_aligned_path) , "Sorting avaperator faces"): + filepath = Path(filename) + if filepath.suffix == '.png': + dflimg = DFLPNG.load( str(filepath) ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filepath) ) + else: + dflimg = None + + if dflimg is None: + io.log_err ("Fatal error: %s is not a dfl image file" % (filepath.name) ) + return + + avatar_image_paths += [ (filename, dflimg.get_source_filename() ) ] + avatar_image_paths = [ p[0] for p in sorted(avatar_image_paths, key=operator.itemgetter(1)) ] + + if len(input_path_image_paths) < len(avatar_image_paths): + io.log_err("Input faces count must be >= avatar operator faces count.") + return + files_processed, faces_processed = ConvertSubprocessor ( converter = converter, - input_path_image_paths = Path_utils.get_image_paths(input_path), + input_path_image_paths = input_path_image_paths, output_path = output_path, alignments = alignments, + avatar_image_paths = avatar_image_paths, debug = args.get('debug',False) ).run() diff --git a/mainscripts/Extractor.py b/mainscripts/Extractor.py index eaee690..58328ad 100644 --- a/mainscripts/Extractor.py +++ b/mainscripts/Extractor.py @@ -10,11 +10,13 @@ import mathlib import imagelib import cv2 from utils import Path_utils +from utils.DFLPNG import DFLPNG from utils.DFLJPG import DFLJPG from utils.cv2_utils import * import facelib from facelib import FaceType from facelib import LandmarksProcessor +from facelib import FANSegmentator from nnlib import nnlib from joblib import Subprocessor from interact import interact as io @@ -79,7 +81,12 @@ class ExtractSubprocessor(Subprocessor): self.second_pass_e.__enter__() else: self.second_pass_e = None - + + elif self.type == 'fanseg': + nnlib.import_all (device_config) + self.e = facelib.FANSegmentator(256, FaceType.toString(FaceType.FULL) ) + self.e.__enter__() + elif self.type == 'final': pass @@ -124,6 +131,8 @@ class ExtractSubprocessor(Subprocessor): h, w, ch = image.shape if h == w: #extracting from already extracted jpg image? + if filename_path.suffix == '.png': + src_dflimg = DFLPNG.load ( str(filename_path) ) if filename_path.suffix == '.jpg': src_dflimg = DFLJPG.load ( str(filename_path) ) @@ -253,15 +262,22 @@ class ExtractSubprocessor(Subprocessor): cv2_imwrite(debug_output_file, debug_image, [int(cv2.IMWRITE_JPEG_QUALITY), 50] ) return data - - + + elif self.type == 'fanseg': + if src_dflimg is not None: + fanseg_mask = self.e.extract( image / 255.0 ) + src_dflimg.embed_and_set( filename_path_str, + fanseg_mask=fanseg_mask, + #fanseg_mask_ver=FANSegmentator.VERSION, + ) + #overridable def get_data_name (self, data): #return string identificator of your data return data.filename #override - def __init__(self, input_data, type, image_size, face_type, debug_dir=None, multi_gpu=False, cpu_only=False, manual=False, manual_window_size=0, final_output_path=None): + def __init__(self, input_data, type, image_size=None, face_type=None, debug_dir=None, multi_gpu=False, cpu_only=False, manual=False, manual_window_size=0, final_output_path=None): self.input_data = input_data self.type = type self.image_size = image_size @@ -561,7 +577,7 @@ class ExtractSubprocessor(Subprocessor): if 'cpu' in backend: cpu_only = True - if 'rects' in type or type == 'landmarks': + if 'rects' in type or type == 'landmarks' or type == 'fanseg': if not cpu_only and type == 'rects-mt' and backend == "plaidML": #plaidML works with MT very slowly cpu_only = True @@ -583,7 +599,7 @@ class ExtractSubprocessor(Subprocessor): dev_name = nnlib.device.getDeviceName(idx) dev_vram = nnlib.device.getDeviceVRAMTotalGb(idx) - if not manual and (type == 'rects-dlib' or type == 'rects-mt'): + if not manual and (type == 'rects-dlib' or type == 'rects-mt' ): for i in range ( int (max (1, dev_vram / 2) ) ): result += [ (idx, 'GPU', '%s #%d' % (dev_name,i) , dev_vram) ] else: @@ -658,7 +674,34 @@ class DeletedFilesSearcherSubprocessor(Subprocessor): return self.result +#currently unused +def extract_fanseg(input_dir, device_args={} ): + multi_gpu = device_args.get('multi_gpu', False) + cpu_only = device_args.get('cpu_only', False) + + input_path = Path(input_dir) + if not input_path.exists(): + raise ValueError('Input directory not found. Please ensure it exists.') + + paths_to_extract = [] + for filename in Path_utils.get_image_paths(input_path) : + filepath = Path(filename) + if filepath.suffix == '.png': + dflimg = DFLPNG.load( str(filepath) ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filepath) ) + else: + dflimg = None + if dflimg is not None: + paths_to_extract.append (filepath) + + paths_to_extract_len = len(paths_to_extract) + if paths_to_extract_len > 0: + io.log_info ("Performing extract fanseg for %d files..." % (paths_to_extract_len) ) + data = ExtractSubprocessor ([ ExtractSubprocessor.Data(filename) for filename in paths_to_extract ], 'fanseg', multi_gpu=multi_gpu, cpu_only=cpu_only).run() + + def main(input_dir, output_dir, debug_dir=None, diff --git a/mainscripts/MaskEditorTool.py b/mainscripts/MaskEditorTool.py index 9c48a7a..1d9750f 100644 --- a/mainscripts/MaskEditorTool.py +++ b/mainscripts/MaskEditorTool.py @@ -397,13 +397,17 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None): else: lmrks = dflimg.get_landmarks() ie_polys = dflimg.get_ie_polys() + fanseg_mask = dflimg.get_fanseg_mask() if filepath.name in cached_images: img = cached_images[filepath.name] else: img = cached_images[filepath.name] = cv2_imread(str(filepath)) / 255.0 - mask = LandmarksProcessor.get_image_hull_mask( img.shape, lmrks) + if fanseg_mask is not None: + mask = fanseg_mask + else: + mask = LandmarksProcessor.get_image_hull_mask( img.shape, lmrks) else: img = np.zeros ( (target_wh,target_wh,3) ) mask = np.ones ( (target_wh,target_wh,3) ) diff --git a/mainscripts/Util.py b/mainscripts/Util.py index a7c4a28..0d3e7c4 100644 --- a/mainscripts/Util.py +++ b/mainscripts/Util.py @@ -7,6 +7,33 @@ from utils.cv2_utils import * from facelib import LandmarksProcessor from interact import interact as io +def remove_fanseg_file (filepath): + filepath = Path(filepath) + + if filepath.suffix == '.png': + dflimg = DFLPNG.load( str(filepath) ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filepath) ) + else: + return + + if dflimg is None: + io.log_err ("%s is not a dfl image file" % (filepath.name) ) + return + + dflimg.remove_fanseg_mask() + dflimg.embed_and_set( str(filepath) ) + + +def remove_fanseg_folder(input_path): + input_path = Path(input_path) + + io.log_info ("Removing fanseg mask...\r\n") + + for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Removing"): + filepath = Path(filepath) + remove_fanseg_file(filepath) + def convert_png_to_jpg_file (filepath): filepath = Path(filepath) diff --git a/models/Model_FANSegmentator/Model.py b/models/Model_FANSEG/Model.py similarity index 66% rename from models/Model_FANSegmentator/Model.py rename to models/Model_FANSEG/Model.py index 0b0d58d..5e84c08 100644 --- a/models/Model_FANSegmentator/Model.py +++ b/models/Model_FANSEG/Model.py @@ -16,14 +16,23 @@ class Model(ModelBase): ask_sort_by_yaw=False, ask_random_flip=False, ask_src_scale_mod=False) - + + #override + def onInitializeOptions(self, is_first_run, ask_override): + default_face_type = 'f' + if is_first_run: + self.options['face_type'] = io.input_str ("Half or Full face? (h/f, ?:help skip:f) : ", default_face_type, ['h','f'], help_message="Half face has better resolution, but covers less area of cheeks.").lower() + else: + self.options['face_type'] = self.options.get('face_type', default_face_type) + #override def onInitialize(self): exec(nnlib.import_all(), locals(), globals()) self.set_vram_batch_requirements( {1.5:4} ) self.resolution = 256 - self.face_type = FaceType.FULL + self.face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF + self.fan_seg = FANSegmentator(self.resolution, FaceType.toString(self.face_type), @@ -33,18 +42,18 @@ class Model(ModelBase): if self.is_training_mode: f = SampleProcessor.TypeFlags - f_type = f.FACE_TYPE_FULL + face_type = f.FACE_TYPE_FULL if self.options['face_type'] == 'f' else f.FACE_TYPE_HALF self.set_training_data_generators ([ SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, - sample_process_options=SampleProcessor.Options(random_flip=True, motion_blur = [25, 1], normalize_tanh = True ), - output_sample_types=[ [f.TRANSFORMED | f_type | f.MODE_BGR_SHUFFLE | f.OPT_APPLY_MOTION_BLUR, self.resolution], - [f.TRANSFORMED | f_type | f.MODE_M | f.FACE_MASK_FULL, self.resolution] + sample_process_options=SampleProcessor.Options(random_flip=True, motion_blur = [25, 1] ), + output_sample_types=[ [f.WARPED_TRANSFORMED | face_type | f.MODE_BGR_SHUFFLE | f.OPT_APPLY_MOTION_BLUR, self.resolution], + [f.WARPED_TRANSFORMED | face_type | f.MODE_M | f.FACE_MASK_FULL, self.resolution] ]), SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, - sample_process_options=SampleProcessor.Options(random_flip=True, normalize_tanh = True ), - output_sample_types=[ [f.TRANSFORMED | f_type | f.MODE_BGR_SHUFFLE, self.resolution] + sample_process_options=SampleProcessor.Options(random_flip=True ), + output_sample_types=[ [f.TRANSFORMED | face_type | f.MODE_BGR_SHUFFLE, self.resolution] ]) ]) @@ -56,20 +65,18 @@ class Model(ModelBase): def onTrainOneIter(self, generators_samples, generators_list): target_src, target_src_mask = generators_samples[0] - loss = self.fan_seg.train_on_batch( [target_src], [target_src_mask] ) + loss,acc = self.fan_seg.train_on_batch( [target_src], [target_src_mask] ) - return ( ('loss', loss), ) + return ( ('loss', loss), ('acc',acc)) #override def onGetPreview(self, sample): test_A = sample[0][0][0:4] #first 4 samples test_B = sample[1][0][0:4] #first 4 samples - mAA = self.fan_seg.extract_from_bgr([test_A]) - mBB = self.fan_seg.extract_from_bgr([test_B]) - - test_A, test_B, = [ np.clip( (x + 1.0)/2.0, 0.0, 1.0) for x in [test_A, test_B] ] - + mAA = self.fan_seg.extract(test_A) + mBB = self.fan_seg.extract(test_B) + mAA = np.repeat ( mAA, (3,), -1) mBB = np.repeat ( mBB, (3,), -1) @@ -89,6 +96,6 @@ class Model(ModelBase): test_B[i,:,:,0:3]*mBB[i], ), axis=1) ) - return [ ('FANSegmentator', np.concatenate ( st, axis=0 ) ), - ('never seen', np.concatenate ( st2, axis=0 ) ), + return [ ('training data', np.concatenate ( st, axis=0 ) ), + ('evaluating data', np.concatenate ( st2, axis=0 ) ), ] diff --git a/models/Model_FANSegmentator/__init__.py b/models/Model_FANSEG/__init__.py similarity index 100% rename from models/Model_FANSegmentator/__init__.py rename to models/Model_FANSEG/__init__.py diff --git a/models/Model_RecycleGAN/Model.py b/models/Model_RecycleGAN/Model.py new file mode 100644 index 0000000..8014b27 --- /dev/null +++ b/models/Model_RecycleGAN/Model.py @@ -0,0 +1,477 @@ +from functools import partial + +import cv2 +import numpy as np + +from facelib import FaceType +from interact import interact as io +from mathlib import get_power_of_two +from models import ModelBase +from nnlib import nnlib +from samplelib import * + +class RecycleGANModel(ModelBase): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs, + ask_sort_by_yaw=False, + ask_random_flip=False, + ask_src_scale_mod=False) + + #override + def onInitializeOptions(self, is_first_run, ask_override): + if is_first_run: + self.options['resolution'] = io.input_int("Resolution ( 128,256 ?:help skip:128) : ", 128, [128,256], help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16.") + else: + self.options['resolution'] = self.options.get('resolution', 128) + + #override + def onInitialize(self, batch_size=-1, **in_options): + exec(nnlib.code_import_all, locals(), globals()) + self.set_vram_batch_requirements({6:16}) + + resolution = self.options['resolution'] + bgr_shape = (resolution, resolution, 3) + ngf = 64 + npf = 32 + ndf = 64 + lambda_A = 10 + lambda_B = 10 + + use_batch_norm = True #created_batch_size > 1 + self.GA = modelify(RecycleGANModel.ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=True))(Input(bgr_shape)) + self.GB = modelify(RecycleGANModel.ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=True))(Input(bgr_shape)) + + #self.GA = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape)) + #self.GB = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape)) + + self.PA = modelify(RecycleGANModel.UNetTemporalPredictor(bgr_shape[2], use_batch_norm, ngf=npf))([Input(bgr_shape), Input(bgr_shape)]) + self.PB = modelify(RecycleGANModel.UNetTemporalPredictor(bgr_shape[2], use_batch_norm, ngf=npf))([Input(bgr_shape), Input(bgr_shape)]) + + self.DA = modelify(RecycleGANModel.PatchDiscriminator(ndf=ndf) ) (Input(bgr_shape)) + self.DB = modelify(RecycleGANModel.PatchDiscriminator(ndf=ndf) ) (Input(bgr_shape)) + + if not self.is_first_run(): + weights_to_load = [ + (self.GA, 'GA.h5'), + (self.DA, 'DA.h5'), + (self.PA, 'PA.h5'), + (self.GB, 'GB.h5'), + (self.DB, 'DB.h5'), + (self.PB, 'PB.h5'), + ] + self.load_weights_safe(weights_to_load) + + real_A0 = Input(bgr_shape, name="real_A0") + real_A1 = Input(bgr_shape, name="real_A1") + real_A2 = Input(bgr_shape, name="real_A2") + + real_B0 = Input(bgr_shape, name="real_B0") + real_B1 = Input(bgr_shape, name="real_B1") + real_B2 = Input(bgr_shape, name="real_B2") + + DA_ones = K.ones_like ( K.shape(self.DA.outputs[0]) ) + DA_zeros = K.zeros_like ( K.shape(self.DA.outputs[0] )) + DB_ones = K.ones_like ( K.shape(self.DB.outputs[0] )) + DB_zeros = K.zeros_like ( K.shape(self.DB.outputs[0] )) + + def DLoss(labels,logits): + return K.mean(K.binary_crossentropy(labels,logits)) + + def CycleLoss (t1,t2): + return K.mean(K.abs(t1 - t2)) + + def RecurrentLOSS(t1,t2): + return K.mean(K.abs(t1 - t2)) + + def RecycleLOSS(t1,t2): + return K.mean(K.abs(t1 - t2)) + + fake_B0 = self.GA(real_A0) + fake_B1 = self.GA(real_A1) + + fake_A0 = self.GB(real_B0) + fake_A1 = self.GB(real_B1) + + real_A0_d = self.DA(real_A0) + real_A0_d_ones = K.ones_like(real_A0_d) + real_A1_d = self.DA(real_A1) + real_A1_d_ones = K.ones_like(real_A1_d) + + fake_A0_d = self.DA(fake_A0) + fake_A0_d_ones = K.ones_like(fake_A0_d) + fake_A0_d_zeros = K.zeros_like(fake_A0_d) + + fake_A1_d = self.DA(fake_A1) + fake_A1_d_ones = K.ones_like(fake_A1_d) + fake_A1_d_zeros = K.zeros_like(fake_A1_d) + + real_B0_d = self.DB(real_B0) + real_B0_d_ones = K.ones_like(real_B0_d) + + real_B1_d = self.DB(real_B1) + real_B1_d_ones = K.ones_like(real_B1_d) + + fake_B0_d = self.DB(fake_B0) + fake_B0_d_ones = K.ones_like(fake_B0_d) + fake_B0_d_zeros = K.zeros_like(fake_B0_d) + + fake_B1_d = self.DB(fake_B1) + fake_B1_d_ones = K.ones_like(fake_B1_d) + fake_B1_d_zeros = K.zeros_like(fake_B1_d) + + pred_A2 = self.PA ( [real_A0, real_A1]) + pred_B2 = self.PB ( [real_B0, real_B1]) + rec_A2 = self.GB ( self.PB ( [fake_B0, fake_B1]) ) + rec_B2 = self.GA ( self.PA ( [fake_A0, fake_A1])) + + + loss_GA = DLoss(fake_B0_d_ones, fake_B0_d ) + \ + DLoss(fake_B1_d_ones, fake_B1_d ) + \ + lambda_A * (RecurrentLOSS(pred_A2, real_A2) + \ + RecycleLOSS(rec_B2, real_B2) ) + + + weights_GA = self.GA.trainable_weights + self.PA.trainable_weights + + loss_GB = DLoss(fake_A0_d_ones, fake_A0_d ) + \ + DLoss(fake_A1_d_ones, fake_A1_d ) + \ + lambda_B * (RecurrentLOSS(pred_B2, real_B2) + \ + RecycleLOSS(rec_A2, real_A2) ) + + weights_GB = self.GB.trainable_weights + self.PB.trainable_weights + + def opt(): + return Adam(lr=2e-4, beta_1=0.5, beta_2=0.999, tf_cpu_mode=2)#, clipnorm=1) + + self.GA_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_GA], + opt().get_updates(loss_GA, weights_GA) ) + + self.GB_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_GB], + opt().get_updates(loss_GB, weights_GB) ) + + ########### + + loss_D_A0 = ( DLoss(real_A0_d_ones, real_A0_d ) + \ + DLoss(fake_A0_d_zeros, fake_A0_d ) ) * 0.5 + + loss_D_A1 = ( DLoss(real_A1_d_ones, real_A1_d ) + \ + DLoss(fake_A1_d_zeros, fake_A1_d ) ) * 0.5 + + loss_D_A = loss_D_A0 + loss_D_A1 + + self.DA_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_D_A], + opt().get_updates(loss_D_A, self.DA.trainable_weights) ) + + ############ + + loss_D_B0 = ( DLoss(real_B0_d_ones, real_B0_d ) + \ + DLoss(fake_B0_d_zeros, fake_B0_d ) ) * 0.5 + + loss_D_B1 = ( DLoss(real_B1_d_ones, real_B1_d ) + \ + DLoss(fake_B1_d_zeros, fake_B1_d ) ) * 0.5 + + loss_D_B = loss_D_B0 + loss_D_B1 + + self.DB_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_D_B], + opt().get_updates(loss_D_B, self.DB.trainable_weights) ) + + ############ + + + self.G_view = K.function([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[fake_A0, fake_A1, pred_A2, rec_A2, fake_B0, fake_B1, pred_B2, rec_B2 ]) + + + + if self.is_training_mode: + f = SampleProcessor.TypeFlags + self.set_training_data_generators ([ + SampleGeneratorImageTemporal(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, + temporal_image_count=3, + sample_process_options=SampleProcessor.Options(random_flip = False, normalize_tanh = True), + output_sample_types=[ [f.SOURCE | f.MODE_BGR, resolution] ] ), + + SampleGeneratorImageTemporal(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, + temporal_image_count=3, + sample_process_options=SampleProcessor.Options(random_flip = False, normalize_tanh = True), + output_sample_types=[ [f.SOURCE | f.MODE_BGR, resolution] ] ), + ]) + else: + self.G_convert = K.function([real_B0],[fake_A0]) + + #override + def onSave(self): + self.save_weights_safe( [[self.GA, 'GA.h5'], + [self.GB, 'GB.h5'], + [self.DA, 'DA.h5'], + [self.DB, 'DB.h5'], + [self.PA, 'PA.h5'], + [self.PB, 'PB.h5'] ]) + + #override + def onTrainOneIter(self, generators_samples, generators_list): + source_src_0, source_src_1, source_src_2, = generators_samples[0] + source_dst_0, source_dst_1, source_dst_2, = generators_samples[1] + + feed = [source_src_0, source_src_1, source_src_2, source_dst_0, source_dst_1, source_dst_2] + + loss_GA, = self.GA_train ( feed ) + loss_GB, = self.GB_train ( feed ) + loss_DA, = self.DA_train( feed ) + loss_DB, = self.DB_train( feed ) + + return ( ('GA', loss_GA), ('GB', loss_GB), ('DA', loss_DA), ('DB', loss_DB) ) + + #override + def onGetPreview(self, sample): + test_A0 = sample[0][0] + test_A1 = sample[0][1] + test_A2 = sample[0][2] + + test_B0 = sample[1][0] + test_B1 = sample[1][1] + test_B2 = sample[1][2] + + G_view_result = self.G_view([test_A0, test_A1, test_A2, test_B0, test_B1, test_B2]) + + fake_A0, fake_A1, pred_A2, rec_A2, fake_B0, fake_B1, pred_B2, rec_B2 = [ x[0] / 2 + 0.5 for x in G_view_result] + test_A0, test_A1, test_A2, test_B0, test_B1, test_B2 = [ x[0] / 2 + 0.5 for x in [test_A0, test_A1, test_A2, test_B0, test_B1, test_B2] ] + + r = np.concatenate ((np.concatenate ( (test_A0, test_A1, test_A2, pred_A2, fake_B0, fake_B1, rec_A2), axis=1), + np.concatenate ( (test_B0, test_B1, test_B2, pred_B2, fake_A0, fake_A1, rec_B2), axis=1) + ), axis=0) + + return [ ('RecycleGAN', r ) ] + + def predictor_func (self, face): + x = self.G_convert ( [ face[np.newaxis,...]*2-1 ] )[0] + return np.clip ( x[0] / 2 + 0.5 , 0, 1) + + #override + def get_converter(self, **in_options): + from converters import ConverterImage + return ConverterImage(self.predictor_func, + predictor_input_size=self.options['resolution'], + **in_options) + + @staticmethod + def ResNet(output_nc, use_batch_norm, ngf=64, n_blocks=6, use_dropout=False): + exec (nnlib.import_all(), locals(), globals()) + + if not use_batch_norm: + use_bias = True + def XNormalization(x): + return InstanceNormalization (axis=-1)(x) + else: + use_bias = False + def XNormalization(x): + return BatchNormalization (axis=-1)(x) + + XConv2D = partial(Conv2D, padding='same', use_bias=use_bias) + XConv2DTranspose = partial(Conv2DTranspose, padding='same', use_bias=use_bias) + + def func(input): + + + def ResnetBlock(dim, use_dropout=False): + def func(input): + x = input + + x = XConv2D(dim, 3, strides=1)(x) + x = XNormalization(x) + x = ReLU()(x) + + if use_dropout: + x = Dropout(0.5)(x) + + x = XConv2D(dim, 3, strides=1)(x) + x = XNormalization(x) + x = ReLU()(x) + return Add()([x,input]) + return func + + x = input + + x = ReLU()(XNormalization(XConv2D(ngf, 7, strides=1)(x))) + + x = ReLU()(XNormalization(XConv2D(ngf*2, 3, strides=2)(x))) + x = ReLU()(XNormalization(XConv2D(ngf*4, 3, strides=2)(x))) + + for i in range(n_blocks): + x = ResnetBlock(ngf*4, use_dropout=use_dropout)(x) + + x = ReLU()(XNormalization(XConv2DTranspose(ngf*2, 3, strides=2)(x))) + x = ReLU()(XNormalization(XConv2DTranspose(ngf , 3, strides=2)(x))) + + x = XConv2D(output_nc, 7, strides=1, activation='tanh', use_bias=True)(x) + + return x + + return func + + @staticmethod + def UNet(output_nc, use_batch_norm, ngf=64, use_dropout=False): + exec (nnlib.import_all(), locals(), globals()) + + if not use_batch_norm: + use_bias = True + def XNormalizationL(): + return InstanceNormalization (axis=-1) + else: + use_bias = False + def XNormalizationL(): + return BatchNormalization (axis=-1) + + def XNormalization(x): + return XNormalizationL()(x) + + XConv2D = partial(Conv2D, padding='same', use_bias=use_bias) + XConv2DTranspose = partial(Conv2DTranspose, padding='same', use_bias=use_bias) + + def func(input): + + b,h,w,c = K.int_shape(input) + + n_downs = get_power_of_two(w) - 4 + + Norm = XNormalizationL() + Norm2 = XNormalizationL() + Norm4 = XNormalizationL() + Norm8 = XNormalizationL() + + x = input + + x = e1 = XConv2D( ngf, 4, strides=2, use_bias=True ) (x) + + x = e2 = Norm2( XConv2D( ngf*2, 4, strides=2 )( LeakyReLU(0.2)(x) ) ) + x = e3 = Norm4( XConv2D( ngf*4, 4, strides=2 )( LeakyReLU(0.2)(x) ) ) + + l = [] + for i in range(n_downs): + x = Norm8( XConv2D( ngf*8, 4, strides=2 )( LeakyReLU(0.2)(x) ) ) + l += [x] + + x = XConv2D( ngf*8, 4, strides=2, use_bias=True )( LeakyReLU(0.2)(x) ) + + for i in range(n_downs): + x = Norm8( XConv2DTranspose( ngf*8, 4, strides=2 )( ReLU()(x) ) ) + if i <= n_downs-2: + x = Dropout(0.5)(x) + x = Concatenate(axis=-1)([x, l[-i-1] ]) + + x = Norm4( XConv2DTranspose( ngf*4, 4, strides=2 )( ReLU()(x) ) ) + x = Concatenate(axis=-1)([x, e3]) + + x = Norm2( XConv2DTranspose( ngf*2, 4, strides=2 )( ReLU()(x) ) ) + x = Concatenate(axis=-1)([x, e2]) + + x = Norm( XConv2DTranspose( ngf, 4, strides=2 )( ReLU()(x) ) ) + x = Concatenate(axis=-1)([x, e1]) + + x = XConv2DTranspose(output_nc, 4, strides=2, activation='tanh', use_bias=True)( ReLU()(x) ) + + return x + return func + nnlib.UNet = UNet + + @staticmethod + def UNetTemporalPredictor(output_nc, use_batch_norm, ngf=64, use_dropout=False): + exec (nnlib.import_all(), locals(), globals()) + def func(inputs): + past_2_image_tensor, past_1_image_tensor = inputs + + x = Concatenate(axis=-1)([ past_2_image_tensor, past_1_image_tensor ]) + x = UNet(3, use_batch_norm, ngf=ngf, use_dropout=use_dropout) (x) + + return x + + return func + + @staticmethod + def PatchDiscriminator(ndf=64): + exec (nnlib.import_all(), locals(), globals()) + + #use_bias = True + #def XNormalization(x): + # return InstanceNormalization (axis=-1)(x) + use_bias = False + def XNormalization(x): + return BatchNormalization (axis=-1)(x) + + XConv2D = partial(Conv2D, use_bias=use_bias) + + def func(input): + b,h,w,c = K.int_shape(input) + + x = input + + x = ZeroPadding2D((1,1))(x) + x = XConv2D( ndf, 4, strides=2, padding='valid', use_bias=True)(x) + x = LeakyReLU(0.2)(x) + + x = ZeroPadding2D((1,1))(x) + x = XConv2D( ndf*2, 4, strides=2, padding='valid')(x) + x = XNormalization(x) + x = LeakyReLU(0.2)(x) + + x = ZeroPadding2D((1,1))(x) + x = XConv2D( ndf*4, 4, strides=2, padding='valid')(x) + x = XNormalization(x) + x = LeakyReLU(0.2)(x) + + x = ZeroPadding2D((1,1))(x) + x = XConv2D( ndf*8, 4, strides=2, padding='valid')(x) + x = XNormalization(x) + x = LeakyReLU(0.2)(x) + + x = ZeroPadding2D((1,1))(x) + x = XConv2D( ndf*8, 4, strides=2, padding='valid')(x) + x = XNormalization(x) + x = LeakyReLU(0.2)(x) + + x = ZeroPadding2D((1,1))(x) + return XConv2D( 1, 4, strides=1, padding='valid', use_bias=True, activation='sigmoid')(x)# + return func + + @staticmethod + def NLayerDiscriminator(ndf=64, n_layers=3): + exec (nnlib.import_all(), locals(), globals()) + + #use_bias = True + #def XNormalization(x): + # return InstanceNormalization (axis=-1)(x) + use_bias = False + def XNormalization(x): + return BatchNormalization (axis=-1)(x) + + XConv2D = partial(Conv2D, use_bias=use_bias) + + def func(input): + b,h,w,c = K.int_shape(input) + + x = input + + f = ndf + + x = ZeroPadding2D((1,1))(x) + x = XConv2D( f, 4, strides=2, padding='valid', use_bias=True)(x) + f = min( ndf*8, f*2 ) + x = LeakyReLU(0.2)(x) + + for i in range(n_layers): + x = ZeroPadding2D((1,1))(x) + x = XConv2D( f, 4, strides=2, padding='valid')(x) + f = min( ndf*8, f*2 ) + x = XNormalization(x) + x = LeakyReLU(0.2)(x) + + x = ZeroPadding2D((1,1))(x) + x = XConv2D( f, 4, strides=1, padding='valid')(x) + x = XNormalization(x) + x = LeakyReLU(0.2)(x) + + x = ZeroPadding2D((1,1))(x) + return XConv2D( 1, 4, strides=1, padding='valid', use_bias=True, activation='sigmoid')(x)# + return func + +Model = RecycleGANModel diff --git a/models/Model_RecycleGAN/__init__.py b/models/Model_RecycleGAN/__init__.py new file mode 100644 index 0000000..cdb3fe7 --- /dev/null +++ b/models/Model_RecycleGAN/__init__.py @@ -0,0 +1 @@ +from .Model import Model \ No newline at end of file diff --git a/models/Model_SAE/Model.py b/models/Model_SAE/Model.py index 3ddbe7e..478fad4 100644 --- a/models/Model_SAE/Model.py +++ b/models/Model_SAE/Model.py @@ -1,3 +1,4 @@ +from functools import partial import numpy as np from nnlib import nnlib @@ -385,20 +386,20 @@ class SAEModel(ModelBase): #override def onGetPreview(self, sample): - test_A = sample[0][1][0:4] #first 4 samples - test_A_m = sample[0][2][0:4] #first 4 samples - test_B = sample[1][1][0:4] - test_B_m = sample[1][2][0:4] + test_S = sample[0][1][0:4] #first 4 samples + test_S_m = sample[0][2][0:4] #first 4 samples + test_D = sample[1][1][0:4] + test_D_m = sample[1][2][0:4] if self.options['learn_mask']: - S, D, SS, DD, DDM, SD, SDM = [ np.clip(x, 0.0, 1.0) for x in ([test_A,test_B] + self.AE_view ([test_A, test_B]) ) ] + S, D, SS, DD, DDM, SD, SDM = [ np.clip(x, 0.0, 1.0) for x in ([test_S,test_D] + self.AE_view ([test_S, test_D]) ) ] DDM, SDM, = [ np.repeat (x, (3,), -1) for x in [DDM, SDM] ] else: - S, D, SS, DD, SD, = [ np.clip(x, 0.0, 1.0) for x in ([test_A,test_B] + self.AE_view ([test_A, test_B]) ) ] + S, D, SS, DD, SD, = [ np.clip(x, 0.0, 1.0) for x in ([test_S,test_D] + self.AE_view ([test_S, test_D]) ) ] result = [] st = [] - for i in range(0, len(test_A)): + for i in range(0, len(test_S)): ar = S[i], SS[i], D[i], DD[i], SD[i] st.append ( np.concatenate ( ar, axis=1) ) @@ -406,12 +407,12 @@ class SAEModel(ModelBase): if self.options['learn_mask']: st_m = [] - for i in range(0, len(test_A)): - ar = S[i], SS[i], D[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i]) + for i in range(0, len(test_S)): + ar = S[i]*test_S_m[i], SS[i], D[i]*test_D_m[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i]) st_m.append ( np.concatenate ( ar, axis=1) ) result += [ ('SAE masked', np.concatenate (st_m, axis=0 )), ] - + return result def predictor_func (self, face): @@ -485,57 +486,29 @@ class SAEModel(ModelBase): return x SAEModel.ResidualBlock = ResidualBlock - def ResidualBlock_pre (**base_kwargs): - def func(*args, **kwargs): - kwargs.update(base_kwargs) - return ResidualBlock(*args, **kwargs) - return func - SAEModel.ResidualBlock_pre = ResidualBlock_pre - def downscale (dim, padding='zero', norm='', act='', **kwargs): def func(x): return Norm(norm)( Act(act) (Conv2D(dim, kernel_size=5, strides=2, padding=padding)(x)) ) return func SAEModel.downscale = downscale - def downscale_pre (**base_kwargs): - def func(*args, **kwargs): - kwargs.update(base_kwargs) - return downscale(*args, **kwargs) - return func - SAEModel.downscale_pre = downscale_pre - def upscale (dim, padding='zero', norm='', act='', **kwargs): def func(x): return SubpixelUpscaler()(Norm(norm)(Act(act)(Conv2D(dim * 4, kernel_size=3, strides=1, padding=padding)(x)))) return func SAEModel.upscale = upscale - def upscale_pre (**base_kwargs): - def func(*args, **kwargs): - kwargs.update(base_kwargs) - return upscale(*args, **kwargs) - return func - SAEModel.upscale_pre = upscale_pre - def to_bgr (output_nc, padding='zero', **kwargs): def func(x): return Conv2D(output_nc, kernel_size=5, padding=padding, activation='sigmoid')(x) return func SAEModel.to_bgr = to_bgr - def to_bgr_pre (**base_kwargs): - def func(*args, **kwargs): - kwargs.update(base_kwargs) - return to_bgr(*args, **kwargs) - return func - SAEModel.to_bgr_pre = to_bgr_pre - @staticmethod def LIAEEncFlow(resolution, ch_dims, **kwargs): exec (nnlib.import_all(), locals(), globals()) - upscale = SAEModel.upscale_pre(**kwargs) - downscale = SAEModel.downscale_pre(**kwargs) + upscale = partial(SAEModel.upscale, **kwargs) + downscale = partial(SAEModel.downscale, **kwargs) def func(input): dims = K.int_shape(input)[-1]*ch_dims @@ -553,7 +526,7 @@ class SAEModel(ModelBase): @staticmethod def LIAEInterFlow(resolution, ae_dims=256, **kwargs): exec (nnlib.import_all(), locals(), globals()) - upscale = SAEModel.upscale_pre(**kwargs) + upscale = partial(SAEModel.upscale, **kwargs) lowest_dense_res=resolution // 16 def func(input): @@ -568,10 +541,10 @@ class SAEModel(ModelBase): @staticmethod def LIAEDecFlow(output_nc,ch_dims, multiscale_count=1, add_residual_blocks=False, padding='zero', norm='', **kwargs): exec (nnlib.import_all(), locals(), globals()) - upscale = SAEModel.upscale_pre(**kwargs) - to_bgr = SAEModel.to_bgr_pre(**kwargs) + upscale = partial(SAEModel.upscale, **kwargs) + to_bgr = partial(SAEModel.to_bgr, **kwargs) dims = output_nc * ch_dims - ResidualBlock = SAEModel.ResidualBlock_pre(**kwargs) + ResidualBlock = partial(SAEModel.ResidualBlock, **kwargs) def func(input): x = input[0] @@ -609,8 +582,8 @@ class SAEModel(ModelBase): @staticmethod def DFEncFlow(resolution, ae_dims, ch_dims, padding='zero', **kwargs): exec (nnlib.import_all(), locals(), globals()) - upscale = SAEModel.upscale_pre(padding=padding) - downscale = SAEModel.downscale_pre(padding=padding) + upscale = partial(SAEModel.upscale, padding=padding) + downscale = partial(SAEModel.downscale, padding=padding) lowest_dense_res = resolution // 16 def func(input): @@ -634,10 +607,10 @@ class SAEModel(ModelBase): @staticmethod def DFDecFlow(output_nc, ch_dims, multiscale_count=1, add_residual_blocks=False, padding='zero', **kwargs): exec (nnlib.import_all(), locals(), globals()) - upscale = SAEModel.upscale_pre(padding=padding) - to_bgr = SAEModel.to_bgr_pre(padding=padding) + upscale = partial(SAEModel.upscale, padding=padding) + to_bgr = partial(SAEModel.to_bgr, padding=padding) dims = output_nc * ch_dims - ResidualBlock = SAEModel.ResidualBlock_pre(padding=padding) + ResidualBlock = partial(SAEModel.ResidualBlock, padding=padding) def func(input): x = input[0] diff --git a/models/archived_models.zip b/models/archived_models.zip index 0829194..6e15f87 100644 Binary files a/models/archived_models.zip and b/models/archived_models.zip differ diff --git a/nnlib/device.py b/nnlib/device.py index 672a63c..5653623 100644 --- a/nnlib/device.py +++ b/nnlib/device.py @@ -274,28 +274,34 @@ has_nvml_cap = False #- CUDA build of DFL has_nvidia_device = os.environ.get("DFL_FORCE_HAS_NVIDIA_DEVICE", "0") == "1" -plaidML_devices = [] - -# Using plaidML OpenCL backend to determine system devices and has_nvidia_device -try: - os.environ['PLAIDML_EXPERIMENTAL'] = 'false' #this enables work plaidML without run 'plaidml-setup' - import plaidml - ctx = plaidml.Context() - for d in plaidml.devices(ctx, return_all=True)[0]: - details = json.loads(d.details) - if details['type'] == 'CPU': #skipping opencl-CPU - continue - if 'nvidia' in details['vendor'].lower(): - has_nvidia_device = True - plaidML_devices += [ {'id':d.id, - 'globalMemSize' : int(details['globalMemSize']), - 'description' : d.description.decode() - }] - ctx.shutdown() -except: - pass - -plaidML_devices_count = len(plaidML_devices) +plaidML_devices = None +def get_plaidML_devices(): + global plaidML_devices + global has_nvidia_device + if plaidML_devices is None: + plaidML_devices = [] + # Using plaidML OpenCL backend to determine system devices and has_nvidia_device + try: + os.environ['PLAIDML_EXPERIMENTAL'] = 'false' #this enables work plaidML without run 'plaidml-setup' + import plaidml + ctx = plaidml.Context() + for d in plaidml.devices(ctx, return_all=True)[0]: + details = json.loads(d.details) + if details['type'] == 'CPU': #skipping opencl-CPU + continue + if 'nvidia' in details['vendor'].lower(): + has_nvidia_device = True + plaidML_devices += [ {'id':d.id, + 'globalMemSize' : int(details['globalMemSize']), + 'description' : d.description.decode() + }] + ctx.shutdown() + except: + pass + return plaidML_devices + +if not has_nvidia_device: + get_plaidML_devices() #choosing backend @@ -324,7 +330,7 @@ if device.backend is None and not force_tf_cpu: if force_plaidML or (device.backend is None and not has_nvidia_device): #tensorflow backend was failed without has_nvidia_device , or forcing plaidML, trying to use plaidML backend - if plaidML_devices_count == 0: + if len(get_plaidML_devices()) == 0: #print ("plaidML: No capable OpenCL devices found. Falling back to tensorflow backend.") device.backend = None else: diff --git a/nnlib/nnlib.py b/nnlib/nnlib.py index 50003a6..7279218 100644 --- a/nnlib/nnlib.py +++ b/nnlib/nnlib.py @@ -52,7 +52,7 @@ Input = KL.Input Dense = KL.Dense Conv2D = nnlib.Conv2D -Conv2DTranspose = KL.Conv2DTranspose +Conv2DTranspose = nnlib.Conv2DTranspose SeparableConv2D = KL.SeparableConv2D MaxPooling2D = KL.MaxPooling2D UpSampling2D = KL.UpSampling2D @@ -695,6 +695,26 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator x = ReflectionPadding2D( self.pad ) (x) return self.func(x) nnlib.Conv2D = Conv2D + + class Conv2DTranspose(): + def __init__ (self, *args, **kwargs): + self.reflect_pad = False + padding = kwargs.get('padding','') + if padding == 'zero': + kwargs['padding'] = 'same' + if padding == 'reflect': + kernel_size = kwargs['kernel_size'] + if (kernel_size % 2) == 1: + self.pad = (kernel_size // 2,)*2 + kwargs['padding'] = 'valid' + self.reflect_pad = True + self.func = keras.layers.Conv2DTranspose (*args, **kwargs) + + def __call__(self,x): + if self.reflect_pad: + x = ReflectionPadding2D( self.pad ) (x) + return self.func(x) + nnlib.Conv2DTranspose = Conv2DTranspose @staticmethod def import_keras_contrib(device_config): diff --git a/samplelib/Sample.py b/samplelib/Sample.py index 1c311e2..7a65ae0 100644 --- a/samplelib/Sample.py +++ b/samplelib/Sample.py @@ -1,7 +1,13 @@ from enum import IntEnum +from pathlib import Path + import cv2 import numpy as np + from utils.cv2_utils import * +from utils.DFLJPG import DFLJPG +from utils.DFLPNG import DFLPNG + class SampleType(IntEnum): IMAGE = 0 #raw image @@ -10,13 +16,12 @@ class SampleType(IntEnum): FACE = 1 #aligned face unsorted FACE_YAW_SORTED = 2 #sorted by yaw FACE_YAW_SORTED_AS_TARGET = 3 #sorted by yaw and included only yaws which exist in TARGET also automatic mirrored - FACE_WITH_CLOSE_TO_SELF = 4 - FACE_END = 4 + FACE_END = 3 - QTY = 5 + QTY = 4 class Sample(object): - def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, mirror=None, close_target_list=None): + def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, source_filename=None, mirror=None, close_target_list=None, fanseg_mask_exist=False): self.sample_type = sample_type if sample_type is not None else SampleType.IMAGE self.filename = filename self.face_type = face_type @@ -25,10 +30,12 @@ class Sample(object): self.ie_polys = ie_polys self.pitch = pitch self.yaw = yaw + self.source_filename = source_filename self.mirror = mirror self.close_target_list = close_target_list + self.fanseg_mask_exist = fanseg_mask_exist - def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, mirror=None, close_target_list=None): + def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, source_filename=None, mirror=None, close_target_list=None, fanseg_mask=None, fanseg_mask_exist=None): return Sample( sample_type=sample_type if sample_type is not None else self.sample_type, filename=filename if filename is not None else self.filename, @@ -38,8 +45,10 @@ class Sample(object): ie_polys=ie_polys if ie_polys is not None else self.ie_polys, pitch=pitch if pitch is not None else self.pitch, yaw=yaw if yaw is not None else self.yaw, + source_filename=source_filename if source_filename is not None else self.source_filename, mirror=mirror if mirror is not None else self.mirror, - close_target_list=close_target_list if close_target_list is not None else self.close_target_list) + close_target_list=close_target_list if close_target_list is not None else self.close_target_list, + fanseg_mask_exist=fanseg_mask_exist if fanseg_mask_exist is not None else self.fanseg_mask_exist) def load_bgr(self): img = cv2_imread (self.filename).astype(np.float32) / 255.0 @@ -47,6 +56,19 @@ class Sample(object): img = img[:,::-1].copy() return img + def load_fanseg_mask(self): + if self.fanseg_mask_exist: + filepath = Path(self.filename) + if filepath.suffix == '.png': + dflimg = DFLPNG.load ( str(filepath) ) + elif filepath.suffix == '.jpg': + dflimg = DFLJPG.load ( str(filepath) ) + else: + dflimg = None + return dflimg.get_fanseg_mask() + + return None + def get_random_close_target_sample(self): if self.close_target_list is None: return None diff --git a/samplelib/SampleGeneratorFace.py b/samplelib/SampleGeneratorFace.py index dd7cf2d..6a9636d 100644 --- a/samplelib/SampleGeneratorFace.py +++ b/samplelib/SampleGeneratorFace.py @@ -15,7 +15,7 @@ output_sample_types = [ ] ''' class SampleGeneratorFace(SampleGeneratorBase): - def __init__ (self, samples_path, debug, batch_size, sort_by_yaw=False, sort_by_yaw_target_samples_path=None, with_close_to_self=False, sample_process_options=SampleProcessor.Options(), output_sample_types=[], add_sample_idx=False, add_pitch=False, add_yaw=False, generators_count=2, generators_random_seed=None, **kwargs): + def __init__ (self, samples_path, debug, batch_size, sort_by_yaw=False, sort_by_yaw_target_samples_path=None, sample_process_options=SampleProcessor.Options(), output_sample_types=[], add_sample_idx=False, add_pitch=False, add_yaw=False, generators_count=2, generators_random_seed=None, **kwargs): super().__init__(samples_path, debug, batch_size) self.sample_process_options = sample_process_options self.output_sample_types = output_sample_types @@ -27,8 +27,6 @@ class SampleGeneratorFace(SampleGeneratorBase): self.sample_type = SampleType.FACE_YAW_SORTED_AS_TARGET elif sort_by_yaw: self.sample_type = SampleType.FACE_YAW_SORTED - elif with_close_to_self: - self.sample_type = SampleType.FACE_WITH_CLOSE_TO_SELF else: self.sample_type = SampleType.FACE @@ -82,7 +80,7 @@ class SampleGeneratorFace(SampleGeneratorBase): if all ( [ samples[idx] == None for idx in samples_idxs] ): raise ValueError('Not enough training data. Gather more faces!') - if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF: + if self.sample_type == SampleType.FACE: shuffle_idxs = [] elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET: shuffle_idxs = [] @@ -102,12 +100,12 @@ class SampleGeneratorFace(SampleGeneratorBase): if len(repeat_samples_idxs) > 0: idx = repeat_samples_idxs.pop() - if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF: + if self.sample_type == SampleType.FACE: sample = samples[idx] elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET: sample = samples[(idx >> 16) & 0xFFFF][idx & 0xFFFF] else: - if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF: + if self.sample_type == SampleType.FACE: if len(shuffle_idxs) == 0: shuffle_idxs = samples_idxs.copy() np.random.shuffle(shuffle_idxs) diff --git a/samplelib/SampleLoader.py b/samplelib/SampleLoader.py index 8184409..742b9d0 100644 --- a/samplelib/SampleLoader.py +++ b/samplelib/SampleLoader.py @@ -1,19 +1,19 @@ +import operator import traceback from enum import IntEnum -import cv2 -import numpy as np from pathlib import Path -from utils import Path_utils -from utils.DFLPNG import DFLPNG -from utils.DFLJPG import DFLJPG +import cv2 +import numpy as np -from .Sample import Sample -from .Sample import SampleType - -from facelib import FaceType -from facelib import LandmarksProcessor +from facelib import FaceType, LandmarksProcessor from interact import interact as io +from utils import Path_utils +from utils.DFLJPG import DFLJPG +from utils.DFLPNG import DFLPNG + +from .Sample import Sample, SampleType + class SampleLoader: cache = dict() @@ -35,6 +35,10 @@ class SampleLoader: if datas[sample_type] is None: datas[sample_type] = SampleLoader.upgradeToFaceSamples( [ Sample(filename=filename) for filename in Path_utils.get_image_paths(samples_path) ] ) + # elif sample_type == SampleType.FACE_TEMPORAL_SORTED: + # if datas[sample_type] is None: + # datas[sample_type] = SampleLoader.upgradeToFaceTemporalSortedSamples( SampleLoader.load(SampleType.FACE, samples_path) ) + elif sample_type == SampleType.FACE_YAW_SORTED: if datas[sample_type] is None: datas[sample_type] = SampleLoader.upgradeToFaceYawSortedSamples( SampleLoader.load(SampleType.FACE, samples_path) ) @@ -44,10 +48,6 @@ class SampleLoader: if target_samples_path is None: raise Exception('target_samples_path is None for FACE_YAW_SORTED_AS_TARGET') datas[sample_type] = SampleLoader.upgradeToFaceYawSortedAsTargetSamples( SampleLoader.load(SampleType.FACE_YAW_SORTED, samples_path), SampleLoader.load(SampleType.FACE_YAW_SORTED, target_samples_path) ) - elif sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF: - if datas[sample_type] is None: - datas[sample_type] = SampleLoader.upgradeToFaceCloseToSelfSamples( SampleLoader.load(SampleType.FACE, samples_path) ) - return datas[sample_type] @@ -77,44 +77,20 @@ class SampleLoader: landmarks=dflimg.get_landmarks(), ie_polys=dflimg.get_ie_polys(), pitch=pitch, - yaw=yaw) ) + yaw=yaw, + source_filename=dflimg.get_source_filename(), + fanseg_mask_exist=dflimg.get_fanseg_mask() is not None, ) ) except: print ("Unable to load %s , error: %s" % (str(s_filename_path), traceback.format_exc() ) ) return sample_list - @staticmethod - def upgradeToFaceCloseToSelfSamples (samples): - yaw_samples = SampleLoader.upgradeToFaceYawSortedSamples(samples) - yaw_samples_len = len(yaw_samples) + # @staticmethod + # def upgradeToFaceTemporalSortedSamples( samples ): + # new_s = [ (s, s.source_filename) for s in samples] + # new_s = sorted(new_s, key=operator.itemgetter(1)) - sample_list = [] - for i in io.progress_bar_generator( range(yaw_samples_len), "Sorting"): - if yaw_samples[i] is not None: - for s in yaw_samples[i]: - s_t = [] - - for n in range(2000): - yaw_idx = np.clip ( i-10 +np.random.randint(20), 0, yaw_samples_len-1 ) - if yaw_samples[yaw_idx] is None: - continue - - yaw_idx_samples_len = len(yaw_samples[yaw_idx]) - - yaw_idx_sample = yaw_samples[yaw_idx][ np.random.randint(yaw_idx_samples_len) ] - if s.filename == yaw_idx_sample.filename: - continue - - s_t.append ( yaw_idx_sample ) - if len(s_t) >= 50: - break - - if len(s_t) == 0: - s_t = [s] - - sample_list.append( s.copy_and_set(close_target_list = s_t) ) - - return sample_list + # return [ s[0] for s in new_s] @staticmethod def upgradeToFaceYawSortedSamples( samples ): diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py index 10aefd7..8924d95 100644 --- a/samplelib/SampleProcessor.py +++ b/samplelib/SampleProcessor.py @@ -13,18 +13,18 @@ class SampleProcessor(object): WARPED_TRANSFORMED = 0x00000004, TRANSFORMED = 0x00000008, LANDMARKS_ARRAY = 0x00000010, #currently unused - + RANDOM_CLOSE = 0x00000020, #currently unused MORPH_TO_RANDOM_CLOSE = 0x00000040, #currently unused - + FACE_TYPE_HALF = 0x00000100, FACE_TYPE_FULL = 0x00000200, FACE_TYPE_HEAD = 0x00000400, #currently unused FACE_TYPE_AVATAR = 0x00000800, #currently unused - + FACE_MASK_FULL = 0x00001000, FACE_MASK_EYES = 0x00002000, #currently unused - + MODE_BGR = 0x00010000, #BGR MODE_G = 0x00020000, #Grayscale MODE_GGG = 0x00040000, #3xGrayscale @@ -35,7 +35,7 @@ class SampleProcessor(object): class Options(object): #motion_blur = [chance_int, range] - chance 0..100 to apply to face (not mask), and range [1..3] where 3 is highest power of motion blur - + def __init__(self, random_flip = True, normalize_tanh = False, rotation_range=[-10,10], scale_range=[-0.05, 0.05], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05], motion_blur=None ): self.random_flip = random_flip self.normalize_tanh = normalize_tanh @@ -49,11 +49,11 @@ class SampleProcessor(object): chance = np.clip(chance, 0, 100) range = [3,5,7,9][ : np.clip(range, 0, 3)+1 ] self.motion_blur = (chance, range) - + @staticmethod def process (sample, sample_process_options, output_sample_types, debug): SPTF = SampleProcessor.TypeFlags - + sample_bgr = sample.load_bgr() h,w,c = sample_bgr.shape @@ -113,7 +113,7 @@ class SampleProcessor(object): target_face_type = FaceType.HEAD elif f & SPTF.FACE_TYPE_AVATAR != 0: target_face_type = FaceType.AVATAR - + apply_motion_blur = f & SPTF.OPT_APPLY_MOTION_BLUR != 0 if img_type == 4: @@ -170,9 +170,16 @@ class SampleProcessor(object): if np.random.randint(100) < chance : dim = mb_range[ np.random.randint(len(mb_range) ) ] img = imagelib.LinearMotionBlur (img, dim, np.random.randint(180) ) - + if face_mask_type == 1: - mask = LandmarksProcessor.get_image_hull_mask (img.shape, cur_sample.landmarks, cur_sample.ie_polys) + mask = cur_sample.load_fanseg_mask() #using fanseg_mask if exist + + if mask is None: + mask = LandmarksProcessor.get_image_hull_mask (img.shape, cur_sample.landmarks) + + if cur_sample.ie_polys is not None: + cur_sample.ie_polys.overlay_mask(mask) + img = np.concatenate( (img, mask ), -1 ) elif face_mask_type == 2: mask = LandmarksProcessor.get_image_eye_mask (img.shape, cur_sample.landmarks) @@ -202,7 +209,7 @@ class SampleProcessor(object): img_mask = img[...,3:4] if f & SPTF.MODE_BGR != 0: - img = img + img = img_bgr elif f & SPTF.MODE_BGR_SHUFFLE != 0: img_bgr = np.take (img_bgr, np.random.permutation(img_bgr.shape[-1]), axis=-1) img = np.concatenate ( (img_bgr,img_mask) , -1 ) diff --git a/utils/DFLJPG.py b/utils/DFLJPG.py index 7aa31b8..a4557c3 100644 --- a/utils/DFLJPG.py +++ b/utils/DFLJPG.py @@ -1,9 +1,13 @@ -import struct import pickle +import struct + +import cv2 import numpy as np + from facelib import FaceType from imagelib import IEPolys from utils.struct_utils import * +from interact import interact as io class DFLJPG(object): def __init__(self): @@ -137,8 +141,15 @@ class DFLJPG(object): if type(chunk['data']) == bytes: inst.dfl_dict = pickle.loads(chunk['data']) - if (inst.dfl_dict is not None) and ('face_type' not in inst.dfl_dict.keys()): - inst.dfl_dict['face_type'] = FaceType.toString (FaceType.FULL) + if (inst.dfl_dict is not None): + if 'face_type' not in inst.dfl_dict: + inst.dfl_dict['face_type'] = FaceType.toString (FaceType.FULL) + + if 'fanseg_mask' in inst.dfl_dict: + fanseg_mask = inst.dfl_dict['fanseg_mask'] + if fanseg_mask is not None: + numpyarray = np.asarray( inst.dfl_dict['fanseg_mask'], dtype=np.uint8) + inst.dfl_dict['fanseg_mask'] = cv2.imdecode(numpyarray, cv2.IMREAD_UNCHANGED) if inst.dfl_dict == None: return None @@ -155,9 +166,21 @@ class DFLJPG(object): source_filename=None, source_rect=None, source_landmarks=None, - image_to_face_mat=None + image_to_face_mat=None, + fanseg_mask=None, **kwargs ): + if fanseg_mask is not None: + fanseg_mask = np.clip ( (fanseg_mask*255).astype(np.uint8), 0, 255 ) + + ret, buf = cv2.imencode( '.jpg', fanseg_mask, [int(cv2.IMWRITE_JPEG_QUALITY), 85] ) + + if ret and len(buf) < 60000: + fanseg_mask = buf + else: + io.log_err("Unable to encode fanseg_mask for %s" % (filename) ) + fanseg_mask = None + inst = DFLJPG.load_raw (filename) inst.setDFLDictData ({ 'face_type': face_type, @@ -166,7 +189,8 @@ class DFLJPG(object): 'source_filename': source_filename, 'source_rect': source_rect, 'source_landmarks': source_landmarks, - 'image_to_face_mat': image_to_face_mat + 'image_to_face_mat': image_to_face_mat, + 'fanseg_mask' : fanseg_mask, }) try: @@ -181,7 +205,8 @@ class DFLJPG(object): source_filename=None, source_rect=None, source_landmarks=None, - image_to_face_mat=None + image_to_face_mat=None, + fanseg_mask=None, **kwargs ): if face_type is None: face_type = self.get_face_type() if landmarks is None: landmarks = self.get_landmarks() @@ -190,14 +215,18 @@ class DFLJPG(object): if source_rect is None: source_rect = self.get_source_rect() if source_landmarks is None: source_landmarks = self.get_source_landmarks() if image_to_face_mat is None: image_to_face_mat = self.get_image_to_face_mat() + if fanseg_mask is None: fanseg_mask = self.get_fanseg_mask() DFLJPG.embed_data (filename, face_type=face_type, landmarks=landmarks, ie_polys=ie_polys, source_filename=source_filename, source_rect=source_rect, source_landmarks=source_landmarks, - image_to_face_mat=image_to_face_mat) - + image_to_face_mat=image_to_face_mat, + fanseg_mask=fanseg_mask) + def remove_fanseg_mask(self): + self.dfl_dict['fanseg_mask'] = None + def dump(self): data = b"" @@ -252,8 +281,13 @@ class DFLJPG(object): def get_source_filename(self): return self.dfl_dict['source_filename'] def get_source_rect(self): return self.dfl_dict['source_rect'] def get_source_landmarks(self): return np.array ( self.dfl_dict['source_landmarks'] ) - def get_image_to_face_mat(self): + def get_image_to_face_mat(self): mat = self.dfl_dict.get ('image_to_face_mat', None) if mat is not None: return np.array (mat) - return None \ No newline at end of file + return None + def get_fanseg_mask(self): + fanseg_mask = self.dfl_dict.get ('fanseg_mask', None) + if fanseg_mask is not None: + return np.clip ( np.array (fanseg_mask) / 255.0, 0.0, 1.0 )[...,np.newaxis] + return None diff --git a/utils/DFLPNG.py b/utils/DFLPNG.py index 79ec9c1..a4c668c 100644 --- a/utils/DFLPNG.py +++ b/utils/DFLPNG.py @@ -1,13 +1,16 @@ -PNG_HEADER = b"\x89PNG\r\n\x1a\n" - +import pickle import string import struct import zlib -import pickle + +import cv2 import numpy as np + from facelib import FaceType from imagelib import IEPolys +PNG_HEADER = b"\x89PNG\r\n\x1a\n" + class Chunk(object): def __init__(self, name=None, data=None): self.length = 0 @@ -219,7 +222,7 @@ class DFLPNG(object): self.data = b"" self.length = 0 self.chunks = [] - self.fcwp_dict = None + self.dfl_dict = None @staticmethod def load_raw(filename): @@ -252,12 +255,19 @@ class DFLPNG(object): def load(filename): try: inst = DFLPNG.load_raw (filename) - inst.fcwp_dict = inst.getDFLDictData() + inst.dfl_dict = inst.getDFLDictData() - if (inst.fcwp_dict is not None) and ('face_type' not in inst.fcwp_dict.keys()): - inst.fcwp_dict['face_type'] = FaceType.toString (FaceType.FULL) + if inst.dfl_dict is not None: + if 'face_type' not in inst.dfl_dict: + inst.dfl_dict['face_type'] = FaceType.toString (FaceType.FULL) - if inst.fcwp_dict == None: + if 'fanseg_mask' in inst.dfl_dict: + fanseg_mask = inst.dfl_dict['fanseg_mask'] + if fanseg_mask is not None: + numpyarray = np.asarray( inst.dfl_dict['fanseg_mask'], dtype=np.uint8) + inst.dfl_dict['fanseg_mask'] = cv2.imdecode(numpyarray, cv2.IMREAD_UNCHANGED) + + if inst.dfl_dict == None: return None return inst @@ -272,9 +282,21 @@ class DFLPNG(object): source_filename=None, source_rect=None, source_landmarks=None, - image_to_face_mat=None + image_to_face_mat=None, + fanseg_mask=None, **kwargs ): + if fanseg_mask is not None: + fanseg_mask = np.clip ( (fanseg_mask*255).astype(np.uint8), 0, 255 ) + + ret, buf = cv2.imencode( '.jpg', fanseg_mask, [int(cv2.IMWRITE_JPEG_QUALITY), 85] ) + + if ret and len(buf) < 60000: + fanseg_mask = buf + else: + io.log_err("Unable to encode fanseg_mask for %s" % (filename) ) + fanseg_mask = None + inst = DFLPNG.load_raw (filename) inst.setDFLDictData ({ 'face_type': face_type, @@ -283,7 +305,8 @@ class DFLPNG(object): 'source_filename': source_filename, 'source_rect': source_rect, 'source_landmarks': source_landmarks, - 'image_to_face_mat':image_to_face_mat + 'image_to_face_mat':image_to_face_mat, + 'fanseg_mask' : fanseg_mask, }) try: @@ -292,13 +315,14 @@ class DFLPNG(object): except: raise Exception( 'cannot save %s' % (filename) ) - def embed_and_set(self, filename, face_type=None, - landmarks=None, - ie_polys=None, - source_filename=None, - source_rect=None, - source_landmarks=None, - image_to_face_mat=None + def embed_and_set(self, filename, face_type=None, + landmarks=None, + ie_polys=None, + source_filename=None, + source_rect=None, + source_landmarks=None, + image_to_face_mat=None, + fanseg_mask=None, **kwargs ): if face_type is None: face_type = self.get_face_type() if landmarks is None: landmarks = self.get_landmarks() @@ -307,13 +331,18 @@ class DFLPNG(object): if source_rect is None: source_rect = self.get_source_rect() if source_landmarks is None: source_landmarks = self.get_source_landmarks() if image_to_face_mat is None: image_to_face_mat = self.get_image_to_face_mat() + if fanseg_mask is None: fanseg_mask = self.get_fanseg_mask() DFLPNG.embed_data (filename, face_type=face_type, landmarks=landmarks, ie_polys=ie_polys, source_filename=source_filename, source_rect=source_rect, source_landmarks=source_landmarks, - image_to_face_mat=image_to_face_mat) + image_to_face_mat=image_to_face_mat, + fanseg_mask=fanseg_mask) + + def remove_fanseg_mask(self): + self.dfl_dict['fanseg_mask'] = None def dump(self): data = PNG_HEADER @@ -352,17 +381,21 @@ class DFLPNG(object): chunk = DFLChunk(dict_data) self.chunks.insert(-1, chunk) - def get_face_type(self): return self.fcwp_dict['face_type'] - def get_landmarks(self): return np.array ( self.fcwp_dict['landmarks'] ) - def get_ie_polys(self): return IEPolys.load(self.fcwp_dict.get('ie_polys',None)) - def get_source_filename(self): return self.fcwp_dict['source_filename'] - def get_source_rect(self): return self.fcwp_dict['source_rect'] - def get_source_landmarks(self): return np.array ( self.fcwp_dict['source_landmarks'] ) + def get_face_type(self): return self.dfl_dict['face_type'] + def get_landmarks(self): return np.array ( self.dfl_dict['landmarks'] ) + def get_ie_polys(self): return IEPolys.load(self.dfl_dict.get('ie_polys',None)) + def get_source_filename(self): return self.dfl_dict['source_filename'] + def get_source_rect(self): return self.dfl_dict['source_rect'] + def get_source_landmarks(self): return np.array ( self.dfl_dict['source_landmarks'] ) def get_image_to_face_mat(self): - mat = self.fcwp_dict.get ('image_to_face_mat', None) + mat = self.dfl_dict.get ('image_to_face_mat', None) if mat is not None: return np.array (mat) return None - + def get_fanseg_mask(self): + fanseg_mask = self.dfl_dict.get ('fanseg_mask', None) + if fanseg_mask is not None: + return np.clip ( np.array (fanseg_mask) / 255.0, 0.0, 1.0 )[...,np.newaxis] + return None def __str__(self): return "".format(len(self.chunks), **self.__dict__)