github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-07-06 21:12:07 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

nothing interesting

Browse source
This commit is contained in:
Colombo 2019-12-11 22:33:23 +04:00
parent 154820a954
commit e8673e3fcc
10 changed files with 339 additions and 262 deletions
Download patch file Download diff file Expand all files Collapse all files

20
imagelib/warp.py
View file

@ -2,18 +2,24 @@ import numpy as np
import cv2 import cv2
from utils import random_utils 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 h,w,c = source.shape
if (h != w): if (h != w):
raise ValueError ('gen_warp_params accepts only square images.') raise ValueError ('gen_warp_params accepts only square images.')
rotation = np.random.uniform( rotation_range[0], rotation_range[1] ) if rnd_seed != None:
scale = np.random.uniform(1 +scale_range[0], 1 +scale_range[1]) rnd_state = np.random.RandomState (rnd_seed)
tx = np.random.uniform( tx_range[0], tx_range[1] ) else:
ty = np.random.uniform( ty_range[0], ty_range[1] ) rnd_state = np.random
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 #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 cell_count = w // cell_size + 1
grid_points = np.linspace( 0, w, cell_count) 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['mapy'] = mapy
params['rmat'] = random_transform_mat params['rmat'] = random_transform_mat
params['w'] = w params['w'] = w
params['flip'] = flip and np.random.randint(10) < 4 params['flip'] = p_flip
return params return params

4
models/Model_DEV_FANSEG/Model.py
View file

@ -47,13 +47,13 @@ class Model(ModelBase):
self.set_training_data_generators ([ self.set_training_data_generators ([
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=True), 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 }, { '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, SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=True ), 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},
]) ])
]) ])

22
models/Model_DEV_FUNIT/Model.py
View file

@ -23,9 +23,9 @@ class FUNITModel(ModelBase):
#override #override
def onInitializeOptions(self, is_first_run, ask_override): def onInitializeOptions(self, is_first_run, ask_override):
default_resolution = 96 default_resolution = 64
if is_first_run: 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: else:
self.options['resolution'] = self.options.get('resolution', default_resolution) self.options['resolution'] = self.options.get('resolution', default_resolution)
@ -48,7 +48,7 @@ class FUNITModel(ModelBase):
resolution = self.options['resolution'] resolution = self.options['resolution']
face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF 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), 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} ] output_sample_types1=[ {'types': (t.IMG_SOURCE, face_type, t.MODE_BGR), 'resolution':resolution, 'normalize_tanh':True} ]
self.set_training_data_generators ([ 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] ), 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] ), 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]), 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]), 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 #override

39
nnlib/nnlib.py
View file

