mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-14 02:37:00 -07:00
removing trailing spaces
This commit is contained in:
parent
fa4e579b95
commit
a3df04999c
61 changed files with 2110 additions and 2103 deletions
|
@ -5,17 +5,17 @@ from utils.cv2_utils import *
|
|||
|
||||
class SampleType(IntEnum):
|
||||
IMAGE = 0 #raw image
|
||||
|
||||
|
||||
FACE_BEGIN = 1
|
||||
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_WITH_CLOSE_TO_SELF = 4
|
||||
FACE_END = 4
|
||||
|
||||
|
||||
QTY = 5
|
||||
|
||||
class Sample(object):
|
||||
|
||||
class Sample(object):
|
||||
def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, pitch=None, yaw=None, mirror=None, close_target_list=None):
|
||||
self.sample_type = sample_type if sample_type is not None else SampleType.IMAGE
|
||||
self.filename = filename
|
||||
|
@ -26,19 +26,19 @@ class Sample(object):
|
|||
self.yaw = yaw
|
||||
self.mirror = mirror
|
||||
self.close_target_list = close_target_list
|
||||
|
||||
|
||||
def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, pitch=None, yaw=None, mirror=None, close_target_list=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,
|
||||
face_type=face_type if face_type is not None else self.face_type,
|
||||
shape=shape if shape is not None else self.shape,
|
||||
landmarks=landmarks if landmarks is not None else self.landmarks.copy(),
|
||||
pitch=pitch if pitch is not None else self.pitch,
|
||||
yaw=yaw if yaw is not None else self.yaw,
|
||||
mirror=mirror if mirror is not None else self.mirror,
|
||||
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,
|
||||
face_type=face_type if face_type is not None else self.face_type,
|
||||
shape=shape if shape is not None else self.shape,
|
||||
landmarks=landmarks if landmarks is not None else self.landmarks.copy(),
|
||||
pitch=pitch if pitch is not None else self.pitch,
|
||||
yaw=yaw if yaw is not None else self.yaw,
|
||||
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)
|
||||
|
||||
|
||||
def load_bgr(self):
|
||||
img = cv2_imread (self.filename).astype(np.float32) / 255.0
|
||||
if self.mirror:
|
||||
|
@ -48,4 +48,4 @@ class Sample(object):
|
|||
def get_random_close_target_sample(self):
|
||||
if self.close_target_list is None:
|
||||
return None
|
||||
return self.close_target_list[randint (0, len(self.close_target_list)-1)]
|
||||
return self.close_target_list[randint (0, len(self.close_target_list)-1)]
|
||||
|
|
|
@ -4,22 +4,21 @@ from pathlib import Path
|
|||
You can implement your own SampleGenerator
|
||||
'''
|
||||
class SampleGeneratorBase(object):
|
||||
|
||||
|
||||
|
||||
|
||||
def __init__ (self, samples_path, debug, batch_size):
|
||||
if samples_path is None:
|
||||
raise Exception('samples_path is None')
|
||||
|
||||
|
||||
self.samples_path = Path(samples_path)
|
||||
self.debug = debug
|
||||
self.batch_size = 1 if self.debug else batch_size
|
||||
|
||||
self.batch_size = 1 if self.debug else batch_size
|
||||
|
||||
#overridable
|
||||
def __iter__(self):
|
||||
#implement your own iterator
|
||||
return self
|
||||
|
||||
|
||||
def __next__(self):
|
||||
#implement your own iterator
|
||||
return None
|
||||
|
||||
|
|
|
@ -12,9 +12,9 @@ from samples import SampleLoader
|
|||
from samples import SampleGeneratorBase
|
||||
|
||||
'''
|
||||
arg
|
||||
arg
|
||||
output_sample_types = [
|
||||
[SampleProcessor.TypeFlags, size, (optional)random_sub_size] ,
|
||||
[SampleProcessor.TypeFlags, size, (optional)random_sub_size] ,
|
||||
...
|
||||
]
|
||||
'''
|
||||
|
@ -26,7 +26,7 @@ class SampleGeneratorFace(SampleGeneratorBase):
|
|||
self.add_sample_idx = add_sample_idx
|
||||
self.add_pitch = add_pitch
|
||||
self.add_yaw = add_yaw
|
||||
|
||||
|
||||
if sort_by_yaw_target_samples_path is not None:
|
||||
self.sample_type = SampleType.FACE_YAW_SORTED_AS_TARGET
|
||||
elif sort_by_yaw:
|
||||
|
@ -34,9 +34,9 @@ class SampleGeneratorFace(SampleGeneratorBase):
|
|||
elif with_close_to_self:
|
||||
self.sample_type = SampleType.FACE_WITH_CLOSE_TO_SELF
|
||||
else:
|
||||
self.sample_type = SampleType.FACE
|
||||
|
||||
self.samples = SampleLoader.load (self.sample_type, self.samples_path, sort_by_yaw_target_samples_path)
|
||||
self.sample_type = SampleType.FACE
|
||||
|
||||
self.samples = SampleLoader.load (self.sample_type, self.samples_path, sort_by_yaw_target_samples_path)
|
||||
|
||||
if self.debug:
|
||||
self.generators_count = 1
|
||||
|
@ -46,24 +46,24 @@ class SampleGeneratorFace(SampleGeneratorBase):
|
|||
self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, i ) for i in range(self.generators_count) ]
|
||||
|
||||
self.generators_sq = [ multiprocessing.Queue() for _ in range(self.generators_count) ]
|
||||
|
||||
|
||||
self.generator_counter = -1
|
||||
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
|
||||
def __next__(self):
|
||||
self.generator_counter += 1
|
||||
generator = self.generators[self.generator_counter % len(self.generators) ]
|
||||
return next(generator)
|
||||
|
||||
|
||||
#forces to repeat these sample idxs as fast as possible
|
||||
#currently unused
|
||||
def repeat_sample_idxs(self, idxs): # [ idx, ... ]
|
||||
#send idxs list to all sub generators.
|
||||
for gen_sq in self.generators_sq:
|
||||
gen_sq.put (idxs)
|
||||
|
||||
gen_sq.put (idxs)
|
||||
|
||||
def batch_func(self, generator_id):
|
||||
gen_sq = self.generators_sq[generator_id]
|
||||
samples = self.samples
|
||||
|
@ -73,11 +73,11 @@ class SampleGeneratorFace(SampleGeneratorBase):
|
|||
|
||||
if len(samples_idxs) == 0:
|
||||
raise ValueError('No training data provided.')
|
||||
|
||||
|
||||
if self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
|
||||
if all ( [ samples[idx] == None for idx in samples_idxs] ):
|
||||
raise ValueError('Not enough training data. Gather more faces!')
|
||||
|
||||
|
||||
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:
|
||||
|
@ -89,25 +89,25 @@ class SampleGeneratorFace(SampleGeneratorBase):
|
|||
idxs = gen_sq.get()
|
||||
for idx in idxs:
|
||||
if idx in samples_idxs:
|
||||
repeat_samples_idxs.append(idx)
|
||||
repeat_samples_idxs.append(idx)
|
||||
|
||||
batches = None
|
||||
for n_batch in range(self.batch_size):
|
||||
while True:
|
||||
sample = None
|
||||
|
||||
|
||||
if len(repeat_samples_idxs) > 0:
|
||||
idx = repeat_samples_idxs.pop()
|
||||
if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
|
||||
idx = repeat_samples_idxs.pop()
|
||||
if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
|
||||
sample = samples[idx]
|
||||
elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
|
||||
sample = samples[(idx >> 16) & 0xFFFF][idx & 0xFFFF]
|
||||
else:
|
||||
else:
|
||||
if self.sample_type == SampleType.FACE or self.sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
|
||||
if len(shuffle_idxs) == 0:
|
||||
shuffle_idxs = samples_idxs.copy()
|
||||
np.random.shuffle(shuffle_idxs)
|
||||
|
||||
|
||||
idx = shuffle_idxs.pop()
|
||||
sample = samples[ idx ]
|
||||
|
||||
|
@ -120,18 +120,18 @@ class SampleGeneratorFace(SampleGeneratorBase):
|
|||
if samples[idx] != None:
|
||||
if len(shuffle_idxs_2D[idx]) == 0:
|
||||
shuffle_idxs_2D[idx] = random.sample( range(len(samples[idx])), len(samples[idx]) )
|
||||
|
||||
idx2 = shuffle_idxs_2D[idx].pop()
|
||||
|
||||
idx2 = shuffle_idxs_2D[idx].pop()
|
||||
sample = samples[idx][idx2]
|
||||
|
||||
|
||||
idx = (idx << 16) | (idx2 & 0xFFFF)
|
||||
|
||||
if sample is not None:
|
||||
|
||||
if sample is not None:
|
||||
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 type(x) != tuple and type(x) != list:
|
||||
raise Exception('SampleProcessor.process returns NOT tuple/list')
|
||||
|
||||
|
@ -144,23 +144,23 @@ class SampleGeneratorFace(SampleGeneratorBase):
|
|||
batches += [ [] ]
|
||||
i_pitch = len(batches)-1
|
||||
if self.add_yaw:
|
||||
batches += [ [] ]
|
||||
batches += [ [] ]
|
||||
i_yaw = len(batches)-1
|
||||
|
||||
|
||||
for i in range(len(x)):
|
||||
batches[i].append ( x[i] )
|
||||
|
||||
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]
|
||||
|
|
|
@ -11,36 +11,36 @@ from samples import SampleLoader
|
|||
from samples import SampleGeneratorBase
|
||||
|
||||
'''
|
||||
output_sample_types = [
|
||||
[SampleProcessor.TypeFlags, size, (optional)random_sub_size] ,
|
||||
output_sample_types = [
|
||||
[SampleProcessor.TypeFlags, size, (optional)random_sub_size] ,
|
||||
...
|
||||
]
|
||||
'''
|
||||
class SampleGeneratorImageTemporal(SampleGeneratorBase):
|
||||
def __init__ (self, samples_path, debug, batch_size, temporal_image_count, sample_process_options=SampleProcessor.Options(), output_sample_types=[], **kwargs):
|
||||
super().__init__(samples_path, debug, batch_size)
|
||||
|
||||
|
||||
self.temporal_image_count = temporal_image_count
|
||||
self.sample_process_options = sample_process_options
|
||||
self.output_sample_types = output_sample_types
|
||||
|
||||
self.samples = SampleLoader.load (SampleType.IMAGE, self.samples_path)
|
||||
|
||||
self.samples = SampleLoader.load (SampleType.IMAGE, self.samples_path)
|
||||
|
||||
self.generator_samples = [ self.samples ]
|
||||
self.generators = [iter_utils.ThisThreadGenerator ( self.batch_func, 0 )] if self.debug else \
|
||||
[iter_utils.SubprocessGenerator ( self.batch_func, 0 )]
|
||||
|
||||
|
||||
self.generator_counter = -1
|
||||
|
||||
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
|
||||
def __next__(self):
|
||||
self.generator_counter += 1
|
||||
generator = self.generators[self.generator_counter % len(self.generators) ]
|
||||
return next(generator)
|
||||
|
||||
def batch_func(self, generator_id):
|
||||
|
||||
def batch_func(self, generator_id):
|
||||
samples = self.generator_samples[generator_id]
|
||||
samples_len = len(samples)
|
||||
if samples_len == 0:
|
||||
|
@ -48,20 +48,20 @@ class SampleGeneratorImageTemporal(SampleGeneratorBase):
|
|||
|
||||
if samples_len - self.temporal_image_count < 0:
|
||||
raise ValueError('Not enough samples to fit temporal line.')
|
||||
|
||||
|
||||
shuffle_idxs = []
|
||||
samples_sub_len = samples_len - self.temporal_image_count + 1
|
||||
|
||||
while True:
|
||||
|
||||
|
||||
while True:
|
||||
|
||||
batches = None
|
||||
for n_batch in range(self.batch_size):
|
||||
|
||||
if len(shuffle_idxs) == 0:
|
||||
shuffle_idxs = random.sample( range(samples_sub_len), samples_sub_len )
|
||||
|
||||
|
||||
idx = shuffle_idxs.pop()
|
||||
|
||||
|
||||
temporal_samples = []
|
||||
|
||||
for i in range( self.temporal_image_count ):
|
||||
|
@ -70,11 +70,11 @@ class SampleGeneratorImageTemporal(SampleGeneratorBase):
|
|||
temporal_samples += 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(temporal_samples)) ]
|
||||
|
||||
|
||||
for i in range(len(temporal_samples)):
|
||||
batches[i].append ( temporal_samples[i] )
|
||||
|
||||
|
||||
yield [ np.array(batch) for batch in batches]
|
||||
|
|
|
@ -17,44 +17,44 @@ from interact import interact as io
|
|||
|
||||
class SampleLoader:
|
||||
cache = dict()
|
||||
|
||||
|
||||
@staticmethod
|
||||
def load(sample_type, samples_path, target_samples_path=None):
|
||||
cache = SampleLoader.cache
|
||||
|
||||
|
||||
if str(samples_path) not in cache.keys():
|
||||
cache[str(samples_path)] = [None]*SampleType.QTY
|
||||
|
||||
|
||||
datas = cache[str(samples_path)]
|
||||
|
||||
if sample_type == SampleType.IMAGE:
|
||||
if datas[sample_type] is None:
|
||||
if datas[sample_type] is None:
|
||||
datas[sample_type] = [ Sample(filename=filename) for filename in io.progress_bar_generator( Path_utils.get_image_paths(samples_path), "Loading") ]
|
||||
|
||||
elif sample_type == SampleType.FACE:
|
||||
if datas[sample_type] is None:
|
||||
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_YAW_SORTED:
|
||||
if datas[sample_type] is None:
|
||||
datas[sample_type] = SampleLoader.upgradeToFaceYawSortedSamples( SampleLoader.load(SampleType.FACE, samples_path) )
|
||||
|
||||
elif sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
|
||||
|
||||
elif sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
|
||||
if datas[sample_type] is None:
|
||||
if target_samples_path is None:
|
||||
raise Exception('target_samples_path is None for FACE_YAW_SORTED_AS_TARGET')
|
||||
datas[sample_type] = SampleLoader.upgradeToFaceYawSortedAsTargetSamples( SampleLoader.load(SampleType.FACE_YAW_SORTED, samples_path), SampleLoader.load(SampleType.FACE_YAW_SORTED, target_samples_path) )
|
||||
elif sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
|
||||
elif sample_type == SampleType.FACE_WITH_CLOSE_TO_SELF:
|
||||
if datas[sample_type] is None:
|
||||
datas[sample_type] = SampleLoader.upgradeToFaceCloseToSelfSamples( SampleLoader.load(SampleType.FACE, samples_path) )
|
||||
|
||||
|
||||
|
||||
|
||||
return datas[sample_type]
|
||||
|
||||
|
||||
@staticmethod
|
||||
def upgradeToFaceSamples ( samples ):
|
||||
sample_list = []
|
||||
|
||||
|
||||
for s in io.progress_bar_generator(samples, "Loading"):
|
||||
s_filename_path = Path(s.filename)
|
||||
try:
|
||||
|
@ -64,57 +64,57 @@ class SampleLoader:
|
|||
dflimg = DFLJPG.load ( str(s_filename_path) )
|
||||
else:
|
||||
dflimg = None
|
||||
|
||||
|
||||
if dflimg is None:
|
||||
print ("%s is not a dfl image file required for training" % (s_filename_path.name) )
|
||||
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(),
|
||||
shape=dflimg.get_shape(),
|
||||
landmarks=dflimg.get_landmarks(),
|
||||
pitch=pitch,
|
||||
yaw=yaw) )
|
||||
except:
|
||||
print ("Unable to load %s , error: %s" % (str(s_filename_path), traceback.format_exc() ) )
|
||||
|
||||
return sample_list
|
||||
|
||||
|
||||
return sample_list
|
||||
|
||||
@staticmethod
|
||||
def upgradeToFaceCloseToSelfSamples (samples):
|
||||
yaw_samples = SampleLoader.upgradeToFaceYawSortedSamples(samples)
|
||||
yaw_samples_len = len(yaw_samples)
|
||||
|
||||
|
||||
sample_list = []
|
||||
for i in io.progress_bar_generator( range(yaw_samples_len), "Sorting"):
|
||||
if yaw_samples[i] is not None:
|
||||
for s in yaw_samples[i]:
|
||||
s_t = []
|
||||
|
||||
for n in range(2000):
|
||||
|
||||
for n in range(2000):
|
||||
yaw_idx = np.clip ( i-10 +np.random.randint(20), 0, yaw_samples_len-1 )
|
||||
if yaw_samples[yaw_idx] is None:
|
||||
continue
|
||||
|
||||
|
||||
yaw_idx_samples_len = len(yaw_samples[yaw_idx])
|
||||
|
||||
|
||||
yaw_idx_sample = yaw_samples[yaw_idx][ np.random.randint(yaw_idx_samples_len) ]
|
||||
if s.filename == yaw_idx_sample.filename:
|
||||
continue
|
||||
|
||||
|
||||
s_t.append ( yaw_idx_sample )
|
||||
if len(s_t) >= 50:
|
||||
break
|
||||
|
||||
|
||||
if len(s_t) == 0:
|
||||
s_t = [s]
|
||||
|
||||
|
||||
sample_list.append( s.copy_and_set(close_target_list = s_t) )
|
||||
|
||||
return sample_list
|
||||
|
||||
|
||||
@staticmethod
|
||||
def upgradeToFaceYawSortedSamples( samples ):
|
||||
|
||||
|
@ -123,50 +123,50 @@ class SampleLoader:
|
|||
diff_rot_per_grad = abs(highest_yaw-lowest_yaw) / gradations
|
||||
|
||||
yaws_sample_list = [None]*gradations
|
||||
|
||||
|
||||
for i in io.progress_bar_generator(range(gradations), "Sorting"):
|
||||
yaw = lowest_yaw + i*diff_rot_per_grad
|
||||
next_yaw = lowest_yaw + (i+1)*diff_rot_per_grad
|
||||
|
||||
yaw_samples = []
|
||||
for s in samples:
|
||||
for s in samples:
|
||||
s_yaw = s.yaw
|
||||
if (i == 0 and s_yaw < next_yaw) or \
|
||||
(i < gradations-1 and s_yaw >= yaw and s_yaw < next_yaw) or \
|
||||
(i == gradations-1 and s_yaw >= yaw):
|
||||
yaw_samples.append ( s.copy_and_set(sample_type=SampleType.FACE_YAW_SORTED) )
|
||||
|
||||
|
||||
if len(yaw_samples) > 0:
|
||||
yaws_sample_list[i] = yaw_samples
|
||||
|
||||
|
||||
return yaws_sample_list
|
||||
|
||||
|
||||
@staticmethod
|
||||
def upgradeToFaceYawSortedAsTargetSamples (s, t):
|
||||
l = len(s)
|
||||
if l != len(t):
|
||||
raise Exception('upgradeToFaceYawSortedAsTargetSamples() s_len != t_len')
|
||||
b = l // 2
|
||||
|
||||
|
||||
s_idxs = np.argwhere ( np.array ( [ 1 if x != None else 0 for x in s] ) == 1 )[:,0]
|
||||
t_idxs = np.argwhere ( np.array ( [ 1 if x != None else 0 for x in t] ) == 1 )[:,0]
|
||||
|
||||
new_s = [None]*l
|
||||
|
||||
|
||||
new_s = [None]*l
|
||||
|
||||
for t_idx in t_idxs:
|
||||
search_idxs = []
|
||||
search_idxs = []
|
||||
for i in range(0,l):
|
||||
search_idxs += [t_idx - i, (l-t_idx-1) - i, t_idx + i, (l-t_idx-1) + i]
|
||||
|
||||
for search_idx in search_idxs:
|
||||
for search_idx in search_idxs:
|
||||
if search_idx in s_idxs:
|
||||
mirrored = ( t_idx != search_idx and ((t_idx < b and search_idx >= b) or (search_idx < b and t_idx >= b)) )
|
||||
new_s[t_idx] = [ sample.copy_and_set(sample_type=SampleType.FACE_YAW_SORTED_AS_TARGET,
|
||||
mirror=True,
|
||||
yaw=-sample.yaw,
|
||||
mirror=True,
|
||||
yaw=-sample.yaw,
|
||||
landmarks=LandmarksProcessor.mirror_landmarks (sample.landmarks, sample.shape[1] ))
|
||||
for sample in s[search_idx]
|
||||
] if mirrored else s[search_idx]
|
||||
for sample in s[search_idx]
|
||||
] if mirrored else s[search_idx]
|
||||
break
|
||||
|
||||
return new_s
|
||||
|
||||
return new_s
|
||||
|
|
|
@ -13,61 +13,61 @@ class SampleProcessor(object):
|
|||
WARPED_TRANSFORMED = 0x00000004,
|
||||
TRANSFORMED = 0x00000008,
|
||||
LANDMARKS_ARRAY = 0x00000010, #currently unused
|
||||
|
||||
|
||||
RANDOM_CLOSE = 0x00000020,
|
||||
MORPH_TO_RANDOM_CLOSE = 0x00000040,
|
||||
|
||||
|
||||
FACE_ALIGN_HALF = 0x00000100,
|
||||
FACE_ALIGN_FULL = 0x00000200,
|
||||
FACE_ALIGN_HEAD = 0x00000400,
|
||||
FACE_ALIGN_AVATAR = 0x00000800,
|
||||
|
||||
FACE_ALIGN_AVATAR = 0x00000800,
|
||||
|
||||
FACE_MASK_FULL = 0x00001000,
|
||||
FACE_MASK_EYES = 0x00002000,
|
||||
|
||||
|
||||
MODE_BGR = 0x01000000, #BGR
|
||||
MODE_G = 0x02000000, #Grayscale
|
||||
MODE_GGG = 0x04000000, #3xGrayscale
|
||||
MODE_GGG = 0x04000000, #3xGrayscale
|
||||
MODE_M = 0x08000000, #mask only
|
||||
MODE_BGR_SHUFFLE = 0x10000000, #BGR shuffle
|
||||
|
||||
class Options(object):
|
||||
|
||||
class Options(object):
|
||||
def __init__(self, random_flip = True, normalize_tanh = False, rotation_range=[-10,10], scale_range=[-0.05, 0.05], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05]):
|
||||
self.random_flip = random_flip
|
||||
self.random_flip = random_flip
|
||||
self.normalize_tanh = normalize_tanh
|
||||
self.rotation_range = rotation_range
|
||||
self.scale_range = scale_range
|
||||
self.tx_range = tx_range
|
||||
self.ty_range = ty_range
|
||||
|
||||
self.ty_range = ty_range
|
||||
|
||||
@staticmethod
|
||||
def process (sample, sample_process_options, output_sample_types, debug):
|
||||
sample_bgr = sample.load_bgr()
|
||||
h,w,c = sample_bgr.shape
|
||||
|
||||
is_face_sample = sample.landmarks is not None
|
||||
|
||||
is_face_sample = sample.landmarks is not None
|
||||
|
||||
if debug and is_face_sample:
|
||||
LandmarksProcessor.draw_landmarks (sample_bgr, sample.landmarks, (0, 1, 0))
|
||||
|
||||
|
||||
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_bgr = close_sample.load_bgr() if close_sample is not None else None
|
||||
|
||||
|
||||
if debug and close_sample_bgr is not None:
|
||||
LandmarksProcessor.draw_landmarks (close_sample_bgr, close_sample.landmarks, (0, 1, 0))
|
||||
|
||||
LandmarksProcessor.draw_landmarks (close_sample_bgr, close_sample.landmarks, (0, 1, 0))
|
||||
|
||||
params = image_utils.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 )
|
||||
|
||||
images = [[None]*3 for _ in range(30)]
|
||||
|
||||
|
||||
sample_rnd_seed = np.random.randint(0x80000000)
|
||||
|
||||
outputs = []
|
||||
|
||||
outputs = []
|
||||
for sample_type in output_sample_types:
|
||||
f = sample_type[0]
|
||||
size = sample_type[1]
|
||||
random_sub_size = 0 if len (sample_type) < 3 else min( sample_type[2] , size)
|
||||
|
||||
|
||||
if f & SampleProcessor.TypeFlags.SOURCE != 0:
|
||||
img_type = 0
|
||||
elif f & SampleProcessor.TypeFlags.WARPED != 0:
|
||||
|
@ -77,53 +77,53 @@ class SampleProcessor(object):
|
|||
elif f & SampleProcessor.TypeFlags.TRANSFORMED != 0:
|
||||
img_type = 3
|
||||
elif f & SampleProcessor.TypeFlags.LANDMARKS_ARRAY != 0:
|
||||
img_type = 4
|
||||
img_type = 4
|
||||
else:
|
||||
raise ValueError ('expected SampleTypeFlags type')
|
||||
|
||||
|
||||
if f & SampleProcessor.TypeFlags.RANDOM_CLOSE != 0:
|
||||
img_type += 10
|
||||
elif f & SampleProcessor.TypeFlags.MORPH_TO_RANDOM_CLOSE != 0:
|
||||
img_type += 20
|
||||
|
||||
|
||||
face_mask_type = 0
|
||||
if f & SampleProcessor.TypeFlags.FACE_MASK_FULL != 0:
|
||||
face_mask_type = 1
|
||||
face_mask_type = 1
|
||||
elif f & SampleProcessor.TypeFlags.FACE_MASK_EYES != 0:
|
||||
face_mask_type = 2
|
||||
|
||||
|
||||
target_face_type = -1
|
||||
if f & SampleProcessor.TypeFlags.FACE_ALIGN_HALF != 0:
|
||||
target_face_type = FaceType.HALF
|
||||
target_face_type = FaceType.HALF
|
||||
elif f & SampleProcessor.TypeFlags.FACE_ALIGN_FULL != 0:
|
||||
target_face_type = FaceType.FULL
|
||||
elif f & SampleProcessor.TypeFlags.FACE_ALIGN_HEAD != 0:
|
||||
target_face_type = FaceType.HEAD
|
||||
elif f & SampleProcessor.TypeFlags.FACE_ALIGN_AVATAR != 0:
|
||||
target_face_type = FaceType.AVATAR
|
||||
|
||||
|
||||
if img_type == 4:
|
||||
l = sample.landmarks
|
||||
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
|
||||
else:
|
||||
else:
|
||||
if images[img_type][face_mask_type] is None:
|
||||
if img_type >= 10 and img_type <= 19: #RANDOM_CLOSE
|
||||
img_type -= 10
|
||||
img = close_sample_bgr
|
||||
cur_sample = close_sample
|
||||
|
||||
|
||||
elif img_type >= 20 and img_type <= 29: #MORPH_TO_RANDOM_CLOSE
|
||||
img_type -= 20
|
||||
res = sample.shape[0]
|
||||
|
||||
s_landmarks = sample.landmarks.copy()
|
||||
d_landmarks = close_sample.landmarks.copy()
|
||||
idxs = list(range(len(s_landmarks)))
|
||||
|
||||
s_landmarks = sample.landmarks.copy()
|
||||
d_landmarks = close_sample.landmarks.copy()
|
||||
idxs = list(range(len(s_landmarks)))
|
||||
#remove landmarks near boundaries
|
||||
for i in idxs[:]:
|
||||
s_l = s_landmarks[i]
|
||||
s_l = s_landmarks[i]
|
||||
d_l = d_landmarks[i]
|
||||
if s_l[0] < 5 or s_l[1] < 5 or s_l[0] >= res-5 or s_l[1] >= res-5 or \
|
||||
d_l[0] < 5 or d_l[1] < 5 or d_l[0] >= res-5 or d_l[1] >= res-5:
|
||||
|
@ -139,39 +139,39 @@ class SampleProcessor(object):
|
|||
diff_l = np.abs(s_l - s_l_2)
|
||||
if np.sqrt(diff_l.dot(diff_l)) < 5:
|
||||
idxs.remove(i)
|
||||
break
|
||||
break
|
||||
s_landmarks = s_landmarks[idxs]
|
||||
d_landmarks = d_landmarks[idxs]
|
||||
s_landmarks = np.concatenate ( [s_landmarks, [ [0,0], [ res // 2, 0], [ res-1, 0], [0, res//2], [res-1, res//2] ,[0,res-1] ,[res//2, res-1] ,[res-1,res-1] ] ] )
|
||||
s_landmarks = np.concatenate ( [s_landmarks, [ [0,0], [ res // 2, 0], [ res-1, 0], [0, res//2], [res-1, res//2] ,[0,res-1] ,[res//2, res-1] ,[res-1,res-1] ] ] )
|
||||
d_landmarks = np.concatenate ( [d_landmarks, [ [0,0], [ res // 2, 0], [ res-1, 0], [0, res//2], [res-1, res//2] ,[0,res-1] ,[res//2, res-1] ,[res-1,res-1] ] ] )
|
||||
img = image_utils.morph_by_points (sample_bgr, s_landmarks, d_landmarks)
|
||||
cur_sample = close_sample
|
||||
else:
|
||||
img = sample_bgr
|
||||
cur_sample = sample
|
||||
|
||||
|
||||
if is_face_sample:
|
||||
if face_mask_type == 1:
|
||||
img = np.concatenate( (img, LandmarksProcessor.get_image_hull_mask (img.shape, cur_sample.landmarks) ), -1 )
|
||||
img = np.concatenate( (img, LandmarksProcessor.get_image_hull_mask (img.shape, cur_sample.landmarks) ), -1 )
|
||||
elif face_mask_type == 2:
|
||||
mask = LandmarksProcessor.get_image_eye_mask (img.shape, cur_sample.landmarks)
|
||||
mask = np.expand_dims (cv2.blur (mask, ( w // 32, w // 32 ) ), -1)
|
||||
mask[mask > 0.0] = 1.0
|
||||
img = np.concatenate( (img, mask ), -1 )
|
||||
img = np.concatenate( (img, mask ), -1 )
|
||||
|
||||
images[img_type][face_mask_type] = image_utils.warp_by_params (params, img, (img_type==1 or img_type==2), (img_type==2 or img_type==3), img_type != 0, face_mask_type == 0)
|
||||
|
||||
|
||||
img = images[img_type][face_mask_type]
|
||||
|
||||
|
||||
if is_face_sample and target_face_type != -1:
|
||||
if target_face_type > 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_face_type) )
|
||||
img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, size, target_face_type), (size,size), flags=cv2.INTER_CUBIC )
|
||||
else:
|
||||
img = cv2.resize( img, (size,size), cv2.INTER_CUBIC )
|
||||
|
||||
|
||||
if random_sub_size != 0:
|
||||
sub_size = size - random_sub_size
|
||||
sub_size = size - random_sub_size
|
||||
rnd_state = np.random.RandomState (sample_rnd_seed+random_sub_size)
|
||||
start_x = rnd_state.randint(sub_size+1)
|
||||
start_y = rnd_state.randint(sub_size+1)
|
||||
|
@ -195,7 +195,7 @@ class SampleProcessor(object):
|
|||
img = img_mask
|
||||
else:
|
||||
raise ValueError ('expected SampleTypeFlags mode')
|
||||
|
||||
|
||||
if not debug:
|
||||
if sample_process_options.normalize_tanh:
|
||||
img = np.clip (img * 2.0 - 1.0, -1.0, 1.0)
|
||||
|
@ -213,6 +213,6 @@ class SampleProcessor(object):
|
|||
elif output.shape[2] == 4:
|
||||
result += [output[...,0:3]*output[...,3:4],]
|
||||
|
||||
return result
|
||||
return result
|
||||
else:
|
||||
return outputs
|
||||
return outputs
|
||||
|
|
|
@ -4,4 +4,4 @@ from .SampleLoader import SampleLoader
|
|||
from .SampleProcessor import SampleProcessor
|
||||
from .SampleGeneratorBase import SampleGeneratorBase
|
||||
from .SampleGeneratorFace import SampleGeneratorFace
|
||||
from .SampleGeneratorImageTemporal import SampleGeneratorImageTemporal
|
||||
from .SampleGeneratorImageTemporal import SampleGeneratorImageTemporal
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue