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

SampleProcessor.py : refactoring and gen mask struct

Browse source
This commit is contained in:
Colombo 2020-01-29 18:08:54 +04:00
parent 0251eb3490
commit 5fe5fa131c
4 changed files with 51 additions and 80 deletions
Download patch file Download diff file Expand all files Collapse all files

14
facelib/LandmarksProcessor.py
View file

@ -609,7 +609,13 @@ def mirror_landmarks (landmarks, val):
result[:,0] = val - result[:,0] - 1 result[:,0] = val - result[:,0] - 1
return result return result
def draw_landmarks (image, image_landmarks, color=(0,255,0), transparent_mask=False, ie_polys=None): def get_face_struct_mask (image_shape, image_landmarks, eyebrows_expand_mod=1.0, ie_polys=None, color=(1,) ):
mask = np.zeros(image_shape[0:2]+( len(color),),dtype=np.float32)
lmrks = expand_eyebrows(image_landmarks, eyebrows_expand_mod)
draw_landmarks (mask, image_landmarks, color=color, draw_circles=False, thickness=2, ie_polys=ie_polys)
return mask
def draw_landmarks (image, image_landmarks, color=(0,255,0), draw_circles=True, thickness=1, transparent_mask=False, ie_polys=None):
if len(image_landmarks) != 68: if len(image_landmarks) != 68:
raise Exception('get_image_eye_mask works only with 68 landmarks') raise Exception('get_image_eye_mask works only with 68 landmarks')
@ -625,10 +631,12 @@ def draw_landmarks (image, image_landmarks, color=(0,255,0), transparent_mask=Fa
# open shapes # open shapes
cv2.polylines(image, tuple(np.array([v]) for v in ( right_eyebrow, jaw, left_eyebrow, np.concatenate((nose, [nose[-6]])) )), cv2.polylines(image, tuple(np.array([v]) for v in ( right_eyebrow, jaw, left_eyebrow, np.concatenate((nose, [nose[-6]])) )),
False, color, lineType=cv2.LINE_AA) False, color, thickness=thickness, lineType=cv2.LINE_AA)
# closed shapes # closed shapes
cv2.polylines(image, tuple(np.array([v]) for v in (right_eye, left_eye, mouth)), cv2.polylines(image, tuple(np.array([v]) for v in (right_eye, left_eye, mouth)),
True, color, lineType=cv2.LINE_AA) True, color, thickness=thickness, lineType=cv2.LINE_AA)
if draw_circles:
# the rest of the cicles # the rest of the cicles
for x, y in np.concatenate((right_eyebrow, left_eyebrow, mouth, right_eye, left_eye, nose), axis=0): for x, y in np.concatenate((right_eyebrow, left_eyebrow, mouth, right_eye, left_eye, nose), axis=0):
cv2.circle(image, (x, y), 1, color, 1, lineType=cv2.LINE_AA) cv2.circle(image, (x, y), 1, color, 1, lineType=cv2.LINE_AA)

4
models/Model_Quick96/Model.py
View file

@ -372,14 +372,14 @@ class QModel(ModelBase):
sample_process_options=SampleProcessor.Options(random_flip=True if self.pretrain else False), sample_process_options=SampleProcessor.Options(random_flip=True if self.pretrain else False),
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution':resolution, }, output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution':resolution, },
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution, }, {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution, },
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'data_format':nn.data_format, 'resolution': resolution } ], {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_FACE_MASK_HULL), 'data_format':nn.data_format, 'resolution': resolution } ],
generators_count=src_generators_count ), generators_count=src_generators_count ),
SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(), SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=True if self.pretrain else False), sample_process_options=SampleProcessor.Options(random_flip=True if self.pretrain else False),
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution':resolution}, output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution':resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution}, {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'data_format':nn.data_format, 'resolution': resolution} ], {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_FACE_MASK_HULL), 'data_format':nn.data_format, 'resolution': resolution} ],
generators_count=dst_generators_count ) generators_count=dst_generators_count )
]) ])

4
models/Model_SAEHD/Model.py
View file

@ -725,14 +725,14 @@ class SAEHDModel(ModelBase):
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip), sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
output_sample_types = [ {'types' : (t_img_warped, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution, 'ct_mode': self.options['ct_mode'] }, output_sample_types = [ {'types' : (t_img_warped, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution, 'ct_mode': self.options['ct_mode'] },
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution, 'ct_mode': self.options['ct_mode'] }, {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution, 'ct_mode': self.options['ct_mode'] },
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'data_format':nn.data_format, 'resolution': resolution } ], {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_FACE_MASK_HULL), 'data_format':nn.data_format, 'resolution': resolution } ],
generators_count=src_generators_count ), generators_count=src_generators_count ),
SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(), SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.get_batch_size(),
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip), sample_process_options=SampleProcessor.Options(random_flip=self.random_flip),
output_sample_types = [ {'types' : (t_img_warped, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution}, output_sample_types = [ {'types' : (t_img_warped, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution}, {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_BGR), 'data_format':nn.data_format, 'resolution': resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'data_format':nn.data_format, 'resolution': resolution} ], {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_FACE_MASK_HULL), 'data_format':nn.data_format, 'resolution': resolution} ],
generators_count=dst_generators_count ) generators_count=dst_generators_count )
]) ])

