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

refactoring

Browse source
This commit is contained in:
iperov 2018-12-24 13:45:40 +04:00
commit 8a223845fb
19 changed files with 963 additions and 468 deletions
Download patch file Download diff file Expand all files Collapse all files

47
samples/Sample.py Normal file
View file

@ -0,0 +1,47 @@
from enum import IntEnum
import cv2
import numpy as np
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_END = 3
QTY = 4
class Sample(object):
def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, yaw=None, mirror=None, nearest_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.yaw = yaw
self.mirror = mirror
self.nearest_target_list = nearest_target_list
def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, yaw=None, mirror=None, nearest_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(),
yaw=yaw if yaw is not None else self.yaw,
mirror=mirror if mirror is not None else self.mirror,
nearest_target_list=nearest_target_list if nearest_target_list is not None else self.nearest_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_nearest_target_sample(self):
if self.nearest_target_list is None:
return None
return self.nearest_target_list[randint (0, len(self.nearest_target_list)-1)]

25
samples/SampleGeneratorBase.py Normal file
View file

@ -0,0 +1,25 @@
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

114
samples/SampleGeneratorFace.py Normal file
View file

@ -0,0 +1,114 @@
import traceback
import numpy as np
import random
import cv2
from utils import iter_utils
from samples import SampleType
from samples import SampleProcessor
from samples import SampleLoader
from samples import SampleGeneratorBase
'''
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, sample_process_options=SampleProcessor.Options(), output_sample_types=[], **kwargs):
super().__init__(samples_path, debug, batch_size)
self.sample_process_options = sample_process_options
self.output_sample_types = output_sample_types
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
else:
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.generator_samples = [ self.samples ]
self.generators = [iter_utils.ThisThreadGenerator ( self.batch_func, 0 )]
else:
if len(self.samples) > 1:
self.generator_samples = [ self.samples[0::2],
self.samples[1::2] ]
self.generators = [iter_utils.SubprocessGenerator ( self.batch_func, 0 ),
iter_utils.SubprocessGenerator ( self.batch_func, 1 )]
else:
self.generator_samples = [ self.samples ]
self.generators = [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]
data_len = len(samples)
if data_len == 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 ( [ x == None for x in samples] ):
raise ValueError('Not enough training data. Gather more faces!')
if self.sample_type == SampleType.FACE:
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 = [[]]*data_len
while True:
batches = None
for n_batch in range(self.batch_size):
while True:
sample = None
if self.sample_type == SampleType.FACE:
if len(shuffle_idxs) == 0:
shuffle_idxs = random.sample( range(data_len), data_len )
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 = random.sample( range(data_len), data_len )
idx = shuffle_idxs.pop()
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()
sample = samples[idx][idx2]
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)) ]
for i in range(len(x)):
batches[i].append ( x[i] )
break
yield [ np.array(batch) for batch in batches]

128
samples/SampleLoader.py Normal file
View file

@ -0,0 +1,128 @@
from enum import IntEnum
import cv2
import numpy as np
from tqdm import tqdm
from pathlib import Path
from utils import Path_utils
from utils.DFLPNG import DFLPNG
from .Sample import Sample
from .Sample import SampleType
from facelib import FaceType
from facelib import LandmarksProcessor
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
if target_samples_path is not None and str(target_samples_path) not in cache.keys():
cache[str(target_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 tqdm( Path_utils.get_image_paths(samples_path), desc="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) )
return datas[sample_type]
@staticmethod
def upgradeToFaceSamples ( samples ):
sample_list = []
for s in tqdm( samples, desc="Loading" ):
s_filename_path = Path(s.filename)
if s_filename_path.suffix != '.png':
print ("%s is not a png file required for training" % (s_filename_path.name) )
continue
dflpng = DFLPNG.load ( str(s_filename_path), print_on_no_embedded_data=True )
if dflpng is None:
continue
sample_list.append( s.copy_and_set(sample_type=SampleType.FACE,
face_type=FaceType.fromString (dflpng.get_face_type()),
shape=dflpng.get_shape(),
landmarks=dflpng.get_landmarks(),
yaw=dflpng.get_yaw_value()) )
return sample_list
@staticmethod
def upgradeToFaceYawSortedSamples( YAW_RAWS ):
lowest_yaw, highest_yaw = -32, +32
gradations = 64
diff_rot_per_grad = abs(highest_yaw-lowest_yaw) / gradations
yaws_sample_list = [None]*gradations
for i in tqdm( range(0, gradations), desc="Sorting" ):
yaw = lowest_yaw + i*diff_rot_per_grad
next_yaw = lowest_yaw + (i+1)*diff_rot_per_grad
yaw_samples = []
for s in YAW_RAWS:
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

152
samples/SampleProcessor.py Normal file
View file

@ -0,0 +1,152 @@
from enum import IntEnum
import numpy as np
import cv2
from utils import image_utils
from facelib import LandmarksProcessor
from facelib import FaceType
class SampleProcessor(object):
class TypeFlags(IntEnum):
SOURCE = 0x00000001,
WARPED = 0x00000002,
WARPED_TRANSFORMED = 0x00000004,
TRANSFORMED = 0x00000008,
FACE_ALIGN_HALF = 0x00000010,
FACE_ALIGN_FULL = 0x00000020,
FACE_ALIGN_HEAD = 0x00000040,
FACE_ALIGN_AVATAR = 0x00000080,
FACE_MASK_FULL = 0x00000100,
FACE_MASK_EYES = 0x00000200,
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):
source = sample.load_bgr()
h,w,c = source.shape
is_face_sample = sample.landmarks is not None
if debug and is_face_sample:
LandmarksProcessor.draw_landmarks (source, sample.landmarks, (0, 1, 0))
params = image_utils.gen_warp_params(source, 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(4)]
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
else:
raise ValueError ('expected SampleTypeFlags type')
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 images[img_type][face_mask_type] is None:
img = source
if is_face_sample:
if face_mask_type == 1:
img = np.concatenate( (img, LandmarksProcessor.get_image_hull_mask (source, sample.landmarks) ), -1 )
elif face_mask_type == 2:
mask = LandmarksProcessor.get_image_eye_mask (source, 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] = image_utils.warp_by_params (params, img, (img_type==1 or img_type==2), (img_type==2 or img_type==3), img_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_LANCZOS4 )
else:
img = cv2.resize( img, (size,size), cv2.INTER_LANCZOS4 )
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 and sample_process_options.normalize_tanh:
img = img * 2.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

6
samples/__init__.py Normal file
View file

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