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

Formatting

Browse source
This commit is contained in:
Jeremy Hummel 2019-08-10 00:49:25 -07:00
commit 9a41965063
1 changed files with 81 additions and 70 deletions
Download patch file Download diff file Expand all files Collapse all files

151
samplelib/SampleProcessor.py
View file

@ -7,7 +7,6 @@ import numpy as np
import imagelib import imagelib
from facelib import FaceType, LandmarksProcessor from facelib import FaceType, LandmarksProcessor
""" """
output_sample_types = [ output_sample_types = [
{} opts, {} opts,
@ -42,38 +41,40 @@ opts:
""" """
class SampleProcessor(object): class SampleProcessor(object):
class Types(IntEnum): class Types(IntEnum):
NONE = 0 NONE = 0
IMG_TYPE_BEGIN = 1 IMG_TYPE_BEGIN = 1
IMG_SOURCE = 1 IMG_SOURCE = 1
IMG_WARPED = 2 IMG_WARPED = 2
IMG_WARPED_TRANSFORMED = 3 IMG_WARPED_TRANSFORMED = 3
IMG_TRANSFORMED = 4 IMG_TRANSFORMED = 4
IMG_LANDMARKS_ARRAY = 5 #currently unused IMG_LANDMARKS_ARRAY = 5 # currently unused
IMG_PITCH_YAW_ROLL = 6 IMG_PITCH_YAW_ROLL = 6
IMG_PITCH_YAW_ROLL_SIGMOID = 7 IMG_PITCH_YAW_ROLL_SIGMOID = 7
IMG_TYPE_END = 10 IMG_TYPE_END = 10
FACE_TYPE_BEGIN = 10 FACE_TYPE_BEGIN = 10
FACE_TYPE_HALF = 10 FACE_TYPE_HALF = 10
FACE_TYPE_FULL = 11 FACE_TYPE_FULL = 11
FACE_TYPE_HEAD = 12 #currently unused FACE_TYPE_HEAD = 12 # currently unused
FACE_TYPE_AVATAR = 13 #currently unused FACE_TYPE_AVATAR = 13 # currently unused
FACE_TYPE_END = 20 FACE_TYPE_END = 20
MODE_BEGIN = 40 MODE_BEGIN = 40
MODE_BGR = 40 #BGR MODE_BGR = 40 # BGR
MODE_G = 41 #Grayscale MODE_G = 41 # Grayscale
MODE_GGG = 42 #3xGrayscale MODE_GGG = 42 # 3xGrayscale
MODE_M = 43 #mask only MODE_M = 43 # mask only
MODE_BGR_SHUFFLE = 44 #BGR shuffle MODE_BGR_SHUFFLE = 44 # BGR shuffle
MODE_END = 50 MODE_END = 50
class Options(object): class Options(object):
def __init__(self, random_flip = True, rotation_range=[-10,10], scale_range=[-0.05, 0.05], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05] ): def __init__(self, random_flip=True, 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.rotation_range = rotation_range self.rotation_range = rotation_range
self.scale_range = scale_range self.scale_range = scale_range
@ -81,35 +82,39 @@ class SampleProcessor(object):
self.ty_range = ty_range self.ty_range = ty_range
@staticmethod @staticmethod
def process (sample, sample_process_options, output_sample_types, debug, ct_sample=None): def process(sample, sample_process_options, output_sample_types, debug, ct_sample=None):
SPTF = SampleProcessor.Types SPTF = SampleProcessor.Types
sample_bgr = sample.load_bgr() sample_bgr = sample.load_bgr()
ct_sample_bgr = None ct_sample_bgr = None
ct_sample_mask = None ct_sample_mask = None
h,w,c = sample_bgr.shape 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: if debug and is_face_sample:
LandmarksProcessor.draw_landmarks (sample_bgr, sample.landmarks, (0, 1, 0)) LandmarksProcessor.draw_landmarks(sample_bgr, 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 ) 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)
cached_images = collections.defaultdict(dict) cached_images = collections.defaultdict(dict)
sample_rnd_seed = np.random.randint(0x80000000) sample_rnd_seed = np.random.randint(0x80000000)
SPTF_FACETYPE_TO_FACETYPE = { SPTF.FACE_TYPE_HALF : FaceType.HALF, SPTF_FACETYPE_TO_FACETYPE = {SPTF.FACE_TYPE_HALF: FaceType.HALF,
SPTF.FACE_TYPE_FULL : FaceType.FULL, SPTF.FACE_TYPE_FULL: FaceType.FULL,
SPTF.FACE_TYPE_HEAD : FaceType.HEAD, SPTF.FACE_TYPE_HEAD: FaceType.HEAD,
SPTF.FACE_TYPE_AVATAR : FaceType.AVATAR } SPTF.FACE_TYPE_AVATAR: FaceType.AVATAR}
outputs = [] outputs = []
for opts in output_sample_types: for opts in output_sample_types:
resolution = opts.get('resolution', 0) resolution = opts.get('resolution', 0)
types = opts.get('types', [] ) types = opts.get('types', [])
random_sub_res = opts.get('random_sub_res', 0) random_sub_res = opts.get('random_sub_res', 0)
normalize_std_dev = opts.get('normalize_std_dev', False) normalize_std_dev = opts.get('normalize_std_dev', False)
@ -131,11 +136,11 @@ class SampleProcessor(object):
mode_type = t mode_type = t
if img_type == SPTF.NONE: if img_type == SPTF.NONE:
raise ValueError ('expected IMG_ type') raise ValueError('expected IMG_ type')
if img_type == SPTF.IMG_LANDMARKS_ARRAY: if img_type == SPTF.IMG_LANDMARKS_ARRAY:
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.concatenate([np.expand_dims(l[:, 0] / w, -1), np.expand_dims(l[:, 1] / h, -1)], -1)
l = np.clip(l, 0.0, 1.0) l = np.clip(l, 0.0, 1.0)
img = l img = l
elif img_type == SPTF.IMG_PITCH_YAW_ROLL or img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID: elif img_type == SPTF.IMG_PITCH_YAW_ROLL or img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID:
@ -143,19 +148,19 @@ class SampleProcessor(object):
if pitch_yaw_roll is not None: if pitch_yaw_roll is not None:
pitch, yaw, roll = pitch_yaw_roll pitch, yaw, roll = pitch_yaw_roll
else: else:
pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll (sample.landmarks) pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll(sample.landmarks)
if params['flip']: if params['flip']:
yaw = -yaw yaw = -yaw
if img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID: if img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID:
pitch = (pitch+1.0) / 2.0 pitch = (pitch + 1.0) / 2.0
yaw = (yaw+1.0) / 2.0 yaw = (yaw + 1.0) / 2.0
roll = (roll+1.0) / 2.0 roll = (roll + 1.0) / 2.0
img = (pitch, yaw, roll) img = (pitch, yaw, roll)
else: else:
if mode_type == SPTF.NONE: if mode_type == SPTF.NONE:
raise ValueError ('expected MODE_ type') raise ValueError('expected MODE_ type')
img = cached_images.get(img_type, None) img = cached_images.get(img_type, None)
if img is None: if img is None:
@ -170,98 +175,104 @@ class SampleProcessor(object):
chance = np.clip(chance, 0, 100) chance = np.clip(chance, 0, 100)
if np.random.randint(100) < chance: if np.random.randint(100) < chance:
mb_range = [3,5,7,9][ : np.clip(mb_range, 0, 3)+1 ] mb_range = [3, 5, 7, 9][: np.clip(mb_range, 0, 3) + 1]
dim = mb_range[ np.random.randint(len(mb_range) ) ] dim = mb_range[np.random.randint(len(mb_range))]
img = imagelib.LinearMotionBlur (img, dim, np.random.randint(180) ) img = imagelib.LinearMotionBlur(img, dim, np.random.randint(180))
mask = cur_sample.load_fanseg_mask() #using fanseg_mask if exist mask = cur_sample.load_fanseg_mask() # using fanseg_mask if exist
if mask is None: if mask is None:
mask = LandmarksProcessor.get_image_hull_mask (img.shape, cur_sample.landmarks) mask = LandmarksProcessor.get_image_hull_mask(img.shape, cur_sample.landmarks)
if cur_sample.ie_polys is not None: if cur_sample.ie_polys is not None:
cur_sample.ie_polys.overlay_mask(mask) cur_sample.ie_polys.overlay_mask(mask)
warp = (img_type==SPTF.IMG_WARPED or img_type==SPTF.IMG_WARPED_TRANSFORMED) warp = (img_type == SPTF.IMG_WARPED or img_type == SPTF.IMG_WARPED_TRANSFORMED)
transform = (img_type==SPTF.IMG_WARPED_TRANSFORMED or img_type==SPTF.IMG_TRANSFORMED) transform = (img_type == SPTF.IMG_WARPED_TRANSFORMED or img_type == SPTF.IMG_TRANSFORMED)
flip = img_type != SPTF.IMG_WARPED flip = img_type != SPTF.IMG_WARPED
img = imagelib.warp_by_params (params, img, warp, transform, flip, True) img = imagelib.warp_by_params(params, img, warp, transform, flip, True)
if mask is not None: if mask is not None:
mask = imagelib.warp_by_params (params, mask, warp, transform, flip, False)[...,np.newaxis] mask = imagelib.warp_by_params(params, mask, warp, transform, flip, False)[..., np.newaxis]
img = np.concatenate( (img, mask ), -1 ) img = np.concatenate((img, mask), -1)
cached_images[img_type] = img cached_images[img_type] = img
if is_face_sample and target_face_type != SPTF.NONE: if is_face_sample and target_face_type != SPTF.NONE:
ft = SPTF_FACETYPE_TO_FACETYPE[target_face_type] ft = SPTF_FACETYPE_TO_FACETYPE[target_face_type]
if ft > sample.face_type: if ft > 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, ft) ) raise Exception(
img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, resolution, ft), (resolution,resolution), flags=cv2.INTER_CUBIC ) 'sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.' % (
sample.filename, sample.face_type, ft))
img = cv2.warpAffine(img, LandmarksProcessor.get_transform_mat(sample.landmarks, resolution, ft),
(resolution, resolution), flags=cv2.INTER_CUBIC)
else: else:
img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC ) img = cv2.resize(img, (resolution, resolution), cv2.INTER_CUBIC)
if random_sub_res != 0: if random_sub_res != 0:
sub_size = resolution - random_sub_res sub_size = resolution - random_sub_res
rnd_state = np.random.RandomState (sample_rnd_seed+random_sub_res) rnd_state = np.random.RandomState(sample_rnd_seed + random_sub_res)
start_x = rnd_state.randint(sub_size+1) start_x = rnd_state.randint(sub_size + 1)
start_y = 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 = img[start_y:start_y + sub_size, start_x:start_x + sub_size, :]
img = np.clip(img, 0, 1) img = np.clip(img, 0, 1)
img_bgr = img[...,0:3] img_bgr = img[..., 0:3]
img_mask = img[...,3:4] img_mask = img[..., 3:4]
if apply_ct and ct_sample is not None: if apply_ct and ct_sample is not None:
if ct_sample_bgr is None: if ct_sample_bgr is None:
ct_sample_bgr = ct_sample.load_bgr() ct_sample_bgr = ct_sample.load_bgr()
ct_sample_bgr_resized = cv2.resize( ct_sample_bgr, (resolution,resolution), cv2.INTER_LINEAR ) ct_sample_bgr_resized = cv2.resize(ct_sample_bgr, (resolution, resolution), cv2.INTER_LINEAR)
img_bgr = imagelib.linear_color_transfer (img_bgr, ct_sample_bgr_resized) img_bgr = imagelib.linear_color_transfer(img_bgr, ct_sample_bgr_resized)
img_bgr = np.clip( img_bgr, 0.0, 1.0) img_bgr = np.clip(img_bgr, 0.0, 1.0)
if normalize_std_dev: if normalize_std_dev:
img_bgr = (img_bgr - img_bgr.mean( (0,1)) ) / img_bgr.std( (0,1) ) img_bgr = (img_bgr - img_bgr.mean((0, 1))) / img_bgr.std((0, 1))
elif normalize_vgg: elif normalize_vgg:
img_bgr = np.clip(img_bgr*255, 0, 255) img_bgr = np.clip(img_bgr * 255, 0, 255)
img_bgr[:,:,0] -= 103.939 img_bgr[:, :, 0] -= 103.939
img_bgr[:,:,1] -= 116.779 img_bgr[:, :, 1] -= 116.779
img_bgr[:,:,2] -= 123.68 img_bgr[:, :, 2] -= 123.68
if mode_type == SPTF.MODE_BGR: if mode_type == SPTF.MODE_BGR:
img = img_bgr img = img_bgr
elif mode_type == SPTF.MODE_BGR_SHUFFLE: elif mode_type == SPTF.MODE_BGR_SHUFFLE:
rnd_state = np.random.RandomState (sample_rnd_seed) rnd_state = np.random.RandomState(sample_rnd_seed)
img = np.take (img_bgr, rnd_state.permutation(img_bgr.shape[-1]), axis=-1) img = np.take(img_bgr, rnd_state.permutation(img_bgr.shape[-1]), axis=-1)
elif mode_type == SPTF.MODE_G: elif mode_type == SPTF.MODE_G:
img = np.concatenate ( (np.expand_dims(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY),-1),img_mask) , -1 ) img = np.concatenate((np.expand_dims(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY), -1), img_mask), -1)
elif mode_type == SPTF.MODE_GGG: elif mode_type == SPTF.MODE_GGG:
img = np.concatenate ( ( np.repeat ( np.expand_dims(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY),-1), (3,), -1), img_mask), -1) img = np.concatenate(
(np.repeat(np.expand_dims(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY), -1), (3,), -1), img_mask),
-1)
elif mode_type == SPTF.MODE_M and is_face_sample: elif mode_type == SPTF.MODE_M and is_face_sample:
img = img_mask img = img_mask
if not debug: if not debug:
if normalize_tanh: if normalize_tanh:
img = np.clip (img * 2.0 - 1.0, -1.0, 1.0) img = np.clip(img * 2.0 - 1.0, -1.0, 1.0)
else: else:
img = np.clip (img, 0.0, 1.0) img = np.clip(img, 0.0, 1.0)
outputs.append ( img ) outputs.append(img)
if debug: if debug:
result = [] result = []
for output in outputs: for output in outputs:
if output.shape[2] < 4: if output.shape[2] < 4:
result += [output,] result += [output, ]
elif output.shape[2] == 4: elif output.shape[2] == 4:
result += [output[...,0:3]*output[...,3:4],] result += [output[..., 0:3] * output[..., 3:4], ]
return result return result
else: else:
return outputs return outputs
""" """
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 = 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 close_sample_bgr = close_sample.load_bgr() if close_sample is not None else None