mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-15 01:23:44 -07:00
initial code to extract umdfaces.io dataset and train pose estimator
This commit is contained in:
iperov
parent
51a917facc
commit
e58197ca22
18 changed files with 437 additions and 57 deletions
|
|
@ -16,26 +16,26 @@ class SampleType(IntEnum):
|
|||
FACE = 1 #aligned face unsorted
|
||||
FACE_YAW_SORTED = 2 #sorted by yaw
|
||||
FACE_YAW_SORTED_AS_TARGET = 3 #sorted by yaw and included only yaws which exist in TARGET also automatic mirrored
|
||||
FACE_END = 3
|
||||
FACE_TEMPORAL_SORTED = 4
|
||||
FACE_END = 4
|
||||
|
||||
QTY = 4
|
||||
QTY = 5
|
||||
|
||||
class Sample(object):
|
||||
def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, source_filename=None, mirror=None, close_target_list=None, fanseg_mask_exist=False):
|
||||
def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch_yaw_roll=None, source_filename=None, mirror=None, close_target_list=None, fanseg_mask_exist=False):
|
||||
self.sample_type = sample_type if sample_type is not None else SampleType.IMAGE
|
||||
self.filename = filename
|
||||
self.face_type = face_type
|
||||
self.shape = shape
|
||||
self.landmarks = np.array(landmarks) if landmarks is not None else None
|
||||
self.ie_polys = ie_polys
|
||||
self.pitch = pitch
|
||||
self.yaw = yaw
|
||||
self.pitch_yaw_roll = pitch_yaw_roll
|
||||
self.source_filename = source_filename
|
||||
self.mirror = mirror
|
||||
self.close_target_list = close_target_list
|
||||
self.fanseg_mask_exist = fanseg_mask_exist
|
||||
|
||||
def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch=None, yaw=None, source_filename=None, mirror=None, close_target_list=None, fanseg_mask=None, fanseg_mask_exist=None):
|
||||
def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch_yaw_roll=None, source_filename=None, mirror=None, close_target_list=None, fanseg_mask=None, fanseg_mask_exist=None):
|
||||
return Sample(
|
||||
sample_type=sample_type if sample_type is not None else self.sample_type,
|
||||
filename=filename if filename is not None else self.filename,
|
||||
|
|
@ -43,8 +43,7 @@ class Sample(object):
|
|||
shape=shape if shape is not None else self.shape,
|
||||
landmarks=landmarks if landmarks is not None else self.landmarks.copy(),
|
||||
ie_polys=ie_polys if ie_polys is not None else self.ie_polys,
|
||||
pitch=pitch if pitch is not None else self.pitch,
|
||||
yaw=yaw if yaw is not None else self.yaw,
|
||||
pitch_yaw_roll=pitch_yaw_roll if pitch_yaw_roll is not None else self.pitch_yaw_roll,
|
||||
source_filename=source_filename if source_filename is not None else self.source_filename,
|
||||
mirror=mirror if mirror is not None else self.mirror,
|
||||
close_target_list=close_target_list if close_target_list is not None else self.close_target_list,
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
|
|
|||
|
|
@ -35,9 +35,9 @@ class SampleLoader:
|
|||
if datas[sample_type] is None:
|
||||
datas[sample_type] = SampleLoader.upgradeToFaceSamples( [ Sample(filename=filename) for filename in Path_utils.get_image_paths(samples_path) ] )
|
||||
|
||||
# elif sample_type == SampleType.FACE_TEMPORAL_SORTED:
|
||||
# if datas[sample_type] is None:
|
||||
# datas[sample_type] = SampleLoader.upgradeToFaceTemporalSortedSamples( SampleLoader.load(SampleType.FACE, samples_path) )
|
||||
elif sample_type == SampleType.FACE_TEMPORAL_SORTED:
|
||||
if datas[sample_type] is None:
|
||||
datas[sample_type] = SampleLoader.upgradeToFaceTemporalSortedSamples( SampleLoader.load(SampleType.FACE, samples_path) )
|
||||
|
||||
elif sample_type == SampleType.FACE_YAW_SORTED:
|
||||
if datas[sample_type] is None:
|
||||
|
|
@ -69,15 +69,12 @@ class SampleLoader:
|
|||
print ("%s is not a dfl image file required for training" % (s_filename_path.name) )
|
||||
continue
|
||||
|
||||
pitch, yaw = LandmarksProcessor.estimate_pitch_yaw ( dflimg.get_landmarks() )
|
||||
|
||||
sample_list.append( s.copy_and_set(sample_type=SampleType.FACE,
|
||||
face_type=FaceType.fromString (dflimg.get_face_type()),
|
||||
shape=dflimg.get_shape(),
|
||||
landmarks=dflimg.get_landmarks(),
|
||||
ie_polys=dflimg.get_ie_polys(),
|
||||
pitch=pitch,
|
||||
yaw=yaw,
|
||||
pitch_yaw_roll=dflimg.get_pitch_yaw_roll(),
|
||||
source_filename=dflimg.get_source_filename(),
|
||||
fanseg_mask_exist=dflimg.get_fanseg_mask() is not None, ) )
|
||||
except:
|
||||
|
|
@ -85,12 +82,12 @@ class SampleLoader:
|
|||
|
||||
return sample_list
|
||||
|
||||
# @staticmethod
|
||||
# def upgradeToFaceTemporalSortedSamples( samples ):
|
||||
# new_s = [ (s, s.source_filename) for s in samples]
|
||||
# new_s = sorted(new_s, key=operator.itemgetter(1))
|
||||
@staticmethod
|
||||
def upgradeToFaceTemporalSortedSamples( samples ):
|
||||
new_s = [ (s, s.source_filename) for s in samples]
|
||||
new_s = sorted(new_s, key=operator.itemgetter(1))
|
||||
|
||||
# return [ s[0] for s in new_s]
|
||||
return [ s[0] for s in new_s]
|
||||
|
||||
@staticmethod
|
||||
def upgradeToFaceYawSortedSamples( samples ):
|
||||
|
|
|
|||
|
|
@ -13,9 +13,10 @@ class SampleProcessor(object):
|
|||
WARPED_TRANSFORMED = 0x00000004,
|
||||
TRANSFORMED = 0x00000008,
|
||||
LANDMARKS_ARRAY = 0x00000010, #currently unused
|
||||
PITCH_YAW_ROLL = 0x00000020,
|
||||
|
||||
RANDOM_CLOSE = 0x00000020, #currently unused
|
||||
MORPH_TO_RANDOM_CLOSE = 0x00000040, #currently unused
|
||||
RANDOM_CLOSE = 0x00000040, #currently unused
|
||||
MORPH_TO_RANDOM_CLOSE = 0x00000080, #currently unused
|
||||
|
||||
FACE_TYPE_HALF = 0x00000100,
|
||||
FACE_TYPE_FULL = 0x00000200,
|
||||
|
|
@ -77,7 +78,7 @@ class SampleProcessor(object):
|
|||
outputs = []
|
||||
for sample_type in output_sample_types:
|
||||
f = sample_type[0]
|
||||
size = sample_type[1]
|
||||
size = 0 if len (sample_type) < 2 else sample_type[1]
|
||||
random_sub_size = 0 if len (sample_type) < 3 else min( sample_type[2] , size)
|
||||
|
||||
if f & SPTF.SOURCE != 0:
|
||||
|
|
@ -90,6 +91,8 @@ class SampleProcessor(object):
|
|||
img_type = 3
|
||||
elif f & SPTF.LANDMARKS_ARRAY != 0:
|
||||
img_type = 4
|
||||
elif f & SPTF.PITCH_YAW_ROLL != 0:
|
||||
img_type = 5
|
||||
else:
|
||||
raise ValueError ('expected SampleTypeFlags type')
|
||||
|
||||
|
|
@ -121,6 +124,16 @@ class SampleProcessor(object):
|
|||
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 == 5:
|
||||
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
|
||||
|
||||
img = (pitch, yaw, roll)
|
||||
else:
|
||||
if images[img_type][face_mask_type] is None:
|
||||
if img_type >= 10 and img_type <= 19: #RANDOM_CLOSE
|
||||
|
|
|
|||
|
|
@ -4,4 +4,5 @@ from .SampleLoader import SampleLoader
|
|||
from .SampleProcessor import SampleProcessor
|
||||
from .SampleGeneratorBase import SampleGeneratorBase
|
||||
from .SampleGeneratorFace import SampleGeneratorFace
|
||||
from .SampleGeneratorFaceTemporal import SampleGeneratorFaceTemporal
|
||||
from .SampleGeneratorImageTemporal import SampleGeneratorImageTemporal
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue