github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-08-21 05:53:24 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

rename samples to samplelib

Browse source
This commit is contained in:
iperov 2019-03-27 10:44:13 +04:00
commit 3cf3bb786e
14 changed files with 9 additions and 15 deletions
Download patch file Download diff file Expand all files Collapse all files

51
samplelib/Sample.py Normal file
View file

@ -0,0 +1,51 @@
from enum import IntEnum
import cv2
import numpy as np
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):
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
self.face_type = face_type
self.shape = shape
self.landmarks = np.array(landmarks) if landmarks is not None else None
self.pitch = pitch
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,
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:
img = img[:,::-1].copy()
return img
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)]

24
samplelib/SampleGeneratorBase.py Normal file
View file

@ -0,0 +1,24 @@
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
#overridable
def __iter__(self):
#implement your own iterator
return self
def __next__(self):
#implement your own iterator
return None

171
samplelib/SampleGeneratorFace.py Normal file
View file

@ -0,0 +1,171 @@
import traceback
import numpy as np
import cv2
import multiprocessing
from utils import iter_utils
from facelib import LandmarksProcessor
from samplelib import SampleType, SampleProcessor, SampleLoader, 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, with_close_to_self=False, 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):
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
if sort_by_yaw_target_samples_path is not None:
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.samples = SampleLoader.load (self.sample_type, self.samples_path, sort_by_yaw_target_samples_path)
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
if self.debug:
self.generators_count = 1
self.generators = [iter_utils.ThisThreadGenerator ( self.batch_func, 0 )]
else:
self.generators_count = min ( generators_count, len(self.samples) )
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)
def batch_func(self, generator_id):
gen_sq = self.generators_sq[generator_id]
if self.generators_random_seed is not None:
np.random.seed ( self.generators_random_seed[generator_id] )
samples = self.samples
samples_len = len(samples)
samples_idxs = [ *range(samples_len) ] [generator_id::self.generators_count]
repeat_samples_idxs = []
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:
shuffle_idxs = []
shuffle_idxs_2D = [[]]*samples_len
while True:
while not gen_sq.empty():
idxs = gen_sq.get()
for idx in idxs:
if idx in samples_idxs:
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:
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:
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 ]
elif self.sample_type == SampleType.FACE_YAW_SORTED or self.sample_type == SampleType.FACE_YAW_SORTED_AS_TARGET:
if len(shuffle_idxs) == 0:
shuffle_idxs = samples_idxs.copy()
np.random.shuffle(shuffle_idxs)
idx = shuffle_idxs.pop()
if samples[idx] != None:
if len(shuffle_idxs_2D[idx]) == 0:
a = shuffle_idxs_2D[idx] = [ *range(len(samples[idx])) ]
np.random.shuffle (a)
idx2 = shuffle_idxs_2D[idx].pop()
sample = samples[idx][idx2]
idx = (idx << 16) | (idx2 & 0xFFFF)
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')
if batches is None:
batches = [ [] for _ in range(len(x)) ]
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] )
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]

77
samplelib/SampleGeneratorImageTemporal.py Normal file
View file

