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

135
samplelib/SampleProcessor.py
View file

@ -7,7 +7,6 @@ import numpy as np
import imagelib
from facelib import FaceType, LandmarksProcessor
"""
output_sample_types = [
{} opts,
@ -42,6 +41,7 @@ opts:
"""
class SampleProcessor(object):
class Types(IntEnum):
NONE = 0
@ -51,7 +51,7 @@ class SampleProcessor(object):
IMG_WARPED = 2
IMG_WARPED_TRANSFORMED = 3
IMG_TRANSFORMED = 4
IMG_LANDMARKS_ARRAY = 5 #currently unused
IMG_LANDMARKS_ARRAY = 5 # currently unused
IMG_PITCH_YAW_ROLL = 6
IMG_PITCH_YAW_ROLL_SIGMOID = 7
IMG_TYPE_END = 10
@ -59,21 +59,22 @@ class SampleProcessor(object):
FACE_TYPE_BEGIN = 10
FACE_TYPE_HALF = 10
FACE_TYPE_FULL = 11
FACE_TYPE_HEAD = 12 #currently unused
FACE_TYPE_AVATAR = 13 #currently unused
FACE_TYPE_HEAD = 12 # currently unused
FACE_TYPE_AVATAR = 13 # currently unused
FACE_TYPE_END = 20
MODE_BEGIN = 40
MODE_BGR = 40 #BGR
MODE_G = 41 #Grayscale
MODE_GGG = 42 #3xGrayscale
MODE_M = 43 #mask only
MODE_BGR_SHUFFLE = 44 #BGR shuffle
MODE_BGR = 40 # BGR
MODE_G = 41 # Grayscale
MODE_GGG = 42 # 3xGrayscale
MODE_M = 43 # mask only
MODE_BGR_SHUFFLE = 44 # BGR shuffle
MODE_END = 50
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.rotation_range = rotation_range
self.scale_range = scale_range
@ -81,35 +82,39 @@ class SampleProcessor(object):
self.ty_range = ty_range
@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
sample_bgr = sample.load_bgr()
ct_sample_bgr = 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
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)
sample_rnd_seed = np.random.randint(0x80000000)
SPTF_FACETYPE_TO_FACETYPE = { SPTF.FACE_TYPE_HALF : FaceType.HALF,
SPTF.FACE_TYPE_FULL : FaceType.FULL,
SPTF.FACE_TYPE_HEAD : FaceType.HEAD,
SPTF.FACE_TYPE_AVATAR : FaceType.AVATAR }
SPTF_FACETYPE_TO_FACETYPE = {SPTF.FACE_TYPE_HALF: FaceType.HALF,
SPTF.FACE_TYPE_FULL: FaceType.FULL,
SPTF.FACE_TYPE_HEAD: FaceType.HEAD,
SPTF.FACE_TYPE_AVATAR: FaceType.AVATAR}
outputs = []
for opts in output_sample_types:
resolution = opts.get('resolution', 0)
types = opts.get('types', [] )
types = opts.get('types', [])
random_sub_res = opts.get('random_sub_res', 0)
normalize_std_dev = opts.get('normalize_std_dev', False)
@ -131,11 +136,11 @@ class SampleProcessor(object):
mode_type = t
if img_type == SPTF.NONE:
raise ValueError ('expected IMG_ type')
raise ValueError('expected IMG_ type')
if img_type == SPTF.IMG_LANDMARKS_ARRAY:
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)
img = l
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:
pitch, yaw, roll = pitch_yaw_roll
else:
pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll (sample.landmarks)
pitch, yaw, roll = LandmarksProcessor.estimate_pitch_yaw_roll(sample.landmarks)
if params['flip']:
yaw = -yaw
if img_type == SPTF.IMG_PITCH_YAW_ROLL_SIGMOID:
pitch = (pitch+1.0) / 2.0
yaw = (yaw+1.0) / 2.0
roll = (roll+1.0) / 2.0
pitch = (pitch + 1.0) / 2.0
yaw = (yaw + 1.0) / 2.0
roll = (roll + 1.0) / 2.0
img = (pitch, yaw, roll)
else:
if mode_type == SPTF.NONE:
raise ValueError ('expected MODE_ type')
raise ValueError('expected MODE_ type')
img = cached_images.get(img_type, None)
if img is None:
@ -170,98 +175,104 @@ class SampleProcessor(object):
chance = np.clip(chance, 0, 100)
if np.random.randint(100) < chance:
mb_range = [3,5,7,9][ : np.clip(mb_range, 0, 3)+1 ]
dim = mb_range[ np.random.randint(len(mb_range) ) ]
img = imagelib.LinearMotionBlur (img, dim, np.random.randint(180) )
mb_range = [3, 5, 7, 9][: np.clip(mb_range, 0, 3) + 1]
dim = mb_range[np.random.randint(len(mb_range))]
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:
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:
cur_sample.ie_polys.overlay_mask(mask)
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)
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)
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:
mask = imagelib.warp_by_params (params, mask, warp, transform, flip, False)[...,np.newaxis]
img = np.concatenate( (img, mask ), -1 )
mask = imagelib.warp_by_params(params, mask, warp, transform, flip, False)[..., np.newaxis]
img = np.concatenate((img, mask), -1)
cached_images[img_type] = img
if is_face_sample and target_face_type != SPTF.NONE:
ft = SPTF_FACETYPE_TO_FACETYPE[target_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) )
img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, resolution, ft), (resolution,resolution), flags=cv2.INTER_CUBIC )
raise Exception(
'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:
img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC )
img = cv2.resize(img, (resolution, resolution), cv2.INTER_CUBIC)
if random_sub_res != 0:
sub_size = resolution - random_sub_res
rnd_state = np.random.RandomState (sample_rnd_seed+random_sub_res)
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,:]
rnd_state = np.random.RandomState(sample_rnd_seed + random_sub_res)
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 = np.clip(img, 0, 1)
img_bgr = img[...,0:3]
img_mask = img[...,3:4]
img_bgr = img[..., 0:3]
img_mask = img[..., 3:4]
if apply_ct and ct_sample is not None:
if ct_sample_bgr is None:
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 = np.clip( img_bgr, 0.0, 1.0)
img_bgr = imagelib.linear_color_transfer(img_bgr, ct_sample_bgr_resized)
img_bgr = np.clip(img_bgr, 0.0, 1.0)
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:
img_bgr = np.clip(img_bgr*255, 0, 255)
img_bgr[:,:,0] -= 103.939
img_bgr[:,:,1] -= 116.779
img_bgr[:,:,2] -= 123.68
img_bgr = np.clip(img_bgr * 255, 0, 255)
img_bgr[:, :, 0] -= 103.939
img_bgr[:, :, 1] -= 116.779
img_bgr[:, :, 2] -= 123.68
if mode_type == SPTF.MODE_BGR:
img = img_bgr
elif mode_type == SPTF.MODE_BGR_SHUFFLE:
rnd_state = np.random.RandomState (sample_rnd_seed)
img = np.take (img_bgr, rnd_state.permutation(img_bgr.shape[-1]), axis=-1)
rnd_state = np.random.RandomState(sample_rnd_seed)
img = np.take(img_bgr, rnd_state.permutation(img_bgr.shape[-1]), axis=-1)
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:
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:
img = img_mask
if not debug:
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:
img = np.clip (img, 0.0, 1.0)
img = np.clip(img, 0.0, 1.0)
outputs.append ( img )
outputs.append(img)
if debug:
result = []
for output in outputs:
if output.shape[2] < 4:
result += [output,]
result += [output, ]
elif output.shape[2] == 4:
result += [output[...,0:3]*output[...,3:4],]
result += [output[..., 0:3] * output[..., 3:4], ]
return result
else:
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_bgr = close_sample.load_bgr() if close_sample is not None else None