@ -93,6 +93,7 @@ Model = keras.models.Model
Adam = nnlib.Adam Adam = nnlib.Adam
RMSprop = nnlib.RMSprop RMSprop = nnlib.RMSprop
LookaheadOptimizer = nnlib.LookaheadOptimizer LookaheadOptimizer = nnlib.LookaheadOptimizer
SGD = nnlib.keras.optimizers.SGD
modelify = nnlib.modelify modelify = nnlib.modelify
gaussian_blur = nnlib.gaussian_blur 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. 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.initial_decay = kwargs.pop('decay', 0.0)
self.epsilon = kwargs.pop('epsilon', K.epsilon()) self.epsilon = kwargs.pop('epsilon', K.epsilon())
self.lr_dropout = lr_dropout
self.tf_cpu_mode = tf_cpu_mode self.tf_cpu_mode = tf_cpu_mode
learning_rate = kwargs.pop('lr', learning_rate) learning_rate = kwargs.pop('lr', learning_rate)
@ -788,6 +790,8 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
dtype=K.dtype(p), dtype=K.dtype(p),
name='accumulator_' + str(i)) name='accumulator_' + str(i))
for (i, p) in enumerate(params)] 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) if e: e.__exit__(None, None, None)
self.weights = [self.iterations] + accumulators self.weights = [self.iterations] + accumulators
@ -798,12 +802,15 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
lr = lr * (1. / (1. + self.decay * K.cast(self.iterations, lr = lr * (1. / (1. + self.decay * K.cast(self.iterations,
K.dtype(self.decay)))) 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 # update accumulator
e = K.tf.device("/cpu:0") if self.tf_cpu_mode == 2 else None e = K.tf.device("/cpu:0") if self.tf_cpu_mode == 2 else None
if e: e.__enter__() if e: e.__enter__()
new_a = self.rho * a + (1. - self.rho) * K.square(g) 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) if e: e.__exit__(None, None, None)
self.updates.append(K.update(a, new_a)) 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)), config = {'learning_rate': float(K.get_value(self.learning_rate)),
'rho': float(K.get_value(self.rho)), 'rho': float(K.get_value(self.rho)),
'decay': float(K.get_value(self.decay)), '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() base_config = super(RMSprop, self).get_config()
return dict(list(base_config.items()) + list(config.items())) return dict(list(base_config.items()) + list(config.items()))
nnlib.RMSprop = RMSprop nnlib.RMSprop = RMSprop
@ -847,6 +855,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
amsgrad: boolean. Whether to apply the AMSGrad variant of this amsgrad: boolean. Whether to apply the AMSGrad variant of this
algorithm from the paper "On the Convergence of Adam and algorithm from the paper "On the Convergence of Adam and
Beyond". Beyond".
lr_dropout: float [0.0 .. 1.0] Learning rate dropout https://arxiv.org/pdf/1912.00144
tf_cpu_mode: only for tensorflow backend tf_cpu_mode: only for tensorflow backend
0 - default, no changes. 0 - default, no changes.
1 - allows to train x2 bigger network on same VRAM consuming RAM 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, 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) super(Adam, self).__init__(**kwargs)
with K.name_scope(self.__class__.__name__): with K.name_scope(self.__class__.__name__):
self.iterations = K.variable(0, dtype='int64', name='iterations') self.iterations = K.variable(0, dtype='int64', name='iterations')
@ -873,6 +882,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
self.epsilon = epsilon self.epsilon = epsilon
self.initial_decay = decay self.initial_decay = decay
self.amsgrad = amsgrad self.amsgrad = amsgrad
self.lr_dropout = lr_dropout
self.tf_cpu_mode = tf_cpu_mode self.tf_cpu_mode = tf_cpu_mode
def get_updates(self, loss, params): 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] vhats = [K.zeros(K.int_shape(p), dtype=K.dtype(p)) for p in params]
else: else:
vhats = [K.zeros(1) for _ in params] 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) if e: e.__exit__(None, None, None)
self.weights = [self.iterations] + ms + vs + vhats 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 e = K.tf.device("/cpu:0") if self.tf_cpu_mode == 2 else None
if e: e.__enter__() if e: e.__enter__()
m_t = (self.beta_1 * m) + (1. - self.beta_1) * g 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 e: e.__exit__(None, None, None)
if self.amsgrad: 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: 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(m, m_t))
self.updates.append(K.update(v, v_t)) self.updates.append(K.update(v, v_t))
new_p = p_t new_p = p + p_diff
# Apply constraints. # Apply constraints.
if getattr(p, 'constraint', None) is not None: if getattr(p, 'constraint', None) is not None:
@ -933,7 +951,8 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
'beta_2': float(K.get_value(self.beta_2)), 'beta_2': float(K.get_value(self.beta_2)),
'decay': float(K.get_value(self.decay)), 'decay': float(K.get_value(self.decay)),
'epsilon': self.epsilon, 'epsilon': self.epsilon,
'amsgrad': self.amsgrad} 'amsgrad': self.amsgrad,
'lr_dropout' : self.lr_dropout}
base_config = super(Adam, self).get_config() base_config = super(Adam, self).get_config()
return dict(list(base_config.items()) + list(config.items())) return dict(list(base_config.items()) + list(config.items()))
nnlib.Adam = Adam nnlib.Adam = Adam

9
samplelib/SampleGeneratorFace.py
View file

