diff --git a/imagelib/warp.py b/imagelib/warp.py index 0bd5e3a..a3c9490 100644 --- a/imagelib/warp.py +++ b/imagelib/warp.py @@ -2,18 +2,24 @@ import numpy as np import cv2 from utils import random_utils -def gen_warp_params (source, flip, rotation_range=[-10,10], scale_range=[-0.5, 0.5], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05] ): +def gen_warp_params (source, flip, rotation_range=[-10,10], scale_range=[-0.5, 0.5], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05], rnd_seed=None ): h,w,c = source.shape if (h != w): raise ValueError ('gen_warp_params accepts only square images.') + + if rnd_seed != None: + rnd_state = np.random.RandomState (rnd_seed) + else: + rnd_state = np.random - rotation = np.random.uniform( rotation_range[0], rotation_range[1] ) - scale = np.random.uniform(1 +scale_range[0], 1 +scale_range[1]) - tx = np.random.uniform( tx_range[0], tx_range[1] ) - ty = np.random.uniform( ty_range[0], ty_range[1] ) - + rotation = rnd_state.uniform( rotation_range[0], rotation_range[1] ) + scale = rnd_state.uniform(1 +scale_range[0], 1 +scale_range[1]) + tx = rnd_state.uniform( tx_range[0], tx_range[1] ) + ty = rnd_state.uniform( ty_range[0], ty_range[1] ) + p_flip = flip and rnd_state.randint(10) < 4 + #random warp by grid - cell_size = [ w // (2**i) for i in range(1,4) ] [ np.random.randint(3) ] + cell_size = [ w // (2**i) for i in range(1,4) ] [ rnd_state.randint(3) ] cell_count = w // cell_size + 1 grid_points = np.linspace( 0, w, cell_count) @@ -37,7 +43,7 @@ def gen_warp_params (source, flip, rotation_range=[-10,10], scale_range=[-0.5, 0 params['mapy'] = mapy params['rmat'] = random_transform_mat params['w'] = w - params['flip'] = flip and np.random.randint(10) < 4 + params['flip'] = p_flip return params diff --git a/models/Model_DEV_FANSEG/Model.py b/models/Model_DEV_FANSEG/Model.py index f66baaf..ae70865 100644 --- a/models/Model_DEV_FANSEG/Model.py +++ b/models/Model_DEV_FANSEG/Model.py @@ -47,13 +47,13 @@ class Model(ModelBase): 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), - output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR), 'resolution' : self.resolution, 'motion_blur':(25, 5), 'gaussian_blur':(25,5), 'border_replicate':False, 'random_hsv_shift' : True }, + output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR_RANDOM_HSV_SHIFT), 'resolution' : self.resolution, 'motion_blur':(25, 5), 'gaussian_blur':(25,5), 'border_replicate':False}, { 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_M), 'resolution': 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 ), - output_sample_types=[ { 'types': (t.IMG_TRANSFORMED , face_type, t.MODE_BGR), 'resolution' : self.resolution, 'random_hsv_shift' : True}, + output_sample_types=[ { 'types': (t.IMG_TRANSFORMED , face_type, t.MODE_BGR_RANDOM_HSV_SHIFT), 'resolution' : self.resolution}, ]) ]) diff --git a/models/Model_DEV_FUNIT/Model.py b/models/Model_DEV_FUNIT/Model.py index 1001219..0667c44 100644 --- a/models/Model_DEV_FUNIT/Model.py +++ b/models/Model_DEV_FUNIT/Model.py @@ -23,9 +23,9 @@ class FUNITModel(ModelBase): #override def onInitializeOptions(self, is_first_run, ask_override): - default_resolution = 96 + default_resolution = 64 if is_first_run: - self.options['resolution'] = io.input_int(f"Resolution ( 96,128,224 ?:help skip:{default_resolution}) : ", default_resolution, [128,224]) + self.options['resolution'] = io.input_int(f"Resolution ( 64,96,128,224 ?:help skip:{default_resolution}) : ", default_resolution, [64,96,128,224]) else: self.options['resolution'] = self.options.get('resolution', default_resolution) @@ -48,7 +48,7 @@ class FUNITModel(ModelBase): resolution = self.options['resolution'] face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF - person_id_max_count = SampleGeneratorFace.get_person_id_max_count(self.training_data_src_path) + person_id_max_count = SampleGeneratorFacePerson.get_person_id_max_count(self.training_data_src_path) self.model = FUNIT( face_type_str=FaceType.toString(face_type), @@ -85,21 +85,21 @@ class FUNITModel(ModelBase): output_sample_types1=[ {'types': (t.IMG_SOURCE, face_type, t.MODE_BGR), 'resolution':resolution, 'normalize_tanh':True} ] self.set_training_data_generators ([ - SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, + SampleGeneratorFacePerson(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, sample_process_options=SampleProcessor.Options(random_flip=True, rotation_range=[0,0] ), - output_sample_types=output_sample_types, person_id_mode=True ), + output_sample_types=output_sample_types, person_id_mode=1, use_caching=True, generators_count=1 ), - SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, + SampleGeneratorFacePerson(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, sample_process_options=SampleProcessor.Options(random_flip=True, rotation_range=[0,0] ), - output_sample_types=output_sample_types, person_id_mode=True ), + output_sample_types=output_sample_types, person_id_mode=1, use_caching=True, generators_count=1 ), - SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, + SampleGeneratorFacePerson(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, sample_process_options=SampleProcessor.Options(random_flip=True, rotation_range=[0,0]), - output_sample_types=output_sample_types1, person_id_mode=True ), + output_sample_types=output_sample_types1, person_id_mode=1, use_caching=True, generators_count=1 ), - SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, + SampleGeneratorFacePerson(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, sample_process_options=SampleProcessor.Options(random_flip=True, rotation_range=[0,0]), - output_sample_types=output_sample_types1, person_id_mode=True ), + output_sample_types=output_sample_types1, person_id_mode=1, use_caching=True, generators_count=1 ), ]) #override @@ -125,7 +125,7 @@ class FUNITModel(ModelBase): xb = generators_samples[1][0] ta = generators_samples[2][0] tb = generators_samples[3][0] - + view_samples = min(4, xa.shape[0]) lines_train = [] diff --git a/nnlib/nnlib.py b/nnlib/nnlib.py index 39a16f1..c5af230 100644 --- a/nnlib/nnlib.py +++ b/nnlib/nnlib.py @@ -93,6 +93,7 @@ Model = keras.models.Model Adam = nnlib.Adam RMSprop = nnlib.RMSprop LookaheadOptimizer = nnlib.LookaheadOptimizer +SGD = nnlib.keras.optimizers.SGD modelify = nnlib.modelify gaussian_blur = nnlib.gaussian_blur @@ -765,9 +766,10 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator 2 - allows to train x3 bigger network on same VRAM consuming RAM*2 and CPU power. """ - def __init__(self, learning_rate=0.001, rho=0.9, tf_cpu_mode=0, **kwargs): + def __init__(self, learning_rate=0.001, rho=0.9, lr_dropout=0, tf_cpu_mode=0, **kwargs): self.initial_decay = kwargs.pop('decay', 0.0) self.epsilon = kwargs.pop('epsilon', K.epsilon()) + self.lr_dropout = lr_dropout self.tf_cpu_mode = tf_cpu_mode learning_rate = kwargs.pop('lr', learning_rate) @@ -788,6 +790,8 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator dtype=K.dtype(p), name='accumulator_' + str(i)) for (i, p) in enumerate(params)] + if self.lr_dropout != 0: + lr_rnds = [ K.random_binomial(K.int_shape(p), p=self.lr_dropout, dtype=K.dtype(p)) for p in params ] if e: e.__exit__(None, None, None) self.weights = [self.iterations] + accumulators @@ -798,12 +802,15 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator lr = lr * (1. / (1. + self.decay * K.cast(self.iterations, K.dtype(self.decay)))) - for p, g, a in zip(params, grads, accumulators): + for i, (p, g, a) in enumerate(zip(params, grads, accumulators)): # update accumulator e = K.tf.device("/cpu:0") if self.tf_cpu_mode == 2 else None if e: e.__enter__() new_a = self.rho * a + (1. - self.rho) * K.square(g) - new_p = p - lr * g / (K.sqrt(new_a) + self.epsilon) + p_diff = - lr * g / (K.sqrt(new_a) + self.epsilon) + if self.lr_dropout != 0: + p_diff *= lr_rnds[i] + new_p = p + p_diff if e: e.__exit__(None, None, None) self.updates.append(K.update(a, new_a)) @@ -828,7 +835,8 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator config = {'learning_rate': float(K.get_value(self.learning_rate)), 'rho': float(K.get_value(self.rho)), 'decay': float(K.get_value(self.decay)), - 'epsilon': self.epsilon} + 'epsilon': self.epsilon, + 'lr_dropout' : self.lr_dropout } base_config = super(RMSprop, self).get_config() return dict(list(base_config.items()) + list(config.items())) nnlib.RMSprop = RMSprop @@ -847,6 +855,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator amsgrad: boolean. Whether to apply the AMSGrad variant of this algorithm from the paper "On the Convergence of Adam and Beyond". + lr_dropout: float [0.0 .. 1.0] Learning rate dropout https://arxiv.org/pdf/1912.00144 tf_cpu_mode: only for tensorflow backend 0 - default, no changes. 1 - allows to train x2 bigger network on same VRAM consuming RAM @@ -860,7 +869,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator """ def __init__(self, lr=0.001, beta_1=0.9, beta_2=0.999, - epsilon=None, decay=0., amsgrad=False, tf_cpu_mode=0, **kwargs): + epsilon=None, decay=0., amsgrad=False, lr_dropout=0, tf_cpu_mode=0, **kwargs): super(Adam, self).__init__(**kwargs) with K.name_scope(self.__class__.__name__): self.iterations = K.variable(0, dtype='int64', name='iterations') @@ -873,6 +882,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator self.epsilon = epsilon self.initial_decay = decay self.amsgrad = amsgrad + self.lr_dropout = lr_dropout self.tf_cpu_mode = tf_cpu_mode def get_updates(self, loss, params): @@ -896,11 +906,16 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator vhats = [K.zeros(K.int_shape(p), dtype=K.dtype(p)) for p in params] else: vhats = [K.zeros(1) for _ in params] + + + if self.lr_dropout != 0: + lr_rnds = [ K.random_binomial(K.int_shape(p), p=self.lr_dropout, dtype=K.dtype(p)) for p in params ] + if e: e.__exit__(None, None, None) self.weights = [self.iterations] + ms + vs + vhats - for p, g, m, v, vhat in zip(params, grads, ms, vs, vhats): + for i, (p, g, m, v, vhat) in enumerate( zip(params, grads, ms, vs, vhats) ): e = K.tf.device("/cpu:0") if self.tf_cpu_mode == 2 else None if e: e.__enter__() m_t = (self.beta_1 * m) + (1. - self.beta_1) * g @@ -912,13 +927,16 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator if e: e.__exit__(None, None, None) if self.amsgrad: - p_t = p - lr_t * m_t / (K.sqrt(vhat_t) + self.epsilon) + p_diff = - lr_t * m_t / (K.sqrt(vhat_t) + self.epsilon) else: - p_t = p - lr_t * m_t / (K.sqrt(v_t) + self.epsilon) + p_diff = - lr_t * m_t / (K.sqrt(v_t) + self.epsilon) + + if self.lr_dropout != 0: + p_diff *= lr_rnds[i] self.updates.append(K.update(m, m_t)) self.updates.append(K.update(v, v_t)) - new_p = p_t + new_p = p + p_diff # Apply constraints. if getattr(p, 'constraint', None) is not None: @@ -933,17 +951,18 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator 'beta_2': float(K.get_value(self.beta_2)), 'decay': float(K.get_value(self.decay)), 'epsilon': self.epsilon, - 'amsgrad': self.amsgrad} + 'amsgrad': self.amsgrad, + 'lr_dropout' : self.lr_dropout} base_config = super(Adam, self).get_config() return dict(list(base_config.items()) + list(config.items())) nnlib.Adam = Adam - + class LookaheadOptimizer(keras.optimizers.Optimizer): def __init__(self, optimizer, sync_period=5, slow_step=0.5, tf_cpu_mode=0, **kwargs): super(LookaheadOptimizer, self).__init__(**kwargs) self.optimizer = optimizer self.tf_cpu_mode = tf_cpu_mode - + with K.name_scope(self.__class__.__name__): self.sync_period = K.variable(sync_period, dtype='int64', name='sync_period') self.slow_step = K.variable(slow_step, name='slow_step') @@ -975,17 +994,17 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator if e: e.__enter__() slow_params = [K.variable(K.get_value(p), name='sp_{}'.format(i)) for i, p in enumerate(params)] if e: e.__exit__(None, None, None) - - + + self.updates = self.optimizer.get_updates(loss, params) slow_updates = [] for p, sp in zip(params, slow_params): - + e = K.tf.device("/cpu:0") if self.tf_cpu_mode == 2 else None if e: e.__enter__() sp_t = sp + self.slow_step * (p - sp) if e: e.__exit__(None, None, None) - + slow_updates.append(K.update(sp, K.switch( sync_cond, sp_t, @@ -996,7 +1015,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator sp_t - p, K.zeros_like(p), ))) - + self.updates += slow_updates self.weights = self.optimizer.weights + slow_params return self.updates @@ -1015,7 +1034,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator optimizer = keras.optimizers.deserialize(config.pop('optimizer')) return cls(optimizer, **config) nnlib.LookaheadOptimizer = LookaheadOptimizer - + class DenseMaxout(keras.layers.Layer): """A dense maxout layer. A `MaxoutDense` layer takes the element-wise maximum of diff --git a/samplelib/SampleGeneratorFace.py b/samplelib/SampleGeneratorFace.py index 81cca20..715c27c 100644 --- a/samplelib/SampleGeneratorFace.py +++ b/samplelib/SampleGeneratorFace.py @@ -55,7 +55,6 @@ class SampleGeneratorFace(SampleGeneratorBase): raise ValueError('No training data provided.') ct_samples = SampleLoader.load (SampleType.FACE, random_ct_samples_path) if random_ct_samples_path is not None else None - self.random_ct_sample_chance = 100 if self.debug: self.generators_count = 1 @@ -133,16 +132,12 @@ class SampleGeneratorFace(SampleGeneratorBase): try: ct_sample=None if ct_samples is not None: - if np.random.randint(100) < self.random_ct_sample_chance: - ct_sample=ct_samples[np.random.randint(ct_samples_len)] + ct_sample=ct_samples[np.random.randint(ct_samples_len)] - x = SampleProcessor.process (sample, self.sample_process_options, self.output_sample_types, self.debug, ct_sample=ct_sample) + x, = SampleProcessor.process ([sample], self.sample_process_options, self.output_sample_types, self.debug, ct_sample=ct_sample) except: raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) ) - if type(x) != tuple and type(x) != list: - raise Exception('SampleProcessor.process returns NOT tuple/list') - if batches is None: batches = [ [] for _ in range(len(x)) ] if self.add_sample_idx: diff --git a/samplelib/SampleGeneratorFacePerson.py b/samplelib/SampleGeneratorFacePerson.py index aa77feb..95d929c 100644 --- a/samplelib/SampleGeneratorFacePerson.py +++ b/samplelib/SampleGeneratorFacePerson.py @@ -1,3 +1,4 @@ +import copy import multiprocessing import traceback @@ -37,7 +38,10 @@ class SampleGeneratorFacePerson(SampleGeneratorBase): self.generators_random_seed = generators_random_seed samples = SampleLoader.load (SampleType.FACE, self.samples_path, person_id_mode=True, use_caching=use_caching) - + samples = copy.copy(samples) + for i in range(len(samples)): + samples[i] = copy.copy(samples[i]) + if person_id_mode==1: np.random.shuffle(samples) @@ -52,6 +56,7 @@ class SampleGeneratorFacePerson(SampleGeneratorBase): if len(sample) == 0: samples.pop(i) samples = new_samples + #new_samples = [] #for s in samples: # new_samples += s @@ -114,7 +119,19 @@ class SampleGeneratorFacePerson(SampleGeneratorBase): for i in range(persons_count): samples_idxs[i] = [*range(len(samples[i]))] shuffle_idxs[i] = [] - + elif self.person_id_mode==3: + persons_count = len(samples) + + person_idxs = [ *range(persons_count) ] + shuffle_person_idxs = [] + + samples_idxs = [None]*persons_count + shuffle_idxs = [None]*persons_count + + for i in range(persons_count): + samples_idxs[i] = [*range(len(samples[i]))] + shuffle_idxs[i] = [] + while True: if self.person_id_mode==2: @@ -122,7 +139,7 @@ class SampleGeneratorFacePerson(SampleGeneratorBase): shuffle_person_idxs = person_idxs.copy() np.random.shuffle(shuffle_person_idxs) person_ids = shuffle_person_idxs.pop() - + batches = None for n_batch in range(self.batch_size): @@ -130,13 +147,13 @@ class SampleGeneratorFacePerson(SampleGeneratorBase): if self.person_id_mode==1: if len(shuffle_idxs) == 0: shuffle_idxs = samples_idxs.copy() - #np.random.shuffle(shuffle_idxs) + np.random.shuffle(shuffle_idxs) ### idx = shuffle_idxs.pop() sample = samples[ idx ] try: - x = SampleProcessor.process (sample, self.sample_process_options, self.output_sample_types, self.debug) + x, = SampleProcessor.process ([sample], self.sample_process_options, self.output_sample_types, self.debug) except: raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) ) @@ -155,7 +172,7 @@ class SampleGeneratorFacePerson(SampleGeneratorBase): batches[i_person_id].append ( np.array([sample.person_id]) ) - else: + elif self.person_id_mode==2: person_id1, person_id2 = person_ids if len(shuffle_idxs[person_id1]) == 0: @@ -174,12 +191,12 @@ class SampleGeneratorFacePerson(SampleGeneratorBase): if sample1 is not None and sample2 is not None: try: - x1 = SampleProcessor.process (sample1, self.sample_process_options, self.output_sample_types, self.debug) + x1, = SampleProcessor.process ([sample1], self.sample_process_options, self.output_sample_types, self.debug) except: raise Exception ("Exception occured in sample %s. Error: %s" % (sample1.filename, traceback.format_exc() ) ) try: - x2 = SampleProcessor.process (sample2, self.sample_process_options, self.output_sample_types, self.debug) + x2, = SampleProcessor.process ([sample2], self.sample_process_options, self.output_sample_types, self.debug) except: raise Exception ("Exception occured in sample %s. Error: %s" % (sample2.filename, traceback.format_exc() ) ) @@ -203,10 +220,59 @@ class SampleGeneratorFacePerson(SampleGeneratorBase): batches[i_person_id2].append ( np.array([sample2.person_id]) ) + elif self.person_id_mode==3: + if len(shuffle_person_idxs) == 0: + shuffle_person_idxs = person_idxs.copy() + np.random.shuffle(shuffle_person_idxs) + person_id = shuffle_person_idxs.pop() + + if len(shuffle_idxs[person_id]) == 0: + shuffle_idxs[person_id] = samples_idxs[person_id].copy() + np.random.shuffle(shuffle_idxs[person_id]) + + idx = shuffle_idxs[person_id].pop() + sample1 = samples[person_id][idx] + if len(shuffle_idxs[person_id]) == 0: + shuffle_idxs[person_id] = samples_idxs[person_id].copy() + np.random.shuffle(shuffle_idxs[person_id]) + + idx = shuffle_idxs[person_id].pop() + sample2 = samples[person_id][idx] + + if sample1 is not None and sample2 is not None: + try: + x1, = SampleProcessor.process ([sample1], self.sample_process_options, self.output_sample_types, self.debug) + except: + raise Exception ("Exception occured in sample %s. Error: %s" % (sample1.filename, traceback.format_exc() ) ) + + try: + x2, = SampleProcessor.process ([sample2], self.sample_process_options, self.output_sample_types, self.debug) + except: + raise Exception ("Exception occured in sample %s. Error: %s" % (sample2.filename, traceback.format_exc() ) ) + + x1_len = len(x1) + if batches is None: + batches = [ [] for _ in range(x1_len) ] + batches += [ [] ] + i_person_id1 = len(batches)-1 + + batches += [ [] for _ in range(len(x2)) ] + batches += [ [] ] + i_person_id2 = len(batches)-1 + + for i in range(x1_len): + batches[i].append ( x1[i] ) + + for i in range(len(x2)): + batches[x1_len+1+i].append ( x2[i] ) + + batches[i_person_id1].append ( np.array([sample1.person_id]) ) + + batches[i_person_id2].append ( np.array([sample2.person_id]) ) yield [ np.array(batch) for batch in batches] @staticmethod def get_person_id_max_count(samples_path): - return SampleLoader.get_person_id_max_count(samples_path) \ No newline at end of file + return SampleLoader.get_person_id_max_count(samples_path) diff --git a/samplelib/SampleGeneratorFaceTemporal.py b/samplelib/SampleGeneratorFaceTemporal.py index add15ae..10a5856 100644 --- a/samplelib/SampleGeneratorFaceTemporal.py +++ b/samplelib/SampleGeneratorFaceTemporal.py @@ -71,7 +71,7 @@ class SampleGeneratorFaceTemporal(SampleGeneratorBase): for i in range( self.temporal_image_count ): sample = samples[ idx+i*mult ] try: - temporal_samples += SampleProcessor.process (sample, self.sample_process_options, self.output_sample_types, self.debug) + temporal_samples += SampleProcessor.process ([sample], self.sample_process_options, self.output_sample_types, self.debug)[0] except: raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) ) diff --git a/samplelib/SampleGeneratorImageTemporal.py b/samplelib/SampleGeneratorImageTemporal.py index 227c09e..a5dd2c5 100644 --- a/samplelib/SampleGeneratorImageTemporal.py +++ b/samplelib/SampleGeneratorImageTemporal.py @@ -66,7 +66,7 @@ class SampleGeneratorImageTemporal(SampleGeneratorBase): for i in range( self.temporal_image_count ): sample = samples[ idx+i*mult ] try: - temporal_samples += SampleProcessor.process (sample, self.sample_process_options, self.output_sample_types, self.debug) + temporal_samples += SampleProcessor.process ([sample], self.sample_process_options, self.output_sample_types, self.debug)[0] except: raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) ) diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py index 3ba0f80..fa9d13f 100644 --- a/samplelib/SampleProcessor.py +++ b/samplelib/SampleProcessor.py @@ -92,231 +92,218 @@ class SampleProcessor(object): } @staticmethod - def process (sample, sample_process_options, output_sample_types, debug, ct_sample=None): + def process (samples, sample_process_options, output_sample_types, debug, ct_sample=None): SPTF = SampleProcessor.Types - - sample_bgr = sample.load_bgr() - ct_sample_bgr = None - ct_sample_mask = None - h,w,c = sample_bgr.shape - - is_face_sample = sample.landmarks is not None - - if debug and is_face_sample: - LandmarksProcessor.draw_landmarks (sample_bgr, sample.landmarks, (0, 1, 0)) - - params = imagelib.gen_warp_params(sample_bgr, sample_process_options.random_flip, rotation_range=sample_process_options.rotation_range, scale_range=sample_process_options.scale_range, tx_range=sample_process_options.tx_range, ty_range=sample_process_options.ty_range ) - - cached_images = collections.defaultdict(dict) - + sample_rnd_seed = np.random.randint(0x80000000) - + outputs = [] - for opts in output_sample_types: + for sample in samples: + sample_bgr = sample.load_bgr() + ct_sample_bgr = None + ct_sample_mask = None + h,w,c = sample_bgr.shape - resolution = opts.get('resolution', 0) - types = opts.get('types', [] ) + is_face_sample = sample.landmarks is not None - border_replicate = opts.get('border_replicate', True) - random_sub_res = opts.get('random_sub_res', 0) - normalize_std_dev = opts.get('normalize_std_dev', False) - normalize_vgg = opts.get('normalize_vgg', False) - motion_blur = opts.get('motion_blur', None) - gaussian_blur = opts.get('gaussian_blur', None) - - random_hsv_shift = opts.get('random_hsv_shift', None) - ct_mode = opts.get('ct_mode', 'None') - normalize_tanh = opts.get('normalize_tanh', False) + if debug and is_face_sample: + LandmarksProcessor.draw_landmarks (sample_bgr, sample.landmarks, (0, 1, 0)) - img_type = SPTF.NONE - target_face_type = SPTF.NONE - face_mask_type = SPTF.NONE - mode_type = SPTF.NONE - for t in types: - if t >= SPTF.IMG_TYPE_BEGIN and t < SPTF.IMG_TYPE_END: - img_type = t - elif t >= SPTF.FACE_TYPE_BEGIN and t < SPTF.FACE_TYPE_END: - target_face_type = t - elif t >= SPTF.MODE_BEGIN and t < SPTF.MODE_END: - mode_type = t + params = imagelib.gen_warp_params(sample_bgr, sample_process_options.random_flip, rotation_range=sample_process_options.rotation_range, scale_range=sample_process_options.scale_range, tx_range=sample_process_options.tx_range, ty_range=sample_process_options.ty_range, rnd_seed=sample_rnd_seed ) - if img_type == SPTF.NONE: - raise ValueError ('expected IMG_ type') + outputs_sample = [] + for opts in output_sample_types: - if img_type == SPTF.IMG_LANDMARKS_ARRAY: - l = sample.landmarks - l = np.concatenate ( [ np.expand_dims(l[:,0] / w,-1), np.expand_dims(l[:,1] / h,-1) ], -1 ) - l = np.clip(l, 0.0, 1.0) - img = l - elif img_type == SPTF.IMG_PITCH_YAW_ROLL or img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID: - pitch_yaw_roll = sample.pitch_yaw_roll - if pitch_yaw_roll is not None: - pitch, yaw, roll = pitch_yaw_roll + resolution = opts.get('resolution', 0) + types = opts.get('types', [] ) + + border_replicate = opts.get('border_replicate', True) + random_sub_res = opts.get('random_sub_res', 0) + normalize_std_dev = opts.get('normalize_std_dev', False) + normalize_vgg = opts.get('normalize_vgg', False) + motion_blur = opts.get('motion_blur', None) + gaussian_blur = opts.get('gaussian_blur', None) + + ct_mode = opts.get('ct_mode', 'None') + normalize_tanh = opts.get('normalize_tanh', False) + + img_type = SPTF.NONE + target_face_type = SPTF.NONE + face_mask_type = SPTF.NONE + mode_type = SPTF.NONE + for t in types: + if t >= SPTF.IMG_TYPE_BEGIN and t < SPTF.IMG_TYPE_END: + img_type = t + elif t >= SPTF.FACE_TYPE_BEGIN and t < SPTF.FACE_TYPE_END: + target_face_type = t + elif t >= SPTF.MODE_BEGIN and t < SPTF.MODE_END: + mode_type = t + + if img_type == SPTF.NONE: + raise ValueError ('expected IMG_ type') + + if img_type == SPTF.IMG_LANDMARKS_ARRAY: + l = sample.landmarks + l = np.concatenate ( [ np.expand_dims(l[:,0] / w,-1), np.expand_dims(l[:,1] / h,-1) ], -1 ) + l = np.clip(l, 0.0, 1.0) + img = l + elif img_type == SPTF.IMG_PITCH_YAW_ROLL or img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID: + pitch_yaw_roll = sample.pitch_yaw_roll + if pitch_yaw_roll is not None: + pitch, yaw, roll = pitch_yaw_roll + else: + pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll (sample.landmarks) + if params['flip']: + yaw = -yaw + + if img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID: + pitch = (pitch+1.0) / 2.0 + yaw = (yaw+1.0) / 2.0 + roll = (roll+1.0) / 2.0 + + img = (pitch, yaw, roll) else: - pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll (sample.landmarks) - if params['flip']: - yaw = -yaw + if mode_type == SPTF.NONE: + raise ValueError ('expected MODE_ type') - if img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID: - pitch = (pitch+1.0) / 2.0 - yaw = (yaw+1.0) / 2.0 - roll = (roll+1.0) / 2.0 + def do_transform(img, mask): + warp = (img_type==SPTF.IMG_WARPED or img_type==SPTF.IMG_WARPED_TRANSFORMED) + transform = (img_type==SPTF.IMG_WARPED_TRANSFORMED or img_type==SPTF.IMG_TRANSFORMED) + flip = img_type != SPTF.IMG_WARPED - img = (pitch, yaw, roll) - else: - if mode_type == SPTF.NONE: - raise ValueError ('expected MODE_ type') + img = imagelib.warp_by_params (params, img, warp, transform, flip, border_replicate) + if mask is not None: + mask = imagelib.warp_by_params (params, mask, warp, transform, flip, False) + if len(mask.shape) == 2: + mask = mask[...,np.newaxis] - def do_transform(img, mask): - warp = (img_type==SPTF.IMG_WARPED or img_type==SPTF.IMG_WARPED_TRANSFORMED) - transform = (img_type==SPTF.IMG_WARPED_TRANSFORMED or img_type==SPTF.IMG_TRANSFORMED) - flip = img_type != SPTF.IMG_WARPED + img = np.concatenate( (img, mask ), -1 ) + return img - img = imagelib.warp_by_params (params, img, warp, transform, flip, border_replicate) - if mask is not None: - mask = imagelib.warp_by_params (params, mask, warp, transform, flip, False) - if len(mask.shape) == 2: - mask = mask[...,np.newaxis] + img = sample_bgr - img = np.concatenate( (img, mask ), -1 ) - return img + ### Prepare a mask + mask = None + if is_face_sample: + mask = sample.load_fanseg_mask() #using fanseg_mask if exist - img = sample_bgr + if mask is None: + if sample.eyebrows_expand_mod is not None: + mask = LandmarksProcessor.get_image_hull_mask (img.shape, sample.landmarks, eyebrows_expand_mod=sample.eyebrows_expand_mod ) + else: + mask = LandmarksProcessor.get_image_hull_mask (img.shape, sample.landmarks) - ### Prepare a mask - mask = None - if is_face_sample: - mask = sample.load_fanseg_mask() #using fanseg_mask if exist + if sample.ie_polys is not None: + sample.ie_polys.overlay_mask(mask) + ################## - if mask is None: - if sample.eyebrows_expand_mod is not None: - mask = LandmarksProcessor.get_image_hull_mask (img.shape, sample.landmarks, eyebrows_expand_mod=sample.eyebrows_expand_mod ) + + if motion_blur is not None: + chance, mb_max_size = motion_blur + chance = np.clip(chance, 0, 100) + + if np.random.randint(100) < chance: + img = imagelib.LinearMotionBlur (img, np.random.randint( mb_max_size )+1, np.random.randint(360) ) + + if gaussian_blur is not None: + chance, kernel_max_size = gaussian_blur + chance = np.clip(chance, 0, 100) + + if np.random.randint(100) < chance: + img = cv2.GaussianBlur(img, ( np.random.randint( kernel_max_size )*2+1 ,) *2 , 0) + + if is_face_sample and target_face_type != SPTF.NONE: + target_ft = SampleProcessor.SPTF_FACETYPE_TO_FACETYPE[target_face_type] + if target_ft > sample.face_type: + raise Exception ('sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.' % (sample.filename, sample.face_type, target_ft) ) + + if sample.face_type == FaceType.MARK_ONLY: + #first warp to target facetype + img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, sample.shape[0], target_ft), (sample.shape[0],sample.shape[0]), flags=cv2.INTER_CUBIC ) + mask = cv2.warpAffine( mask, LandmarksProcessor.get_transform_mat (sample.landmarks, sample.shape[0], target_ft), (sample.shape[0],sample.shape[0]), flags=cv2.INTER_CUBIC ) + #then apply transforms + img = do_transform (img, mask) + img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC ) else: - mask = LandmarksProcessor.get_image_hull_mask (img.shape, sample.landmarks) + img = do_transform (img, mask) + img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, resolution, target_ft), (resolution,resolution), borderMode=(cv2.BORDER_REPLICATE if border_replicate else cv2.BORDER_CONSTANT), flags=cv2.INTER_CUBIC ) - if sample.ie_polys is not None: - sample.ie_polys.overlay_mask(mask) - ################## - - - if motion_blur is not None: - chance, mb_max_size = motion_blur - chance = np.clip(chance, 0, 100) - - if np.random.randint(100) < chance: - img = imagelib.LinearMotionBlur (img, np.random.randint( mb_max_size )+1, np.random.randint(360) ) - - if gaussian_blur is not None: - chance, kernel_max_size = gaussian_blur - chance = np.clip(chance, 0, 100) - - if np.random.randint(100) < chance: - img = cv2.GaussianBlur(img, ( np.random.randint( kernel_max_size )*2+1 ,) *2 , 0) - - if is_face_sample and target_face_type != SPTF.NONE: - target_ft = SampleProcessor.SPTF_FACETYPE_TO_FACETYPE[target_face_type] - if target_ft > sample.face_type: - raise Exception ('sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.' % (sample.filename, sample.face_type, target_ft) ) - - if sample.face_type == FaceType.MARK_ONLY: - #first warp to target facetype - img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, sample.shape[0], target_ft), (sample.shape[0],sample.shape[0]), flags=cv2.INTER_CUBIC ) - mask = cv2.warpAffine( mask, LandmarksProcessor.get_transform_mat (sample.landmarks, sample.shape[0], target_ft), (sample.shape[0],sample.shape[0]), flags=cv2.INTER_CUBIC ) - #then apply transforms + else: img = do_transform (img, mask) img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC ) - else: - img = do_transform (img, mask) - img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, resolution, target_ft), (resolution,resolution), borderMode=(cv2.BORDER_REPLICATE if border_replicate else cv2.BORDER_CONSTANT), flags=cv2.INTER_CUBIC ) - else: - img = do_transform (img, mask) - img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC ) + if random_sub_res != 0: + sub_size = resolution - random_sub_res + rnd_state = np.random.RandomState (sample_rnd_seed+random_sub_res) + start_x = rnd_state.randint(sub_size+1) + start_y = rnd_state.randint(sub_size+1) + img = img[start_y:start_y+sub_size,start_x:start_x+sub_size,:] - if random_sub_res != 0: - sub_size = resolution - random_sub_res - rnd_state = np.random.RandomState (sample_rnd_seed+random_sub_res) - start_x = rnd_state.randint(sub_size+1) - start_y = rnd_state.randint(sub_size+1) - img = img[start_y:start_y+sub_size,start_x:start_x+sub_size,:] + img = np.clip(img, 0, 1).astype(np.float32) + img_bgr = img[...,0:3] + img_mask = img[...,3:4] - img = np.clip(img, 0, 1).astype(np.float32) - img_bgr = img[...,0:3] - img_mask = img[...,3:4] + if ct_mode is not None and ct_sample is not None: + if ct_sample_bgr is None: + ct_sample_bgr = ct_sample.load_bgr() - if ct_mode is not None and ct_sample is not None: - if ct_sample_bgr is None: - ct_sample_bgr = ct_sample.load_bgr() + ct_sample_bgr_resized = cv2.resize( ct_sample_bgr, (resolution,resolution), cv2.INTER_LINEAR ) - ct_sample_bgr_resized = cv2.resize( ct_sample_bgr, (resolution,resolution), cv2.INTER_LINEAR ) + if ct_mode == 'lct': + img_bgr = imagelib.linear_color_transfer (img_bgr, ct_sample_bgr_resized) + img_bgr = np.clip( img_bgr, 0.0, 1.0) + elif ct_mode == 'rct': + img_bgr = imagelib.reinhard_color_transfer ( np.clip( (img_bgr*255).astype(np.uint8), 0, 255), + np.clip( (ct_sample_bgr_resized*255).astype(np.uint8), 0, 255) ) + img_bgr = np.clip( img_bgr.astype(np.float32) / 255.0, 0.0, 1.0) + elif ct_mode == 'mkl': + img_bgr = imagelib.color_transfer_mkl (img_bgr, ct_sample_bgr_resized) + elif ct_mode == 'idt': + img_bgr = imagelib.color_transfer_idt (img_bgr, ct_sample_bgr_resized) + elif ct_mode == 'sot': + img_bgr = imagelib.color_transfer_sot (img_bgr, ct_sample_bgr_resized) + img_bgr = np.clip( img_bgr, 0.0, 1.0) - if ct_mode == 'lct': - img_bgr = imagelib.linear_color_transfer (img_bgr, ct_sample_bgr_resized) - img_bgr = np.clip( img_bgr, 0.0, 1.0) - elif ct_mode == 'rct': - img_bgr = imagelib.reinhard_color_transfer ( np.clip( (img_bgr*255).astype(np.uint8), 0, 255), - np.clip( (ct_sample_bgr_resized*255).astype(np.uint8), 0, 255) ) - img_bgr = np.clip( img_bgr.astype(np.float32) / 255.0, 0.0, 1.0) - elif ct_mode == 'mkl': - img_bgr = imagelib.color_transfer_mkl (img_bgr, ct_sample_bgr_resized) - elif ct_mode == 'idt': - img_bgr = imagelib.color_transfer_idt (img_bgr, ct_sample_bgr_resized) - elif ct_mode == 'sot': - img_bgr = imagelib.color_transfer_sot (img_bgr, ct_sample_bgr_resized) - img_bgr = np.clip( img_bgr, 0.0, 1.0) - - if random_hsv_shift: - rnd_state = np.random.RandomState (sample_rnd_seed) - hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV) - h, s, v = cv2.split(hsv) - - h = (h + rnd_state.randint(360) ) % 360 - s = np.clip ( s + rnd_state.random()-0.5, 0, 1 ) - v = np.clip ( v + rnd_state.random()-0.5, 0, 1 ) - hsv = cv2.merge([h, s, v]) - - img_bgr = np.clip( cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) , 0, 1 ) + if normalize_std_dev: + img_bgr = (img_bgr - img_bgr.mean( (0,1)) ) / img_bgr.std( (0,1) ) + elif normalize_vgg: + img_bgr = np.clip(img_bgr*255, 0, 255) + img_bgr[:,:,0] -= 103.939 + img_bgr[:,:,1] -= 116.779 + img_bgr[:,:,2] -= 123.68 + + if mode_type == SPTF.MODE_BGR: + img = img_bgr + elif mode_type == SPTF.MODE_BGR_SHUFFLE: + rnd_state = np.random.RandomState (sample_rnd_seed) + img = np.take (img_bgr, rnd_state.permutation(img_bgr.shape[-1]), axis=-1) + + elif mode_type == SPTF.MODE_BGR_RANDOM_HSV_SHIFT: + rnd_state = np.random.RandomState (sample_rnd_seed) + hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV) + h, s, v = cv2.split(hsv) + h = (h + rnd_state.randint(360) ) % 360 + s = np.clip ( s + rnd_state.random()-0.5, 0, 1 ) + v = np.clip ( v + rnd_state.random()-0.5, 0, 1 ) + hsv = cv2.merge([h, s, v]) + img = np.clip( cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) , 0, 1 ) + elif mode_type == SPTF.MODE_G: + img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)[...,None] + elif mode_type == SPTF.MODE_GGG: + img = np.repeat ( np.expand_dims(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY),-1), (3,), -1) + elif mode_type == SPTF.MODE_M and is_face_sample: + img = img_mask - if normalize_std_dev: - img_bgr = (img_bgr - img_bgr.mean( (0,1)) ) / img_bgr.std( (0,1) ) - elif normalize_vgg: - img_bgr = np.clip(img_bgr*255, 0, 255) - img_bgr[:,:,0] -= 103.939 - img_bgr[:,:,1] -= 116.779 - img_bgr[:,:,2] -= 123.68 - - if mode_type == SPTF.MODE_BGR: - img = img_bgr - elif mode_type == SPTF.MODE_BGR_SHUFFLE: - rnd_state = np.random.RandomState (sample_rnd_seed) - img = np.take (img_bgr, rnd_state.permutation(img_bgr.shape[-1]), axis=-1) - elif mode_type == SPTF.MODE_G: - img = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)[...,None] - elif mode_type == SPTF.MODE_GGG: - img = np.repeat ( np.expand_dims(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY),-1), (3,), -1) - elif mode_type == SPTF.MODE_M and is_face_sample: - img = img_mask + if not debug: + if normalize_tanh: + img = np.clip (img * 2.0 - 1.0, -1.0, 1.0) + else: + img = np.clip (img, 0.0, 1.0) - if not debug: - if normalize_tanh: - img = np.clip (img * 2.0 - 1.0, -1.0, 1.0) - else: - img = np.clip (img, 0.0, 1.0) - - outputs.append ( img ) - - if debug: - result = [] - - for output in outputs: - if output.shape[2] < 4: - result += [output,] - elif output.shape[2] == 4: - result += [output[...,0:3]*output[...,3:4],] - - return result - else: - return outputs + outputs_sample.append ( img ) + outputs += [outputs_sample] + + return outputs """ close_sample = sample.close_target_list[ np.random.randint(0, len(sample.close_target_list)) ] if sample.close_target_list is not None else None diff --git a/utils/iter_utils.py b/utils/iter_utils.py index eb00c55..e690e3b 100644 --- a/utils/iter_utils.py +++ b/utils/iter_utils.py @@ -22,7 +22,7 @@ class ThisThreadGenerator(object): return next(self.generator_func) class SubprocessGenerator(object): - def __init__(self, generator_func, user_param=None, prefetch=2): + def __init__(self, generator_func, user_param=None, prefetch=2, start_now=False): super().__init__() self.prefetch = prefetch self.generator_func = generator_func @@ -30,6 +30,16 @@ class SubprocessGenerator(object): self.sc_queue = multiprocessing.Queue() self.cs_queue = multiprocessing.Queue() self.p = None + if start_now: + self._start() + + def _start(self): + if self.p == None: + user_param = self.user_param + self.user_param = None + self.p = multiprocessing.Process(target=self.process_func, args=(user_param,) ) + self.p.daemon = True + self.p.start() def process_func(self, user_param): self.generator_func = self.generator_func(user_param) @@ -54,13 +64,7 @@ class SubprocessGenerator(object): return self_dict def __next__(self): - if self.p == None: - user_param = self.user_param - self.user_param = None - self.p = multiprocessing.Process(target=self.process_func, args=(user_param,) ) - self.p.daemon = True - self.p.start() - + self._start() gen_data = self.cs_queue.get() if gen_data is None: self.p.terminate()