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

initial code to extract umdfaces.io dataset and train pose estimator

Browse source
This commit is contained in:
iperov 2019-04-23 08:14:09 +04:00
parent 51a917facc
commit e58197ca22
18 changed files with 437 additions and 57 deletions
Download patch file Download diff file Expand all files Collapse all files

20
samplelib/SampleGeneratorFace.py
View file

@ -15,13 +15,12 @@ output_sample_types = [
]
'''
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=[], add_sample_idx=False, add_pitch=False, add_yaw=False, generators_count=2, generators_random_seed=None, **kwargs):
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=[], add_sample_idx=False, generators_count=2, generators_random_seed=None, **kwargs):
super().__init__(samples_path, debug, batch_size)
self.sample_process_options = sample_process_options
self.output_sample_types = output_sample_types
self.add_sample_idx = add_sample_idx
self.add_pitch = add_pitch
self.add_yaw = add_yaw
# self.add_pitch_yaw_roll = add_pitch_yaw_roll
if sort_by_yaw_target_samples_path is not None:
self.sample_type = SampleType.FACE_YAW_SORTED_AS_TARGET
@ -143,12 +142,6 @@ class SampleGeneratorFace(SampleGeneratorBase):
if self.add_sample_idx:
batches += [ [] ]
i_sample_idx = len(batches)-1
if self.add_pitch:
batches += [ [] ]
i_pitch = len(batches)-1
if self.add_yaw:
batches += [ [] ]
i_yaw = len(batches)-1
for i in range(len(x)):
batches[i].append ( x[i] )
@ -156,14 +149,5 @@ class SampleGeneratorFace(SampleGeneratorBase):
if self.add_sample_idx:
batches[i_sample_idx].append (idx)
if self.add_pitch or self.add_yaw:
pitch, yaw = LandmarksProcessor.estimate_pitch_yaw (sample.landmarks)
if self.add_pitch:
batches[i_pitch].append ([pitch])
if self.add_yaw:
batches[i_yaw].append ([yaw])
break
yield [ np.array(batch) for batch in batches]