mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-07 05:22:06 -07:00
nothing interesting
This commit is contained in:
Colombo
parent
154820a954
commit
e8673e3fcc
10 changed files with 339 additions and 262 deletions
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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},
|
||||
])
|
||||
])
|
||||
|
||||
|
|
|
|||
|
|
@ -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 = []
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
return SampleLoader.get_person_id_max_count(samples_path)
|
||||
|
|
|
|||
|
|
@ -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() ) )
|
||||
|
||||
|
|
|
|||
|
|
@ -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() ) )
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue