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

refactorings

Browse source
This commit is contained in:
Colombo 2019-12-22 19:00:59 +04:00
parent e0e8970ab9
commit 754d6c385c
13 changed files with 243 additions and 104 deletions
Download patch file Download diff file Expand all files Collapse all files

114
samplelib/SampleGeneratorFacePerson.py
View file

@ -8,7 +8,7 @@ import numpy as np
from facelib import LandmarksProcessor
from samplelib import (SampleGeneratorBase, SampleHost, SampleProcessor,
SampleType)
from utils import iter_utils
from utils import iter_utils, mp_utils
'''
@ -23,9 +23,6 @@ class SampleGeneratorFacePerson(SampleGeneratorBase):
sample_process_options=SampleProcessor.Options(),
output_sample_types=[],
person_id_mode=1,
use_caching=False,
generators_count=2,
generators_random_seed=None,
**kwargs):
super().__init__(samples_path, debug, batch_size)
@ -33,51 +30,32 @@ class SampleGeneratorFacePerson(SampleGeneratorBase):
self.output_sample_types = output_sample_types
self.person_id_mode = person_id_mode
if generators_random_seed is not None and len(generators_random_seed) != generators_count:
raise ValueError("len(generators_random_seed) != generators_count")
self.generators_random_seed = generators_random_seed
samples = SampleHost.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)
#
#new_samples = []
#while len(samples) > 0:
# for i in range( len(samples)-1, -1, -1):
# sample = samples[i]
#
# if len(sample) > 0:
# new_samples.append(sample.pop(0))
#
# if len(sample) == 0:
# samples.pop(i)
# i -= 1
#samples = new_samples
new_samples = []
for s in samples:
new_samples += s
samples = new_samples
np.random.shuffle(samples)
samples_host = SampleHost.mp_host (SampleType.FACE, self.samples_path)
samples = samples_host.get_list()
self.samples_len = len(samples)
if self.samples_len == 0:
raise ValueError('No training data provided.')
raise ValueError('No training data provided.')
persons_name_idxs = {}
for i,sample in enumerate(samples):
person_name = sample.person_name
if person_name not in persons_name_idxs:
persons_name_idxs[person_name] = []
persons_name_idxs[person_name].append (i)
indexes2D = [ persons_name_idxs[person_name] for person_name in sorted(list(persons_name_idxs.keys())) ]
index2d_host = mp_utils.Index2DHost(indexes2D)
if self.debug:
self.generators_count = 1
self.generators = [iter_utils.ThisThreadGenerator ( self.batch_func, (0, samples) )]
self.generators = [iter_utils.ThisThreadGenerator ( self.batch_func, (samples_host.create_cli(), index2d_host.create_cli(),) )]
else:
self.generators_count = min ( generators_count, self.samples_len )
if person_id_mode==1:
self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, (i, samples[i::self.generators_count]) ) for i in range(self.generators_count) ]
else:
self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, (i, samples) ) for i in range(self.generators_count) ]
self.generators_count = np.clip(multiprocessing.cpu_count(), 2, 4)
self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, (samples_host.create_cli(), index2d_host.create_cli(),), start_now=True ) for i in range(self.generators_count) ]
self.generator_counter = -1
@ -94,12 +72,43 @@ class SampleGeneratorFacePerson(SampleGeneratorBase):
return next(generator)
def batch_func(self, param ):
generator_id, samples = param
samples, index2d_host, = param
bs = self.batch_size
if self.generators_random_seed is not None:
np.random.seed ( self.generators_random_seed[generator_id] )
while True:
person_idxs = index2d_host.get_1D(bs)
samples_idxs = index2d_host.get_2D(person_idxs, 1)
batches = None
for n_batch in range(bs):
person_id = person_idxs[n_batch]
sample_idx = samples_idxs[n_batch][0]
if self.person_id_mode==1:
sample = samples[ sample_idx ]
try:
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() ) )
if batches is None:
batches = [ [] for _ in range(len(x)) ]
batches += [ [] ]
i_person_id = len(batches)-1
for i in range(len(x)):
batches[i].append ( x[i] )
batches[i_person_id].append ( np.array([person_id]) )
yield [ np.array(batch) for batch in batches]
@staticmethod
def get_person_id_max_count(samples_path):
return SampleHost.get_person_id_max_count(samples_path)
"""
if self.person_id_mode==1:
samples_len = len(samples)
samples_idxs = [*range(samples_len)]
shuffle_idxs = []
@ -132,9 +141,7 @@ class SampleGeneratorFacePerson(SampleGeneratorBase):
samples_idxs[i] = [*range(len(samples[i]))]
shuffle_idxs[i] = []
while True:
if self.person_id_mode==2:
if self.person_id_mode==2:
if len(shuffle_person_idxs) == 0:
shuffle_person_idxs = person_idxs.copy()
np.random.shuffle(shuffle_person_idxs)
@ -270,9 +277,4 @@ class SampleGeneratorFacePerson(SampleGeneratorBase):
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 SampleHost.get_person_id_max_count(samples_path)
"""