97
samplelib/SampleProcessor.py
View file

@ -7,41 +7,6 @@ import numpy as np
from core import imagelib from core import imagelib
from facelib import FaceType, LandmarksProcessor from facelib import FaceType, LandmarksProcessor
"""
output_sample_types = [
{} opts,
...
]
opts:
'types' : (S,S,...,S)
where S:
'IMG_SOURCE'
'IMG_WARPED'
'IMG_WARPED_TRANSFORMED''
'IMG_TRANSFORMED'
'IMG_LANDMARKS_ARRAY' #currently unused
'IMG_PITCH_YAW_ROLL'
'FACE_TYPE_HALF'
'FACE_TYPE_FULL'
'FACE_TYPE_HEAD' #currently unused
'FACE_TYPE_AVATAR' #currently unused
'MODE_BGR' #BGR
'MODE_G' #Grayscale
'MODE_GGG' #3xGrayscale
'MODE_M' #mask only
'MODE_BGR_SHUFFLE' #BGR shuffle
'resolution' : N
'motion_blur' : (chance_int, range) - chance 0..100 to apply to face (not mask), and max_size of motion blur
'ct_mode' :
'normalize_tanh' : bool
"""
class SampleProcessor(object): class SampleProcessor(object):
class Types(IntEnum): class Types(IntEnum):
NONE = 0 NONE = 0
@ -70,9 +35,10 @@ class SampleProcessor(object):
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_FACE_MASK_HULL = 43 #mask hull as grayscale
MODE_BGR_SHUFFLE = 44 #BGR shuffle MODE_FACE_MASK_STRUCT = 44 #mask structure as grayscale
MODE_BGR_RANDOM_HSV_SHIFT = 45 MODE_BGR_SHUFFLE = 45 #BGR shuffle
MODE_BGR_RANDOM_HSV_SHIFT = 46
MODE_END = 50 MODE_END = 50
class Options(object): class Options(object):
@ -135,8 +101,10 @@ class SampleProcessor(object):
elif t >= SPTF.MODE_BEGIN and t < SPTF.MODE_END: elif t >= SPTF.MODE_BEGIN and t < SPTF.MODE_END:
mode_type = t mode_type = t
if mode_type == SPTF.MODE_M and not is_face_sample: if mode_type == SPTF.MODE_FACE_MASK_HULL and not is_face_sample:
raise ValueError("MODE_M applicable only for face samples") raise ValueError("MODE_FACE_MASK_HULL applicable only for face samples")
if mode_type == SPTF.MODE_FACE_MASK_STRUCT and not is_face_sample:
raise ValueError("MODE_FACE_MASK_STRUCT applicable only for face samples")
can_warp = (img_type==SPTF.IMG_WARPED or img_type==SPTF.IMG_WARPED_TRANSFORMED) can_warp = (img_type==SPTF.IMG_WARPED or img_type==SPTF.IMG_WARPED_TRANSFORMED)
can_transform = (img_type==SPTF.IMG_WARPED_TRANSFORMED or img_type==SPTF.IMG_TRANSFORMED) can_transform = (img_type==SPTF.IMG_WARPED_TRANSFORMED or img_type==SPTF.IMG_TRANSFORMED)
@ -165,10 +133,7 @@ class SampleProcessor(object):
if mode_type == SPTF.NONE: if mode_type == SPTF.NONE:
raise ValueError ('expected MODE_ type') raise ValueError ('expected MODE_ type')
need_img = mode_type != SPTF.MODE_M if mode_type == SPTF.MODE_FACE_MASK_HULL:
need_mask = mode_type == SPTF.MODE_M
if need_mask:
if sample.eyebrows_expand_mod is not None: if sample.eyebrows_expand_mod is not None:
mask = LandmarksProcessor.get_image_hull_mask (sample_bgr.shape, sample.landmarks, eyebrows_expand_mod=sample.eyebrows_expand_mod ) mask = LandmarksProcessor.get_image_hull_mask (sample_bgr.shape, sample.landmarks, eyebrows_expand_mod=sample.eyebrows_expand_mod )
else: else:
@ -176,8 +141,12 @@ class SampleProcessor(object):
if sample.ie_polys is not None: if sample.ie_polys is not None:
sample.ie_polys.overlay_mask(mask) sample.ie_polys.overlay_mask(mask)
elif mode_type == SPTF.MODE_FACE_MASK_STRUCT:
if need_img: if sample.eyebrows_expand_mod is not None:
mask = LandmarksProcessor.get_face_struct_mask (sample_bgr.shape, sample.landmarks, eyebrows_expand_mod=sample.eyebrows_expand_mod )
else:
mask = LandmarksProcessor.get_face_struct_mask (sample_bgr.shape, sample.landmarks)
else:
img = sample_bgr img = sample_bgr
if motion_blur is not None: if motion_blur is not None:
chance, mb_max_size = motion_blur chance, mb_max_size = motion_blur
@ -201,37 +170,31 @@ class SampleProcessor(object):
if sample.face_type == FaceType.MARK_ONLY: if sample.face_type == FaceType.MARK_ONLY:
mat = LandmarksProcessor.get_transform_mat (sample.landmarks, sample.shape[0], target_ft), (sample.shape[0],sample.shape[0]) mat = LandmarksProcessor.get_transform_mat (sample.landmarks, sample.shape[0], target_ft), (sample.shape[0],sample.shape[0])
if need_img: if mode_type == SPTF.MODE_FACE_MASK_HULL or mode_type == SPTF.MODE_FACE_MASK_STRUCT:
img = cv2.warpAffine( img, mat, flags=cv2.INTER_CUBIC )
img = imagelib.warp_by_params (params, img, can_warp, can_transform, can_flip=True, border_replicate=True)
img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC )
if need_mask:
mask = cv2.warpAffine( mask, mat, flags=cv2.INTER_CUBIC ) mask = cv2.warpAffine( mask, mat, flags=cv2.INTER_CUBIC )
mask = imagelib.warp_by_params (params, mask, can_warp, can_transform, can_flip=True, border_replicate=False) mask = imagelib.warp_by_params (params, mask, can_warp, can_transform, can_flip=True, border_replicate=False)
mask = cv2.resize( mask, (resolution,resolution), cv2.INTER_CUBIC )[...,None] mask = cv2.resize( mask, (resolution,resolution), cv2.INTER_CUBIC )[...,None]
else: else:
mat = LandmarksProcessor.get_transform_mat (sample.landmarks, resolution, target_ft) img = cv2.warpAffine( img, mat, flags=cv2.INTER_CUBIC )
if need_img:
img = imagelib.warp_by_params (params, img, can_warp, can_transform, can_flip=True, border_replicate=True)
img = cv2.warpAffine( img, mat, (resolution,resolution), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC )
if need_mask:
mask = imagelib.warp_by_params (params, mask, can_warp, can_transform, can_flip=True, border_replicate=False)
mask = cv2.warpAffine( mask, mat, (resolution,resolution), borderMode=cv2.BORDER_CONSTANT, flags=cv2.INTER_CUBIC )[...,None]
else:
if need_img:
img = imagelib.warp_by_params (params, img, can_warp, can_transform, can_flip=True, border_replicate=True) img = imagelib.warp_by_params (params, img, can_warp, can_transform, can_flip=True, border_replicate=True)
img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC ) img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC )
if need_mask: else:
mat = LandmarksProcessor.get_transform_mat (sample.landmarks, resolution, target_ft)
if mode_type == SPTF.MODE_FACE_MASK_HULL or mode_type == SPTF.MODE_FACE_MASK_STRUCT:
mask = imagelib.warp_by_params (params, mask, can_warp, can_transform, can_flip=True, border_replicate=False) mask = imagelib.warp_by_params (params, mask, can_warp, can_transform, can_flip=True, border_replicate=False)
mask = cv2.resize( mask, (resolution,resolution), cv2.INTER_CUBIC )[...,None] mask = cv2.warpAffine( mask, mat, (resolution,resolution), borderMode=cv2.BORDER_CONSTANT, flags=cv2.INTER_CUBIC )[...,None]
else:
img = imagelib.warp_by_params (params, img, can_warp, can_transform, can_flip=True, border_replicate=True)
img = cv2.warpAffine( img, mat, (resolution,resolution), borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC )
else:
img = imagelib.warp_by_params (params, img, can_warp, can_transform, can_flip=True, border_replicate=True)
img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC )
if mode_type == SPTF.MODE_M: if mode_type == SPTF.MODE_FACE_MASK_HULL or mode_type == SPTF.MODE_FACE_MASK_STRUCT:
out_sample = np.clip(mask, 0, 1).astype(np.float32) out_sample = np.clip(mask, 0, 1).astype(np.float32)
else: else:
img = np.clip(img, 0, 1).astype(np.float32) img = np.clip(img, 0, 1).astype(np.float32)