@ -55,7 +55,6 @@ class SampleGeneratorFace(SampleGeneratorBase):
raise ValueError('No training data provided.') 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 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: if self.debug:
self.generators_count = 1 self.generators_count = 1
@ -133,16 +132,12 @@ class SampleGeneratorFace(SampleGeneratorBase):
try: try:
ct_sample=None ct_sample=None
if ct_samples is not 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: except:
raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) ) 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: if batches is None:
batches = [ [] for _ in range(len(x)) ] batches = [ [] for _ in range(len(x)) ]
if self.add_sample_idx: if self.add_sample_idx:

76
samplelib/SampleGeneratorFacePerson.py
View file

@ -1,3 +1,4 @@
import copy
import multiprocessing import multiprocessing
import traceback import traceback
@ -37,6 +38,9 @@ class SampleGeneratorFacePerson(SampleGeneratorBase):
self.generators_random_seed = generators_random_seed self.generators_random_seed = generators_random_seed
samples = SampleLoader.load (SampleType.FACE, self.samples_path, person_id_mode=True, use_caching=use_caching) 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: if person_id_mode==1:
np.random.shuffle(samples) np.random.shuffle(samples)
@ -52,6 +56,7 @@ class SampleGeneratorFacePerson(SampleGeneratorBase):
if len(sample) == 0: if len(sample) == 0:
samples.pop(i) samples.pop(i)
samples = new_samples samples = new_samples
#new_samples = [] #new_samples = []
#for s in samples: #for s in samples:
# new_samples += s # new_samples += s
@ -111,6 +116,18 @@ class SampleGeneratorFacePerson(SampleGeneratorBase):
samples_idxs = [None]*persons_count samples_idxs = [None]*persons_count
shuffle_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] = []
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): for i in range(persons_count):
samples_idxs[i] = [*range(len(samples[i]))] samples_idxs[i] = [*range(len(samples[i]))]
shuffle_idxs[i] = [] shuffle_idxs[i] = []
@ -130,13 +147,13 @@ class SampleGeneratorFacePerson(SampleGeneratorBase):
if self.person_id_mode==1: if self.person_id_mode==1:
if len(shuffle_idxs) == 0: if len(shuffle_idxs) == 0:
shuffle_idxs = samples_idxs.copy() shuffle_idxs = samples_idxs.copy()
#np.random.shuffle(shuffle_idxs) np.random.shuffle(shuffle_idxs) ###
idx = shuffle_idxs.pop() idx = shuffle_idxs.pop()
sample = samples[ idx ] sample = samples[ idx ]
try: 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: except:
raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) ) 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]) ) batches[i_person_id].append ( np.array([sample.person_id]) )
else: elif self.person_id_mode==2:
person_id1, person_id2 = person_ids person_id1, person_id2 = person_ids
if len(shuffle_idxs[person_id1]) == 0: if len(shuffle_idxs[person_id1]) == 0:
@ -174,12 +191,12 @@ class SampleGeneratorFacePerson(SampleGeneratorBase):
if sample1 is not None and sample2 is not None: if sample1 is not None and sample2 is not None:
try: 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: except:
raise Exception ("Exception occured in sample %s. Error: %s" % (sample1.filename, traceback.format_exc() ) ) raise Exception ("Exception occured in sample %s. Error: %s" % (sample1.filename, traceback.format_exc() ) )
try: 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: except:
raise Exception ("Exception occured in sample %s. Error: %s" % (sample2.filename, traceback.format_exc() ) ) raise Exception ("Exception occured in sample %s. Error: %s" % (sample2.filename, traceback.format_exc() ) )
@ -203,7 +220,56 @@ class SampleGeneratorFacePerson(SampleGeneratorBase):
batches[i_person_id2].append ( np.array([sample2.person_id]) ) 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] yield [ np.array(batch) for batch in batches]

2
samplelib/SampleGeneratorFaceTemporal.py
View file

@ -71,7 +71,7 @@ class SampleGeneratorFaceTemporal(SampleGeneratorBase):
for i in range( self.temporal_image_count ): for i in range( self.temporal_image_count ):
sample = samples[ idx+i*mult ] sample = samples[ idx+i*mult ]
try: 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: except:
raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) ) raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) )

2
samplelib/SampleGeneratorImageTemporal.py
View file

