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

optimizations of nnlib and SampleGeneratorFace,

Browse source 
refactorings
This commit is contained in:
iperov 2019-01-22 11:52:04 +04:00
parent 2de45083a4
commit b6c4171ea1
9 changed files with 175 additions and 79 deletions
Download patch file Download diff file Expand all files Collapse all files

30
samples/SampleGeneratorFace.py
View file

@ -11,13 +11,14 @@ from samples import SampleLoader
from samples import SampleGeneratorBase
'''
arg
output_sample_types = [
[SampleProcessor.TypeFlags, size, (optional)random_sub_size] ,
...
]
'''
class SampleGeneratorFace(SampleGeneratorBase):
def __init__ (self, samples_path, debug, batch_size, sort_by_yaw=False, sort_by_yaw_target_samples_path=None, sample_process_options=SampleProcessor.Options(), output_sample_types=[], **kwargs):
def __init__ (self, samples_path, debug, batch_size, sort_by_yaw=False, sort_by_yaw_target_samples_path=None, with_close_to_self=False, sample_process_options=SampleProcessor.Options(), output_sample_types=[], generators_count=2, **kwargs):
super().__init__(samples_path, debug, batch_size)
self.sample_process_options = sample_process_options
self.output_sample_types = output_sample_types
@ -26,24 +27,20 @@ class SampleGeneratorFace(SampleGeneratorBase):
self.sample_type = SampleType.FACE_YAW_SORTED_AS_TARGET
elif sort_by_yaw:
self.sample_type = SampleType.FACE_YAW_SORTED
elif with_close_to_self:
self.sample_type = SampleType.FACE_WITH_CLOSE_TO_SELF
else:
self.sample_type = SampleType.FACE
self.sample_type = SampleType.FACE
self.samples = SampleLoader.load (self.sample_type, self.samples_path, sort_by_yaw_target_samples_path)
self.generators_count = min ( generators_count, len(self.samples) )
if self.debug:
self.generator_samples = [ self.samples ]
self.generators = [iter_utils.ThisThreadGenerator ( self.batch_func, 0 )]
else:
if len(self.samples) > 1:
self.generator_samples = [ self.samples[0::2],
self.samples[1::2] ]
self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, 0 ),
iter_utils.SubprocessGenerator ( self.batch_func, 1 )]
else:
self.generator_samples = [ self.samples ]
self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, 0 )]
self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, i ) for i in range(self.generators_count) ]
self.generator_counter = -1
def __iter__(self):
@ -55,7 +52,8 @@ class SampleGeneratorFace(SampleGeneratorBase):
return next(generator)
def batch_func(self, generator_id):
samples = self.generator_samples[generator_id]
samples = self.samples[generator_id::self.generators_count]
data_len = len(samples)
if data_len == 0:
raise ValueError('No training data provided.')
@ -64,7 +62,7 @@ class SampleGeneratorFace(SampleGeneratorBase):
if all ( [ x == None for x in samples] ):
raise ValueError('Not enough training data. Gather more faces!')
if self.sample_type == SampleType.FACE:
if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
shuffle_idxs = []
elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
shuffle_idxs = []
@ -77,7 +75,7 @@ class SampleGeneratorFace(SampleGeneratorBase):
while True:
sample = None
if self.sample_type == SampleType.FACE:
if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
if len(shuffle_idxs) == 0:
shuffle_idxs = random.sample( range(data_len), data_len )
idx = shuffle_idxs.pop()