@ -0,0 +1,77 @@
import traceback
import numpy as np
import cv2
from utils import iter_utils
from samplelib import SampleType, SampleProcessor, SampleLoader, SampleGeneratorBase
'''
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.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):
samples = self.generator_samples[generator_id]
samples_len = len(samples)
if samples_len == 0:
raise ValueError('No training data provided.')
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:
batches = None
for n_batch in range(self.batch_size):
if len(shuffle_idxs) == 0:
shuffle_idxs = [ *range(samples_sub_len) ]
np.random.shuffle (shuffle_idxs)
idx = shuffle_idxs.pop()
temporal_samples = []
for i in range( self.temporal_image_count ):
sample = samples[ idx+i ]
try:
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]

172
samplelib/SampleLoader.py Normal file
View file

@ -0,0 +1,172 @@
import traceback
from enum import IntEnum
import cv2
import numpy as np
from pathlib import Path
from utils import Path_utils
from utils.DFLPNG import DFLPNG
from utils.DFLJPG import DFLJPG
from .Sample import Sample
from .Sample import SampleType
from facelib import FaceType
from facelib import LandmarksProcessor
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:
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:
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:
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:
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:
if s_filename_path.suffix == '.png':
dflimg = DFLPNG.load ( str(s_filename_path) )
elif s_filename_path.suffix == '.jpg':
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) )
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(),
pitch=pitch,
yaw=yaw) )
except:
print ("Unable to load %s , error: %s" % (str(s_filename_path), traceback.format_exc() ) )
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):
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 ):
lowest_yaw, highest_yaw = -1.0, 1.0
gradations = 64
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:
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
for t_idx in t_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:
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,
landmarks=LandmarksProcessor.mirror_landmarks (sample.landmarks, sample.shape[1] ))
for sample in s[search_idx]
] if mirrored else s[search_idx]
break
return new_s

218
samplelib/SampleProcessor.py Normal file
View file

@ -0,0 +1,218 @@
from enum import IntEnum
import numpy as np
import cv2
import imagelib
from facelib import LandmarksProcessor
from facelib import FaceType
class SampleProcessor(object):
class TypeFlags(IntEnum):
SOURCE = 0x00000001,
WARPED = 0x00000002,
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_MASK_FULL = 0x00001000,
FACE_MASK_EYES = 0x00002000,
MODE_BGR = 0x01000000, #BGR
MODE_G = 0x02000000, #Grayscale
MODE_GGG = 0x04000000, #3xGrayscale
MODE_M = 0x08000000, #mask only
MODE_BGR_SHUFFLE = 0x10000000, #BGR shuffle
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.normalize_tanh = normalize_tanh
self.rotation_range = rotation_range
self.scale_range = scale_range
self.tx_range = tx_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
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))
params = imagelib.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 = []
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:
img_type = 1
elif f & SampleProcessor.TypeFlags.WARPED_TRANSFORMED != 0:
img_type = 2
elif f & SampleProcessor.TypeFlags.TRANSFORMED != 0:
img_type = 3
elif f & SampleProcessor.TypeFlags.LANDMARKS_ARRAY != 0:
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
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
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 = 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:
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)))
#remove landmarks near boundaries
for i in idxs[:]:
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:
idxs.remove(i)
#remove landmarks that close to each other in 5 dist
for landmarks in [s_landmarks, d_landmarks]:
for i in idxs[:]:
s_l = landmarks[i]
for j in idxs[:]:
if i == j:
continue
s_l_2 = landmarks[j]
diff_l = np.abs(s_l - s_l_2)
if np.sqrt(diff_l.dot(diff_l)) < 5:
idxs.remove(i)
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] ] ] )
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 = imagelib.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 )
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 )
images[img_type][face_mask_type] = imagelib.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
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)
img = img[start_y:start_y+sub_size,start_x:start_x+sub_size,:]
img_bgr = img[...,0:3]
img_mask = img[...,3:4]
if f & SampleProcessor.TypeFlags.MODE_BGR != 0:
img = img
elif f & SampleProcessor.TypeFlags.MODE_BGR_SHUFFLE != 0:
img_bgr = np.take (img_bgr, np.random.permutation(img_bgr.shape[-1]), axis=-1)
img = np.concatenate ( (img_bgr,img_mask) , -1 )
elif f & SampleProcessor.TypeFlags.MODE_G != 0:
img = np.concatenate ( (np.expand_dims(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY),-1),img_mask) , -1 )
elif f & SampleProcessor.TypeFlags.MODE_GGG != 0:
img = np.concatenate ( ( np.repeat ( np.expand_dims(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY),-1), (3,), -1), img_mask), -1)
elif is_face_sample and f & SampleProcessor.TypeFlags.MODE_M != 0:
if face_mask_type== 0:
raise ValueError ('no face_mask_type defined')
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)
else:
img = np.clip (img, 0.0, 1.0)
outputs.append ( img )
if debug:
result = []
for output in outputs:
if output.shape[2] < 4:
result += [output,]
elif output.shape[2] == 4:
result += [output[...,0:3]*output[...,3:4],]
return result
else:
return outputs

7
samplelib/__init__.py Normal file
View file

@ -0,0 +1,7 @@
from .Sample import Sample
from .Sample import SampleType
from .SampleLoader import SampleLoader
from .SampleProcessor import SampleProcessor
from .SampleGeneratorBase import SampleGeneratorBase
from .SampleGeneratorFace import SampleGeneratorFace
from .SampleGeneratorImageTemporal import SampleGeneratorImageTemporal