@ -66,7 +66,7 @@ class SampleGeneratorImageTemporal(SampleGeneratorBase):
for i in range( self.temporal_image_count ): for i in range( self.temporal_image_count ):
sample = samples[ idx+i*mult ] sample = samples[ idx+i*mult ]
try: 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: except:
raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) ) raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) )

367
samplelib/SampleProcessor.py
View file

@ -92,231 +92,218 @@ class SampleProcessor(object):
} }
@staticmethod @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 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) sample_rnd_seed = np.random.randint(0x80000000)
outputs = [] 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) is_face_sample = sample.landmarks is not None
types = opts.get('types', [] )
border_replicate = opts.get('border_replicate', True) if debug and is_face_sample:
random_sub_res = opts.get('random_sub_res', 0) LandmarksProcessor.draw_landmarks (sample_bgr, sample.landmarks, (0, 1, 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) 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 )
ct_mode = opts.get('ct_mode', 'None')
normalize_tanh = opts.get('normalize_tanh', False)
img_type = SPTF.NONE outputs_sample = []
target_face_type = SPTF.NONE for opts in output_sample_types:
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: resolution = opts.get('resolution', 0)
raise ValueError ('expected IMG_ type') types = opts.get('types', [] )
if img_type == SPTF.IMG_LANDMARKS_ARRAY: border_replicate = opts.get('border_replicate', True)
l = sample.landmarks random_sub_res = opts.get('random_sub_res', 0)
l = np.concatenate ( [ np.expand_dims(l[:,0] / w,-1), np.expand_dims(l[:,1] / h,-1) ], -1 ) normalize_std_dev = opts.get('normalize_std_dev', False)
l = np.clip(l, 0.0, 1.0) normalize_vgg = opts.get('normalize_vgg', False)
img = l motion_blur = opts.get('motion_blur', None)
elif img_type == SPTF.IMG_PITCH_YAW_ROLL or img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID: gaussian_blur = opts.get('gaussian_blur', None)
pitch_yaw_roll = sample.pitch_yaw_roll
if pitch_yaw_roll is not None: ct_mode = opts.get('ct_mode', 'None')
pitch, yaw, roll = pitch_yaw_roll 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: else:
pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll (sample.landmarks) if mode_type == SPTF.NONE:
if params['flip']: raise ValueError ('expected MODE_ type')
yaw = -yaw
if img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID: def do_transform(img, mask):
pitch = (pitch+1.0) / 2.0 warp = (img_type==SPTF.IMG_WARPED or img_type==SPTF.IMG_WARPED_TRANSFORMED)
yaw = (yaw+1.0) / 2.0 transform = (img_type==SPTF.IMG_WARPED_TRANSFORMED or img_type==SPTF.IMG_TRANSFORMED)
roll = (roll+1.0) / 2.0 flip = img_type != SPTF.IMG_WARPED
img = (pitch, yaw, roll) img = imagelib.warp_by_params (params, img, warp, transform, flip, border_replicate)
else: if mask is not None:
if mode_type == SPTF.NONE: mask = imagelib.warp_by_params (params, mask, warp, transform, flip, False)
raise ValueError ('expected MODE_ type') if len(mask.shape) == 2:
mask = mask[...,np.newaxis]
def do_transform(img, mask): img = np.concatenate( (img, mask ), -1 )
warp = (img_type==SPTF.IMG_WARPED or img_type==SPTF.IMG_WARPED_TRANSFORMED) return img
transform = (img_type==SPTF.IMG_WARPED_TRANSFORMED or img_type==SPTF.IMG_TRANSFORMED)
flip = img_type != SPTF.IMG_WARPED
img = imagelib.warp_by_params (params, img, warp, transform, flip, border_replicate) img = sample_bgr
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 = np.concatenate( (img, mask ), -1 ) ### Prepare a mask
return img 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 if sample.ie_polys is not None:
mask = None sample.ie_polys.overlay_mask(mask)
if is_face_sample: ##################
mask = sample.load_fanseg_mask() #using fanseg_mask if exist
if mask is None:
if sample.eyebrows_expand_mod is not None: if motion_blur is not None:
mask = LandmarksProcessor.get_image_hull_mask (img.shape, sample.landmarks, eyebrows_expand_mod=sample.eyebrows_expand_mod ) 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: 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: else:
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
img = do_transform (img, mask) img = do_transform (img, mask)
img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC ) 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: if random_sub_res != 0:
img = do_transform (img, mask) sub_size = resolution - random_sub_res
img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC ) 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: img = np.clip(img, 0, 1).astype(np.float32)
sub_size = resolution - random_sub_res img_bgr = img[...,0:3]
rnd_state = np.random.RandomState (sample_rnd_seed+random_sub_res) img_mask = img[...,3:4]
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) if ct_mode is not None and ct_sample is not None:
img_bgr = img[...,0:3] if ct_sample_bgr is None:
img_mask = img[...,3:4] ct_sample_bgr = ct_sample.load_bgr()
if ct_mode is not None and ct_sample is not None: ct_sample_bgr_resized = cv2.resize( ct_sample_bgr, (resolution,resolution), cv2.INTER_LINEAR )
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 ) 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': if normalize_std_dev:
img_bgr = imagelib.linear_color_transfer (img_bgr, ct_sample_bgr_resized) img_bgr = (img_bgr - img_bgr.mean( (0,1)) ) / img_bgr.std( (0,1) )
img_bgr = np.clip( img_bgr, 0.0, 1.0) elif normalize_vgg:
elif ct_mode == 'rct': img_bgr = np.clip(img_bgr*255, 0, 255)
img_bgr = imagelib.reinhard_color_transfer ( np.clip( (img_bgr*255).astype(np.uint8), 0, 255), img_bgr[:,:,0] -= 103.939
np.clip( (ct_sample_bgr_resized*255).astype(np.uint8), 0, 255) ) img_bgr[:,:,1] -= 116.779
img_bgr = np.clip( img_bgr.astype(np.float32) / 255.0, 0.0, 1.0) img_bgr[:,:,2] -= 123.68
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: if mode_type == SPTF.MODE_BGR:
rnd_state = np.random.RandomState (sample_rnd_seed) img = img_bgr
hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV) elif mode_type == SPTF.MODE_BGR_SHUFFLE:
h, s, v = cv2.split(hsv) rnd_state = np.random.RandomState (sample_rnd_seed)
img = np.take (img_bgr, rnd_state.permutation(img_bgr.shape[-1]), axis=-1)
h = (h + rnd_state.randint(360) ) % 360 elif mode_type == SPTF.MODE_BGR_RANDOM_HSV_SHIFT:
s = np.clip ( s + rnd_state.random()-0.5, 0, 1 ) rnd_state = np.random.RandomState (sample_rnd_seed)
v = np.clip ( v + rnd_state.random()-0.5, 0, 1 ) hsv = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2HSV)
hsv = cv2.merge([h, s, v]) 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
img_bgr = np.clip( cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) , 0, 1 ) 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 normalize_std_dev: outputs_sample.append ( img )
img_bgr = (img_bgr - img_bgr.mean( (0,1)) ) / img_bgr.std( (0,1) ) outputs += [outputs_sample]
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: return outputs
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)
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
""" """
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 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

20
utils/iter_utils.py
View file

@ -22,7 +22,7 @@ class ThisThreadGenerator(object):
return next(self.generator_func) return next(self.generator_func)
class SubprocessGenerator(object): 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__() super().__init__()
self.prefetch = prefetch self.prefetch = prefetch
self.generator_func = generator_func self.generator_func = generator_func
@ -30,6 +30,16 @@ class SubprocessGenerator(object):
self.sc_queue = multiprocessing.Queue() self.sc_queue = multiprocessing.Queue()
self.cs_queue = multiprocessing.Queue() self.cs_queue = multiprocessing.Queue()
self.p = None 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): def process_func(self, user_param):
self.generator_func = self.generator_func(user_param) self.generator_func = self.generator_func(user_param)
@ -54,13 +64,7 @@ class SubprocessGenerator(object):
return self_dict return self_dict
def __next__(self): def __next__(self):
if self.p == None: self._start()
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()
gen_data = self.cs_queue.get() gen_data = self.cs_queue.get()
if gen_data is None: if gen_data is None:
self.p.terminate() self.p.terminate()