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

Merge pull request #16 from faceshiftlabs/build/update-from-parent

Browse source 
Build/update from parent
This commit is contained in:
Jeremy Hummel 2019-08-15 20:18:04 +00:00 • committed by GitHub
commit a3319b5e9f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
17 changed files with 389 additions and 248 deletions
Download patch file Download diff file Expand all files Collapse all files

45
converters/ConverterAvatar.py
View file

@ -20,8 +20,9 @@ class ConverterAvatar(Converter):
self.predictor_input_size = predictor_input_size self.predictor_input_size = predictor_input_size
#dummy predict and sleep, tensorflow caching kernels. If remove it, conversion speed will be x2 slower #dummy predict and sleep, tensorflow caching kernels. If remove it, conversion speed will be x2 slower
predictor_func ( np.zeros ( (predictor_input_size,predictor_input_size,3), dtype=np.float32 ), predictor_func ( np.zeros ( (predictor_input_size,predictor_input_size,3), dtype=np.float32 ),
np.zeros ( (predictor_input_size,predictor_input_size,1), dtype=np.float32 ) ) np.zeros ( (predictor_input_size,predictor_input_size,3), dtype=np.float32 ),
np.zeros ( (predictor_input_size,predictor_input_size,3), dtype=np.float32 ) )
time.sleep(2) time.sleep(2)
predictor_func_host, predictor_func = SubprocessFunctionCaller.make_pair(predictor_func) predictor_func_host, predictor_func = SubprocessFunctionCaller.make_pair(predictor_func)
@ -33,38 +34,28 @@ class ConverterAvatar(Converter):
self.predictor_func_host.obj.process_messages() self.predictor_func_host.obj.process_messages()
#override #override
def cli_convert_face (self, img_bgr, img_face_landmarks, debug, avaperator_face_bgr=None, **kwargs): def cli_convert_face (self, f0, f0_lmrk, f1, f1_lmrk, f2, f2_lmrk, debug, **kwargs):
if debug: if debug:
debugs = [img_bgr.copy()] debugs = []
img_size = img_bgr.shape[1], img_bgr.shape[0] inp_size = self.predictor_input_size
img_face_mask_a = LandmarksProcessor.get_image_hull_mask (img_bgr.shape, img_face_landmarks) f0_mat = LandmarksProcessor.get_transform_mat (f0_lmrk, inp_size, face_type=FaceType.FULL_NO_ALIGN)
img_face_mask_aaa = np.repeat(img_face_mask_a, 3, -1) f1_mat = LandmarksProcessor.get_transform_mat (f1_lmrk, inp_size, face_type=FaceType.FULL_NO_ALIGN)
f2_mat = LandmarksProcessor.get_transform_mat (f2_lmrk, inp_size, face_type=FaceType.FULL_NO_ALIGN)
output_size = self.predictor_input_size inp_f0 = cv2.warpAffine( f0, f0_mat, (inp_size, inp_size), flags=cv2.INTER_CUBIC )
face_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, output_size, face_type=FaceType.FULL) inp_f1 = cv2.warpAffine( f1, f1_mat, (inp_size, inp_size), flags=cv2.INTER_CUBIC )
inp_f2 = cv2.warpAffine( f2, f2_mat, (inp_size, inp_size), flags=cv2.INTER_CUBIC )
dst_face_mask_a_0 = cv2.warpAffine( img_face_mask_a, face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC ) prd_f = self.predictor_func ( inp_f0, inp_f1, inp_f2 )
predictor_input_dst_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (self.predictor_input_size,self.predictor_input_size), cv2.INTER_CUBIC ) out_img = np.clip(prd_f, 0.0, 1.0)
prd_inp_dst_face_mask_a = predictor_input_dst_face_mask_a_0[...,np.newaxis]
prd_inp_avaperator_face_bgr = cv2.resize (avaperator_face_bgr, (self.predictor_input_size,self.predictor_input_size), cv2.INTER_CUBIC ) out_img = np.concatenate ( [cv2.resize ( inp_f1, (prd_f.shape[1], prd_f.shape[0]) ),
out_img], axis=1 )
prd_face_bgr = self.predictor_func ( prd_inp_avaperator_face_bgr, prd_inp_dst_face_mask_a )
out_img = img_bgr.copy()
out_img = cv2.warpAffine( prd_face_bgr, face_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT )
out_img = np.clip(out_img, 0.0, 1.0)
if debug: if debug:
debugs += [out_img.copy()] debugs += [out_img.copy()]
out_img = np.clip( img_bgr*(1-img_face_mask_aaa) + (out_img*img_face_mask_aaa) , 0, 1.0 )
if debug:
debugs += [out_img.copy()]
return debugs if debug else out_img return debugs if debug else out_img

15
converters/ConverterMasked.py
View file

@ -328,14 +328,13 @@ class ConverterMasked(Converter):
img_mask_blurry_aaa = img_face_mask_aaa img_mask_blurry_aaa = img_face_mask_aaa
if self.clip_hborder_mask_per > 0: # clip hborder before blur if self.clip_hborder_mask_per > 0: #clip hborder before blur
prd_hborder_rect_mask_a = np.ones(prd_face_mask_a.shape, dtype=np.float32) prd_hborder_rect_mask_a = np.ones ( prd_face_mask_a.shape, dtype=np.float32)
prd_border_size = int(prd_hborder_rect_mask_a.shape[1] * self.clip_hborder_mask_per) prd_border_size = int ( prd_hborder_rect_mask_a.shape[1] * self.clip_hborder_mask_per )
prd_hborder_rect_mask_a[:, 0:prd_border_size, :] = 0 prd_hborder_rect_mask_a[:,0:prd_border_size,:] = 0
prd_hborder_rect_mask_a[:, -prd_border_size:, :] = 0 prd_hborder_rect_mask_a[:,-prd_border_size:,:] = 0
prd_hborder_rect_mask_a[-prd_border_size:, :, :] = 0 prd_hborder_rect_mask_a[-prd_border_size:,:,:] = 0
prd_hborder_rect_mask_a = np.expand_dims( prd_hborder_rect_mask_a = np.expand_dims(cv2.blur(prd_hborder_rect_mask_a, (prd_border_size, prd_border_size) ),-1)
cv2.blur(prd_hborder_rect_mask_a, (prd_border_size, prd_border_size)), -1)
img_prd_hborder_rect_mask_a = cv2.warpAffine(prd_hborder_rect_mask_a, face_output_mat, img_size, img_prd_hborder_rect_mask_a = cv2.warpAffine(prd_hborder_rect_mask_a, face_output_mat, img_size,
np.zeros(img_bgr.shape, dtype=np.float32), np.zeros(img_bgr.shape, dtype=np.float32),

24
facelib/FaceType.py
View file

@ -4,10 +4,10 @@ class FaceType(IntEnum):
HALF = 0, HALF = 0,
FULL = 1, FULL = 1,
HEAD = 2, HEAD = 2,
AVATAR = 3, #centered nose only
MARK_ONLY = 4, #no align at all, just embedded faceinfo FULL_NO_ALIGN = 5,
QTY = 5 MARK_ONLY = 10, #no align at all, just embedded faceinfo
@staticmethod @staticmethod
def fromString (s): def fromString (s):
r = from_string_dict.get (s.lower()) r = from_string_dict.get (s.lower())
@ -17,17 +17,17 @@ class FaceType(IntEnum):
@staticmethod @staticmethod
def toString (face_type): def toString (face_type):
return to_string_list[face_type] return to_string_dict[face_type]
from_string_dict = {'half_face': FaceType.HALF, from_string_dict = {'half_face': FaceType.HALF,
'full_face': FaceType.FULL, 'full_face': FaceType.FULL,
'head' : FaceType.HEAD, 'head' : FaceType.HEAD,
'avatar' : FaceType.AVATAR,
'mark_only' : FaceType.MARK_ONLY, 'mark_only' : FaceType.MARK_ONLY,
'full_face_no_align' : FaceType.FULL_NO_ALIGN,
} }
to_string_list = [ 'half_face', to_string_dict = { FaceType.HALF : 'half_face',
'full_face', FaceType.FULL : 'full_face',
'head', FaceType.HEAD : 'head',
'avatar', FaceType.MARK_ONLY :'mark_only',
'mark_only' FaceType.FULL_NO_ALIGN : 'full_face_no_align'
] }

122
facelib/LandmarksProcessor.py
View file

@ -109,10 +109,19 @@ landmarks_68_3D = np.array( [
[0.205322 , 31.408738 , -21.903670 ], [0.205322 , 31.408738 , -21.903670 ],
[-7.198266 , 30.844876 , -20.328022 ] ], dtype=np.float32) [-7.198266 , 30.844876 , -20.328022 ] ], dtype=np.float32)
def transform_points(points, mat, invert=False):
if invert:
mat = cv2.invertAffineTransform (mat)
points = np.expand_dims(points, axis=1)
points = cv2.transform(points, mat, points.shape)
points = np.squeeze(points)
return points
def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0): def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0):
if not isinstance(image_landmarks, np.ndarray): if not isinstance(image_landmarks, np.ndarray):
image_landmarks = np.array (image_landmarks) image_landmarks = np.array (image_landmarks)
"""
if face_type == FaceType.AVATAR: if face_type == FaceType.AVATAR:
centroid = np.mean (image_landmarks, axis=0) centroid = np.mean (image_landmarks, axis=0)
@ -128,76 +137,79 @@ def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0):
mat = mat * scale * (output_size / 3) mat = mat * scale * (output_size / 3)
mat[:,2] += output_size / 2 mat[:,2] += output_size / 2
else: else:
if face_type == FaceType.HALF: """
padding = 0 remove_align = False
elif face_type == FaceType.FULL: if face_type == FaceType.FULL_NO_ALIGN:
padding = (output_size / 64) * 12 face_type = FaceType.FULL
elif face_type == FaceType.HEAD: remove_align = True
padding = (output_size / 64) * 24
else: if face_type == FaceType.HALF:
raise ValueError ('wrong face_type: ', face_type) padding = 0
elif face_type == FaceType.FULL:
padding = (output_size / 64) * 12
elif face_type == FaceType.HEAD:
padding = (output_size / 64) * 24
else:
raise ValueError ('wrong face_type: ', face_type)
mat = umeyama(image_landmarks[17:], landmarks_2D, True)[0:2] mat = umeyama(image_landmarks[17:], landmarks_2D, True)[0:2]
mat = mat * (output_size - 2 * padding) mat = mat * (output_size - 2 * padding)
mat[:,2] += padding mat[:,2] += padding
mat *= (1 / scale) mat *= (1 / scale)
mat[:,2] += -output_size*( ( (1 / scale) - 1.0 ) / 2 ) mat[:,2] += -output_size*( ( (1 / scale) - 1.0 ) / 2 )
if remove_align:
bbox = transform_points ( [ (0,0), (0,output_size-1), (output_size-1, output_size-1), (output_size-1,0) ], mat, True)
area = mathlib.polygon_area(bbox[:,0], bbox[:,1] )
side = math.sqrt(area) / 2
center = transform_points ( [(output_size/2,output_size/2)], mat, True)
pts1 = np.float32([ center+[-side,-side], center+[side,-side], center+[-side,side] ])
pts2 = np.float32([[0,0],[output_size-1,0],[0,output_size-1]])
mat = cv2.getAffineTransform(pts1,pts2)
return mat return mat
def transform_points(points, mat, invert=False):
if invert:
mat = cv2.invertAffineTransform (mat)
points = np.expand_dims(points, axis=1)
points = cv2.transform(points, mat, points.shape)
points = np.squeeze(points)
return points
def get_image_hull_mask (image_shape, image_landmarks, ie_polys=None): def get_image_hull_mask (image_shape, image_landmarks, ie_polys=None):
if len(image_landmarks) != 68: if len(image_landmarks) != 68:
raise Exception('get_image_hull_mask works only with 68 landmarks') raise Exception('get_image_hull_mask works only with 68 landmarks')
int_lmrks = np.array(image_landmarks, dtype=np.int) int_lmrks = np.array(image_landmarks.copy(), dtype=np.int)
hull_mask = np.zeros(image_shape[0:2]+(1,),dtype=np.float32) hull_mask = np.zeros(image_shape[0:2]+(1,),dtype=np.float32)
cv2.fillConvexPoly( hull_mask, cv2.convexHull( # #nose
np.concatenate ( (int_lmrks[0:9], ml_pnt = (int_lmrks[36] + int_lmrks[0]) // 2
int_lmrks[17:18]))) , (1,) ) mr_pnt = (int_lmrks[16] + int_lmrks[45]) // 2
cv2.fillConvexPoly( hull_mask, cv2.convexHull( # mid points between the mid points and eye
np.concatenate ( (int_lmrks[8:17], ql_pnt = (int_lmrks[36] + ml_pnt) // 2
int_lmrks[26:27]))) , (1,) ) qr_pnt = (int_lmrks[45] + mr_pnt) // 2
cv2.fillConvexPoly( hull_mask, cv2.convexHull( # Top of the eye arrays
np.concatenate ( (int_lmrks[17:20], bot_l = np.array((ql_pnt, int_lmrks[36], int_lmrks[37], int_lmrks[38], int_lmrks[39]))
int_lmrks[8:9]))) , (1,) ) bot_r = np.array((int_lmrks[42], int_lmrks[43], int_lmrks[44], int_lmrks[45], qr_pnt))
cv2.fillConvexPoly( hull_mask, cv2.convexHull( # Eyebrow arrays
np.concatenate ( (int_lmrks[24:27], top_l = int_lmrks[17:22]
int_lmrks[8:9]))) , (1,) ) top_r = int_lmrks[22:27]
cv2.fillConvexPoly( hull_mask, cv2.convexHull( # Adjust eyebrow arrays
np.concatenate ( (int_lmrks[19:25], int_lmrks[17:22] = top_l + ((top_l - bot_l) // 2)
int_lmrks[8:9], int_lmrks[22:27] = top_r + ((top_r - bot_r) // 2)
))) , (1,) )
cv2.fillConvexPoly( hull_mask, cv2.convexHull( r_jaw = (int_lmrks[0:9], int_lmrks[17:18])
np.concatenate ( (int_lmrks[17:22], l_jaw = (int_lmrks[8:17], int_lmrks[26:27])
int_lmrks[27:28], r_cheek = (int_lmrks[17:20], int_lmrks[8:9])
int_lmrks[31:36], l_cheek = (int_lmrks[24:27], int_lmrks[8:9])
int_lmrks[8:9] nose_ridge = (int_lmrks[19:25], int_lmrks[8:9],)
))) , (1,) ) r_eye = (int_lmrks[17:22], int_lmrks[27:28], int_lmrks[31:36], int_lmrks[8:9])
l_eye = (int_lmrks[22:27], int_lmrks[27:28], int_lmrks[31:36], int_lmrks[8:9])
nose = (int_lmrks[27:31], int_lmrks[31:36])
parts = [r_jaw, l_jaw, r_cheek, l_cheek, nose_ridge, r_eye, l_eye, nose]
cv2.fillConvexPoly( hull_mask, cv2.convexHull( for item in parts:
np.concatenate ( (int_lmrks[22:27], merged = np.concatenate(item)
int_lmrks[27:28], cv2.fillConvexPoly(hull_mask, cv2.convexHull(merged), 1)
int_lmrks[31:36],
int_lmrks[8:9]
))) , (1,) )
#nose
cv2.fillConvexPoly( hull_mask, cv2.convexHull(int_lmrks[27:36]), (1,) )
if ie_polys is not None: if ie_polys is not None:
ie_polys.overlay_mask(hull_mask) ie_polys.overlay_mask(hull_mask)
@ -309,7 +321,7 @@ def draw_landmarks (image, image_landmarks, color=(0,255,0), transparent_mask=Fa
mask = get_image_hull_mask (image.shape, image_landmarks, ie_polys) mask = get_image_hull_mask (image.shape, image_landmarks, ie_polys)
image[...] = ( image * (1-mask) + image * mask / 2 )[...] image[...] = ( image * (1-mask) + image * mask / 2 )[...]
def draw_rect_landmarks (image, rect, image_landmarks, face_size, face_type, transparent_mask=False, ie_polys=None, landmarks_color=(0,255,0) ): def draw_rect_landmarks (image, rect, image_landmarks, face_size, face_type, transparent_mask=False, ie_polys=None, landmarks_color=(0,255,0)):
draw_landmarks(image, image_landmarks, color=landmarks_color, transparent_mask=transparent_mask, ie_polys=ie_polys) draw_landmarks(image, image_landmarks, color=landmarks_color, transparent_mask=transparent_mask, ie_polys=ie_polys)
imagelib.draw_rect (image, rect, (255,0,0), 2 ) imagelib.draw_rect (image, rect, (255,0,0), 2 )

4
main.py
View file

@ -40,7 +40,7 @@ if __name__ == "__main__":
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the extracted files will be stored.") p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the extracted files will be stored.")
p.add_argument('--debug-dir', action=fixPathAction, dest="debug_dir", help="Writes debug images to this directory.") p.add_argument('--debug-dir', action=fixPathAction, dest="debug_dir", help="Writes debug images to this directory.")
p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'head', 'avatar', 'mark_only'], default='full_face', help="Default 'full_face'. Don't change this option, currently all models uses 'full_face'") p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'head', 'full_face_no_align', 'mark_only'], default='full_face', help="Default 'full_face'. Don't change this option, currently all models uses 'full_face'")
p.add_argument('--detector', dest="detector", choices=['dlib','mt','s3fd','manual'], default='dlib', help="Type of detector. Default 'dlib'. 'mt' (MTCNNv1) - faster, better, almost no jitter, perfect for gathering thousands faces for src-set. It is also good for dst-set, but can generate false faces in frames where main face not recognized! In this case for dst-set use either 'dlib' with '--manual-fix' or '--detector manual'. Manual detector suitable only for dst-set.") p.add_argument('--detector', dest="detector", choices=['dlib','mt','s3fd','manual'], default='dlib', help="Type of detector. Default 'dlib'. 'mt' (MTCNNv1) - faster, better, almost no jitter, perfect for gathering thousands faces for src-set. It is also good for dst-set, but can generate false faces in frames where main face not recognized! In this case for dst-set use either 'dlib' with '--manual-fix' or '--detector manual'. Manual detector suitable only for dst-set.")
p.add_argument('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="Enables multi GPU.") p.add_argument('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="Enables multi GPU.")
p.add_argument('--manual-fix', action="store_true", dest="manual_fix", default=False, help="Enables manual extract only frames where faces were not recognized.") p.add_argument('--manual-fix', action="store_true", dest="manual_fix", default=False, help="Enables manual extract only frames where faces were not recognized.")
@ -158,7 +158,6 @@ if __name__ == "__main__":
args = {'input_dir' : arguments.input_dir, args = {'input_dir' : arguments.input_dir,
'output_dir' : arguments.output_dir, 'output_dir' : arguments.output_dir,
'aligned_dir' : arguments.aligned_dir, 'aligned_dir' : arguments.aligned_dir,
'avaperator_aligned_dir' : arguments.avaperator_aligned_dir,
'model_dir' : arguments.model_dir, 'model_dir' : arguments.model_dir,
'model_name' : arguments.model_name, 'model_name' : arguments.model_name,
'debug' : arguments.debug, 'debug' : arguments.debug,
@ -173,7 +172,6 @@ if __name__ == "__main__":
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.") p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the converted files will be stored.") p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the converted files will be stored.")
p.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", help="Aligned directory. This is where the extracted of dst faces stored.") p.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", help="Aligned directory. This is where the extracted of dst faces stored.")
p.add_argument('--avaperator-aligned-dir', action=fixPathAction, dest="avaperator_aligned_dir", help="Only for AVATAR model. Directory of aligned avatar operator faces.")
p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.") p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.")
p.add_argument('--model', required=True, dest="model_name", choices=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Type of model") p.add_argument('--model', required=True, dest="model_name", choices=Path_utils.get_all_dir_names_startswith ( Path(__file__).parent / 'models' , 'Model_'), help="Type of model")
p.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug converter.") p.add_argument('--debug', action="store_true", dest="debug", default=False, help="Debug converter.")

199
mainscripts/Converter.py
View file

@ -28,9 +28,9 @@ class ConvertSubprocessor(Subprocessor):
self.device_idx = client_dict['device_idx'] self.device_idx = client_dict['device_idx']
self.device_name = client_dict['device_name'] self.device_name = client_dict['device_name']
self.converter = client_dict['converter'] self.converter = client_dict['converter']
self.input_data = client_dict['input_data']
self.output_path = Path(client_dict['output_dir']) if 'output_dir' in client_dict.keys() else None self.output_path = Path(client_dict['output_dir']) if 'output_dir' in client_dict.keys() else None
self.alignments = client_dict['alignments'] self.alignments = client_dict['alignments']
self.avatar_image_paths = client_dict['avatar_image_paths']
self.debug = client_dict['debug'] self.debug = client_dict['debug']
#transfer and set stdin in order to work code.interact in debug subprocess #transfer and set stdin in order to work code.interact in debug subprocess
@ -50,65 +50,30 @@ class ConvertSubprocessor(Subprocessor):
#override #override
def process_data(self, data): def process_data(self, data):
idx, filename = data
filename_path = Path(filename)
files_processed = 1 files_processed = 1
faces_processed = 0 faces_processed = 0
idx, = data
filename = self.input_data[idx][0]
filename_path = Path(filename)
output_filename_path = self.output_path / (filename_path.stem + '.png') output_filename_path = self.output_path / (filename_path.stem + '.png')
image = None
if self.converter.type == Converter.TYPE_FACE:
if filename_path.stem not in self.alignments.keys():
if not self.debug:
self.log_info ( 'no faces found for %s, copying without faces' % (filename_path.name) )
if (self.converter.type == Converter.TYPE_FACE or self.converter.type == Converter.TYPE_FACE_AVATAR ) \ if filename_path.suffix == '.png':
and filename_path.stem not in self.alignments.keys(): shutil.copy ( str(filename_path), str(output_filename_path) )
if not self.debug: else:
self.log_info ( 'no faces found for %s, copying without faces' % (filename_path.name) ) image = cv2_imread(str(filename_path))
cv2_imwrite ( str(output_filename_path), image )
if filename_path.suffix == '.png': else:
shutil.copy ( str(filename_path), str(output_filename_path) ) image = (cv2_imread(str(filename_path)) / 255.0).astype(np.float32)
else: image = normalize_channels (image, 3)
image = cv2_imread(str(filename_path))
cv2_imwrite ( str(output_filename_path), image )
else:
image = (cv2_imread(str(filename_path)) / 255.0).astype(np.float32)
image = normalize_channels (image, 3)
if self.converter.type == Converter.TYPE_IMAGE:
image = self.converter.cli_convert_image(image, None, self.debug)
if self.debug:
return (1, image)
faces_processed = 1
elif self.converter.type == Converter.TYPE_IMAGE_WITH_LANDMARKS:
#currently unused
if filename_path.suffix == '.png':
dflimg = DFLPNG.load( str(filename_path) )
elif filename_path.suffix == '.jpg':
dflimg = DFLJPG.load ( str(filename_path) )
else:
dflimg = None
if dflimg is not None:
image_landmarks = dflimg.get_landmarks()
image = self.converter.convert_image(image, image_landmarks, self.debug)
if self.debug:
raise NotImplementedError
#for img in image:
# io.show_image ('Debug convert', img )
# cv2.waitKey(0)
faces_processed = 1
else:
self.log_err ("%s is not a dfl image file" % (filename_path.name) )
elif self.converter.type == Converter.TYPE_FACE or self.converter.type == Converter.TYPE_FACE_AVATAR:
ava_face = None
if self.converter.type == Converter.TYPE_FACE_AVATAR:
ava_filename_path = self.avatar_image_paths[idx]
ava_face = (cv2_imread(str(ava_filename_path)) / 255.0).astype(np.float32)
ava_face = normalize_channels (ava_face, 3)
faces = self.alignments[filename_path.stem] faces = self.alignments[filename_path.stem]
if self.debug: if self.debug:
@ -120,9 +85,9 @@ class ConvertSubprocessor(Subprocessor):
self.log_info ( '\nConverting face_num [%d] in file [%s]' % (face_num, filename_path) ) self.log_info ( '\nConverting face_num [%d] in file [%s]' % (face_num, filename_path) )
if self.debug: if self.debug:
debug_images += self.converter.cli_convert_face(image, image_landmarks, self.debug, avaperator_face_bgr=ava_face) debug_images += self.converter.cli_convert_face(image, image_landmarks, self.debug)
else: else:
image = self.converter.cli_convert_face(image, image_landmarks, self.debug, avaperator_face_bgr=ava_face) image = self.converter.cli_convert_face(image, image_landmarks, self.debug)
except Exception as e: except Exception as e:
e_str = traceback.format_exc() e_str = traceback.format_exc()
@ -135,29 +100,60 @@ class ConvertSubprocessor(Subprocessor):
return (1, debug_images) return (1, debug_images)
faces_processed = len(faces) faces_processed = len(faces)
elif self.converter.type == Converter.TYPE_IMAGE:
image = (cv2_imread(str(filename_path)) / 255.0).astype(np.float32)
image = normalize_channels (image, 3)
image = self.converter.cli_convert_image(image, None, self.debug)
if not self.debug: if self.debug:
cv2_imwrite (str(output_filename_path), (image*255).astype(np.uint8) ) return (1, image)
faces_processed = 1
elif self.converter.type == Converter.TYPE_FACE_AVATAR:
max_idx = len(self.input_data)-1
i0 = max (idx-1, 0)
i1 = idx
i2 = min (max_idx, idx+1)
f0 = (cv2_imread( self.input_data[i0][0] ) / 255.0).astype(np.float32)
f0_lmrk = self.input_data[i0][1]
f1 = (cv2_imread( self.input_data[i1][0] ) / 255.0).astype(np.float32)
f1_lmrk = self.input_data[i1][1]
f2 = (cv2_imread( self.input_data[i2][0] ) / 255.0).astype(np.float32)
f2_lmrk = self.input_data[i2][1]
f0, f1, f2 = [ normalize_channels (f, 3) for f in [f0,f1,f2] ]
image = self.converter.cli_convert_face(f0, f0_lmrk, f1, f1_lmrk, f2, f2_lmrk, self.debug)
output_filename_path = self.output_path / self.input_data[idx][2]
if self.debug:
return (1, image)
faces_processed = 1
if image is not None and not self.debug:
cv2_imwrite (str(output_filename_path), (image*255).astype(np.uint8) )
return (0, files_processed, faces_processed) return (0, files_processed, faces_processed)
#overridable #overridable
def get_data_name (self, data): def get_data_name (self, data):
#return string identificator of your data #return string identificator of your data
idx, filename = data idx, = data
return filename return self.input_data[idx][0]
#override #override
def __init__(self, converter, input_path_image_paths, output_path, alignments, avatar_image_paths=None, debug = False): def __init__(self, converter, input_data, output_path, alignments, debug = False):
super().__init__('Converter', ConvertSubprocessor.Cli, 86400 if debug == True else 60) super().__init__('Converter', ConvertSubprocessor.Cli, 86400 if debug == True else 60)
self.converter = converter self.converter = converter
self.input_data = self.input_path_image_paths = input_path_image_paths self.input_data = input_data
self.input_data_idxs = [ *range(len(self.input_data)) ] self.input_data_idxs = [ *range(len(self.input_data)) ]
self.output_path = output_path self.output_path = output_path
self.alignments = alignments self.alignments = alignments
self.avatar_image_paths = avatar_image_paths
self.debug = debug self.debug = debug
self.files_processed = 0 self.files_processed = 0
@ -171,9 +167,9 @@ class ConvertSubprocessor(Subprocessor):
yield 'CPU%d' % (i), {}, {'device_idx': i, yield 'CPU%d' % (i), {}, {'device_idx': i,
'device_name': 'CPU%d' % (i), 'device_name': 'CPU%d' % (i),
'converter' : self.converter, 'converter' : self.converter,
'input_data' : self.input_data,
'output_dir' : str(self.output_path), 'output_dir' : str(self.output_path),
'alignments' : self.alignments, 'alignments' : self.alignments,
'avatar_image_paths' : self.avatar_image_paths,
'debug': self.debug, 'debug': self.debug,
'stdin_fd': sys.stdin.fileno() if self.debug else None 'stdin_fd': sys.stdin.fileno() if self.debug else None
} }
@ -196,12 +192,12 @@ class ConvertSubprocessor(Subprocessor):
def get_data(self, host_dict): def get_data(self, host_dict):
if len (self.input_data_idxs) > 0: if len (self.input_data_idxs) > 0:
idx = self.input_data_idxs.pop(0) idx = self.input_data_idxs.pop(0)
return (idx, self.input_data[idx]) return (idx, )
return None return None
#override #override
def on_data_return (self, host_dict, data): def on_data_return (self, host_dict, data):
idx, filename = data idx, = data
self.input_data_idxs.insert(0, idx) self.input_data_idxs.insert(0, idx)
#override #override
@ -253,9 +249,9 @@ def main (args, device_args):
converter = model.get_converter() converter = model.get_converter()
input_path_image_paths = Path_utils.get_image_paths(input_path) input_path_image_paths = Path_utils.get_image_paths(input_path)
alignments = None alignments = None
avatar_image_paths = None if converter.type == Converter.TYPE_FACE:
if converter.type == Converter.TYPE_FACE or converter.type == Converter.TYPE_FACE_AVATAR:
if aligned_dir is None: if aligned_dir is None:
io.log_err('Aligned directory not found. Please ensure it exists.') io.log_err('Aligned directory not found. Please ensure it exists.')
return return
@ -287,21 +283,15 @@ def main (args, device_args):
alignments[ source_filename_stem ] = [] alignments[ source_filename_stem ] = []
alignments[ source_filename_stem ].append (dflimg.get_source_landmarks()) alignments[ source_filename_stem ].append (dflimg.get_source_landmarks())
#avatar_alignments += [ ( str(filepath), dflimg.get_source_landmarks(), dflimg.get_source_filename() ) ]
if converter.type == Converter.TYPE_FACE_AVATAR: input_data = [ (p,) for p in input_path_image_paths ]
if avaperator_aligned_dir is None: elif converter.type == Converter.TYPE_FACE_AVATAR:
io.log_err('Avatar operator aligned directory not found. Please ensure it exists.')
return input_data = []
for filepath in io.progress_bar_generator(input_path_image_paths, "Collecting info"):
filepath = Path(filepath)
avaperator_aligned_path = Path(avaperator_aligned_dir)
if not avaperator_aligned_path.exists():
io.log_err('Avatar operator aligned directory not found. Please ensure it exists.')
return
avatar_image_paths = []
for filename in io.progress_bar_generator( Path_utils.get_image_paths(avaperator_aligned_path) , "Sorting avaperator faces"):
filepath = Path(filename)
if filepath.suffix == '.png': if filepath.suffix == '.png':
dflimg = DFLPNG.load( str(filepath) ) dflimg = DFLPNG.load( str(filepath) )
elif filepath.suffix == '.jpg': elif filepath.suffix == '.jpg':
@ -310,22 +300,19 @@ def main (args, device_args):
dflimg = None dflimg = None
if dflimg is None: if dflimg is None:
io.log_err ("Fatal error: %s is not a dfl image file" % (filepath.name) ) io.log_err ("%s is not a dfl image file" % (filepath.name) )
return continue
input_data += [ ( str(filepath), dflimg.get_landmarks(), dflimg.get_source_filename() ) ]
avatar_image_paths += [ (filename, dflimg.get_source_filename() ) ]
avatar_image_paths = [ p[0] for p in sorted(avatar_image_paths, key=operator.itemgetter(1)) ]
if len(input_path_image_paths) < len(avatar_image_paths): input_data = sorted(input_data, key=operator.itemgetter(2))
io.log_err("Input faces count must be >= avatar operator faces count.") else:
return input_data = [ (p,) for p in input_path_image_paths ]
files_processed, faces_processed = ConvertSubprocessor ( files_processed, faces_processed = ConvertSubprocessor (
converter = converter, converter = converter,
input_path_image_paths = input_path_image_paths, input_data = input_data,
output_path = output_path, output_path = output_path,
alignments = alignments, alignments = alignments,
avatar_image_paths = avatar_image_paths,
debug = args.get('debug',False) debug = args.get('debug',False)
).run() ).run()
@ -389,3 +376,29 @@ if model_name == 'AVATAR':
# new_points = np.concatenate( [np.expand_dims(p1,-1),np.expand_dims(p2,-1)], -1 ) # new_points = np.concatenate( [np.expand_dims(p1,-1),np.expand_dims(p2,-1)], -1 )
# #
# alignments[ a[i] ][0] = LandmarksProcessor.transform_points (new_points, m0, True).astype(np.int32) # alignments[ a[i] ][0] = LandmarksProcessor.transform_points (new_points, m0, True).astype(np.int32)
"""
elif self.converter.type == Converter.TYPE_IMAGE_WITH_LANDMARKS:
#currently unused
if filename_path.suffix == '.png':
dflimg = DFLPNG.load( str(filename_path) )
elif filename_path.suffix == '.jpg':
dflimg = DFLJPG.load ( str(filename_path) )
else:
dflimg = None
if dflimg is not None:
image_landmarks = dflimg.get_landmarks()
image = self.converter.convert_image(image, image_landmarks, self.debug)
if self.debug:
raise NotImplementedError
#for img in image:
# io.show_image ('Debug convert', img )
# cv2.waitKey(0)
faces_processed = 1
else:
self.log_err ("%s is not a dfl image file" % (filename_path.name) )
"""

21
mainscripts/Extractor.py
View file

@ -6,6 +6,7 @@ import multiprocessing
import shutil import shutil
from pathlib import Path from pathlib import Path
import numpy as np import numpy as np
import math
import mathlib import mathlib
import imagelib import imagelib
import cv2 import cv2
@ -21,6 +22,8 @@ from nnlib import nnlib
from joblib import Subprocessor from joblib import Subprocessor
from interact import interact as io from interact import interact as io
DEBUG = False
class ExtractSubprocessor(Subprocessor): class ExtractSubprocessor(Subprocessor):
class Data(object): class Data(object):
def __init__(self, filename=None, rects=None, landmarks = None, landmarks_accurate=True, pitch_yaw_roll=None, final_output_files = None): def __init__(self, filename=None, rects=None, landmarks = None, landmarks_accurate=True, pitch_yaw_roll=None, final_output_files = None):
@ -44,6 +47,11 @@ class ExtractSubprocessor(Subprocessor):
self.cpu_only = client_dict['device_type'] == 'CPU' self.cpu_only = client_dict['device_type'] == 'CPU'
self.final_output_path = Path(client_dict['final_output_dir']) if 'final_output_dir' in client_dict.keys() else None self.final_output_path = Path(client_dict['final_output_dir']) if 'final_output_dir' in client_dict.keys() else None
self.debug_dir = client_dict['debug_dir'] self.debug_dir = client_dict['debug_dir']
#transfer and set stdin in order to work code.interact in debug subprocess
stdin_fd = client_dict['stdin_fd']
if stdin_fd is not None and DEBUG:
sys.stdin = os.fdopen(stdin_fd)
self.cached_image = (None, None) self.cached_image = (None, None)
@ -224,10 +232,12 @@ class ExtractSubprocessor(Subprocessor):
rect = np.array(rect) rect = np.array(rect)
if self.face_type == FaceType.MARK_ONLY: if self.face_type == FaceType.MARK_ONLY:
image_to_face_mat = None
face_image = image face_image = image
face_image_landmarks = image_landmarks face_image_landmarks = image_landmarks
else: else:
image_to_face_mat = LandmarksProcessor.get_transform_mat (image_landmarks, self.image_size, self.face_type) image_to_face_mat = LandmarksProcessor.get_transform_mat (image_landmarks, self.image_size, self.face_type)
face_image = cv2.warpAffine(image, image_to_face_mat, (self.image_size, self.image_size), cv2.INTER_LANCZOS4) face_image = cv2.warpAffine(image, image_to_face_mat, (self.image_size, self.image_size), cv2.INTER_LANCZOS4)
face_image_landmarks = LandmarksProcessor.transform_points (image_landmarks, image_to_face_mat) face_image_landmarks = LandmarksProcessor.transform_points (image_landmarks, image_to_face_mat)
@ -239,8 +249,8 @@ class ExtractSubprocessor(Subprocessor):
if landmarks_area > 4*rect_area: #get rid of faces which umeyama-landmark-area > 4*detector-rect-area if landmarks_area > 4*rect_area: #get rid of faces which umeyama-landmark-area > 4*detector-rect-area
continue continue
if self.debug_dir is not None: if self.debug_dir is not None:
LandmarksProcessor.draw_rect_landmarks (debug_image, rect, image_landmarks, self.image_size, self.face_type, transparent_mask=True) LandmarksProcessor.draw_rect_landmarks (debug_image, rect, image_landmarks, self.image_size, self.face_type, transparent_mask=True)
if filename_path.suffix == '.jpg': if filename_path.suffix == '.jpg':
@ -308,7 +318,7 @@ class ExtractSubprocessor(Subprocessor):
self.devices = ExtractSubprocessor.get_devices_for_config(self.manual, self.type, multi_gpu, cpu_only) self.devices = ExtractSubprocessor.get_devices_for_config(self.manual, self.type, multi_gpu, cpu_only)
no_response_time_sec = 60 if not self.manual else 999999 no_response_time_sec = 60 if not self.manual and not DEBUG else 999999
super().__init__('Extractor', ExtractSubprocessor.Cli, no_response_time_sec) super().__init__('Extractor', ExtractSubprocessor.Cli, no_response_time_sec)
#override #override
@ -354,7 +364,8 @@ class ExtractSubprocessor(Subprocessor):
'image_size': self.image_size, 'image_size': self.image_size,
'face_type': self.face_type, 'face_type': self.face_type,
'debug_dir': self.debug_dir, 'debug_dir': self.debug_dir,
'final_output_dir': str(self.final_output_path)} 'final_output_dir': str(self.final_output_path),
'stdin_fd': sys.stdin.fileno() }
for (device_idx, device_type, device_name, device_total_vram_gb) in self.devices: for (device_idx, device_type, device_name, device_total_vram_gb) in self.devices:
@ -632,7 +643,7 @@ class ExtractSubprocessor(Subprocessor):
return [ (i, 'CPU', 'CPU%d' % (i), 0 ) for i in range( min(8, multiprocessing.cpu_count() // 2) ) ] return [ (i, 'CPU', 'CPU%d' % (i), 0 ) for i in range( min(8, multiprocessing.cpu_count() // 2) ) ]
elif type == 'final': elif type == 'final':
return [ (i, 'CPU', 'CPU%d' % (i), 0 ) for i in range(min(8, multiprocessing.cpu_count())) ] return [ (i, 'CPU', 'CPU%d' % (i), 0 ) for i in (range(min(8, multiprocessing.cpu_count())) if not DEBUG else [0]) ]
class DeletedFilesSearcherSubprocessor(Subprocessor): class DeletedFilesSearcherSubprocessor(Subprocessor):
class Cli(Subprocessor.Cli): class Cli(Subprocessor.Cli):

11
mainscripts/Sorter.py
View file

@ -423,14 +423,11 @@ def sort_by_hist_dissim(input_path):
else: else:
dflimg = None dflimg = None
if dflimg is None:
io.log_err ("%s is not a dfl image file" % (filepath.name) )
trash_img_list.append ([str(filepath)])
continue
image = cv2_imread(str(filepath)) image = cv2_imread(str(filepath))
face_mask = LandmarksProcessor.get_image_hull_mask (image.shape, dflimg.get_landmarks())
image = (image*face_mask).astype(np.uint8) if dflimg is not None:
face_mask = LandmarksProcessor.get_image_hull_mask (image.shape, dflimg.get_landmarks())
image = (image*face_mask).astype(np.uint8)
img_list.append ([str(filepath), cv2.calcHist([cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)], [0], None, [256], [0, 256]), 0 ]) img_list.append ([str(filepath), cv2.calcHist([cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)], [0], None, [256], [0, 256]), 0 ])

2
models/ModelBase.py
View file

@ -521,7 +521,7 @@ class ModelBase(object):
return imagelib.equalize_and_stack_square(images) return imagelib.equalize_and_stack_square(images)
def generate_next_sample(self): def generate_next_sample(self):
return [next(generator) for generator in self.generator_list] return [ generator.generate_next() for generator in self.generator_list]
def train_one_iter(self): def train_one_iter(self):

1
models/Model_RecycleGAN/Model.py
View file

@ -378,7 +378,6 @@ class RecycleGANModel(ModelBase):
return x return x
return func return func
nnlib.UNet = UNet
@staticmethod @staticmethod
def UNetTemporalPredictor(output_nc, use_batch_norm, ngf=64, use_dropout=False): def UNetTemporalPredictor(output_nc, use_batch_norm, ngf=64, use_dropout=False):

6
models/Model_SAE/Model.py
View file

@ -705,6 +705,12 @@ class SAEModel(ModelBase):
return func return func
SAEModel.downscale = downscale SAEModel.downscale = downscale
#def downscale (dim, padding='zero', norm='', act='', **kwargs):
# def func(x):
# return BlurPool()( Norm(norm)( Act(act) (Conv2D(dim, kernel_size=5, strides=1, padding=padding)(x)) ) )
# return func
#SAEModel.downscale = downscale
def upscale(dim, padding='zero', norm='', act='', **kwargs): def upscale(dim, padding='zero', norm='', act='', **kwargs):
def func(x): def func(x):

84
nnlib/nnlib.py
View file

@ -89,6 +89,8 @@ dssim = nnlib.dssim
PixelShuffler = nnlib.PixelShuffler PixelShuffler = nnlib.PixelShuffler
SubpixelUpscaler = nnlib.SubpixelUpscaler SubpixelUpscaler = nnlib.SubpixelUpscaler
Scale = nnlib.Scale Scale = nnlib.Scale
BlurPool = nnlib.BlurPool
SelfAttention = nnlib.SelfAttention
CAInitializerMP = nnlib.CAInitializerMP CAInitializerMP = nnlib.CAInitializerMP
@ -462,6 +464,51 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
nnlib.PixelShuffler = PixelShuffler nnlib.PixelShuffler = PixelShuffler
nnlib.SubpixelUpscaler = PixelShuffler nnlib.SubpixelUpscaler = PixelShuffler
class BlurPool(KL.Layer):
"""
https://arxiv.org/abs/1904.11486 https://github.com/adobe/antialiased-cnns
"""
def __init__(self, filt_size=3, stride=2, **kwargs):
self.strides = (stride,stride)
self.filt_size = filt_size
self.padding = ( (int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)) ), (int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)) ) )
if(self.filt_size==1):
self.a = np.array([1.,])
elif(self.filt_size==2):
self.a = np.array([1., 1.])
elif(self.filt_size==3):
self.a = np.array([1., 2., 1.])
elif(self.filt_size==4):
self.a = np.array([1., 3., 3., 1.])
elif(self.filt_size==5):
self.a = np.array([1., 4., 6., 4., 1.])
elif(self.filt_size==6):
self.a = np.array([1., 5., 10., 10., 5., 1.])
elif(self.filt_size==7):
self.a = np.array([1., 6., 15., 20., 15., 6., 1.])
super(BlurPool, self).__init__(**kwargs)
def compute_output_shape(self, input_shape):
height = input_shape[1] // self.strides[0]
width = input_shape[2] // self.strides[1]
channels = input_shape[3]
return (input_shape[0], height, width, channels)
def call(self, x):
k = self.a
k = k[:,None]*k[None,:]
k = k / np.sum(k)
k = np.tile (k[:,:,None,None], (1,1,K.int_shape(x)[-1],1) )
k = K.constant (k, dtype=K.floatx() )
x = K.spatial_2d_padding(x, padding=self.padding)
x = K.depthwise_conv2d(x, k, strides=self.strides, padding='valid')
return x
nnlib.BlurPool = BlurPool
class Scale(KL.Layer): class Scale(KL.Layer):
""" """
@ -496,6 +543,43 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
nnlib.Scale = Scale nnlib.Scale = Scale
class SelfAttention(KL.Layer):
def __init__(self, nc, squeeze_factor=8, **kwargs):
assert nc//squeeze_factor > 0, f"Input channels must be >= {squeeze_factor}, recieved nc={nc}"
self.nc = nc
self.squeeze_factor = squeeze_factor
super(SelfAttention, self).__init__(**kwargs)
def compute_output_shape(self, input_shape):
return (input_shape[0], input_shape[1], input_shape[2], self.nc)
def call(self, inp):
x = inp
shape_x = x.get_shape().as_list()
f = Conv2D(self.nc//self.squeeze_factor, 1, kernel_regularizer=keras.regularizers.l2(1e-4))(x)
g = Conv2D(self.nc//self.squeeze_factor, 1, kernel_regularizer=keras.regularizers.l2(1e-4))(x)
h = Conv2D(self.nc, 1, kernel_regularizer=keras.regularizers.l2(1e-4))(x)
shape_f = f.get_shape().as_list()
shape_g = g.get_shape().as_list()
shape_h = h.get_shape().as_list()
flat_f = Reshape( (-1, shape_f[-1]) )(f)
flat_g = Reshape( (-1, shape_g[-1]) )(g)
flat_h = Reshape( (-1, shape_h[-1]) )(h)
s = Lambda(lambda x: K.batch_dot(x[0], keras.layers.Permute((2,1))(x[1]) ))([flat_g, flat_f])
beta = keras.layers.Softmax(axis=-1)(s)
o = Lambda(lambda x: K.batch_dot(x[0], x[1]))([beta, flat_h])
o = Reshape(shape_x[1:])(o)
o = Scale()(o)
out = Add()([o, inp])
return out
nnlib.SelfAttention = SelfAttention
class Adam(keras.optimizers.Optimizer): class Adam(keras.optimizers.Optimizer):
"""Adam optimizer. """Adam optimizer.

13
samplelib/SampleGeneratorBase.py
View file

@ -13,7 +13,18 @@ class SampleGeneratorBase(object):
self.samples_path = Path(samples_path) self.samples_path = Path(samples_path)
self.debug = debug self.debug = debug
self.batch_size = 1 if self.debug else batch_size self.batch_size = 1 if self.debug else batch_size
self.last_generation = None
self.active = True
def set_active(self, is_active):
self.active = is_active
def generate_next(self):
if not self.active and self.last_generation is not None:
return self.last_generation
self.last_generation = next(self)
return self.last_generation
#overridable #overridable
def __iter__(self): def __iter__(self):
#implement your own iterator #implement your own iterator

8
samplelib/SampleGeneratorFaceTemporal.py
View file

@ -8,7 +8,7 @@ from samplelib import SampleType, SampleProcessor, SampleLoader, SampleGenerator
''' '''
output_sample_types = [ output_sample_types = [
[SampleProcessor.TypeFlags, size, (optional)random_sub_size] , [SampleProcessor.TypeFlags, size, (optional) {} opts ] ,
... ...
] ]
''' '''
@ -46,9 +46,9 @@ class SampleGeneratorFaceTemporal(SampleGeneratorBase):
raise ValueError('No training data provided.') raise ValueError('No training data provided.')
mult_max = 1 mult_max = 1
l = samples_len - (self.temporal_image_count-1)*mult_max + 1 l = samples_len - ( (self.temporal_image_count)*mult_max - (mult_max-1) )
samples_idxs = [ *range(l) ] [generator_id::self.generators_count] samples_idxs = [ *range(l+1) ] [generator_id::self.generators_count]
if len(samples_idxs) - self.temporal_image_count < 0: if len(samples_idxs) - self.temporal_image_count < 0:
raise ValueError('Not enough samples to fit temporal line.') raise ValueError('Not enough samples to fit temporal line.')
@ -67,7 +67,7 @@ class SampleGeneratorFaceTemporal(SampleGeneratorBase):
idx = shuffle_idxs.pop() idx = shuffle_idxs.pop()
temporal_samples = [] temporal_samples = []
mult = np.random.randint(mult_max) mult = np.random.randint(mult_max)+1
for i in range( self.temporal_image_count ): for i in range( self.temporal_image_count ):
sample = samples[ idx+i*mult ] sample = samples[ idx+i*mult ]
try: try:

5
samplelib/SampleGeneratorImageTemporal.py
View file

@ -43,7 +43,8 @@ class SampleGeneratorImageTemporal(SampleGeneratorBase):
raise ValueError('No training data provided.') raise ValueError('No training data provided.')
mult_max = 4 mult_max = 4
samples_sub_len = samples_len - (self.temporal_image_count-1)*mult_max samples_sub_len = samples_len - ( (self.temporal_image_count)*mult_max - (mult_max-1) )
if samples_sub_len <= 0: if samples_sub_len <= 0:
raise ValueError('Not enough samples to fit temporal line.') raise ValueError('Not enough samples to fit temporal line.')
@ -61,7 +62,7 @@ class SampleGeneratorImageTemporal(SampleGeneratorBase):
idx = shuffle_idxs.pop() idx = shuffle_idxs.pop()
temporal_samples = [] temporal_samples = []
mult = np.random.randint(mult_max) mult = np.random.randint(mult_max)+1
for i in range( self.temporal_image_count ): for i in range( self.temporal_image_count ):
sample = samples[ idx+i*mult ] sample = samples[ idx+i*mult ]
try: try:

63
samplelib/SampleProcessor.py
View file

@ -70,10 +70,11 @@ class SampleProcessor(object):
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_FULL_NO_ALIGN = 14
FACE_TYPE_END = 20 FACE_TYPE_END = 20
MODE_BEGIN = 40 MODE_BEGIN = 40
@ -118,10 +119,10 @@ class SampleProcessor(object):
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_FULL_NO_ALIGN : FaceType.FULL_NO_ALIGN }
outputs = [] outputs = []
for opts in output_sample_types: for opts in output_sample_types:
@ -175,6 +176,20 @@ class SampleProcessor(object):
if mode_type == SPTF.NONE: if mode_type == SPTF.NONE:
raise ValueError('expected MODE_ type') raise ValueError('expected MODE_ type')
def do_transform(img, 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)
flip = img_type != SPTF.IMG_WARPED
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)
if len(mask.shape) == 2:
mask = mask[...,np.newaxis]
img = np.concatenate( (img, mask ), -1 )
return img
img = cached_images.get(img_type, None) img = cached_images.get(img_type, None)
if img is None: if img is None:
@ -199,26 +214,30 @@ class SampleProcessor(object):
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) if sample.face_type == FaceType.MARK_ONLY:
transform = (img_type == SPTF.IMG_WARPED_TRANSFORMED or img_type == SPTF.IMG_TRANSFORMED) if mask is not None:
flip = img_type != SPTF.IMG_WARPED img = np.concatenate( (img, mask), -1 )
else:
img = imagelib.warp_by_params(params, img, warp, transform, flip, True) img = do_transform (img, mask)
if mask is not None:
mask = imagelib.warp_by_params(params, mask, warp, transform, flip, False)[..., np.newaxis]
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( raise Exception ('sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.' % (sample.filename, sample.face_type, ft) )
'sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.' % (
sample.filename, sample.face_type, ft)) if sample.face_type == FaceType.MARK_ONLY:
img = cv2.warpAffine(img, LandmarksProcessor.get_transform_mat(sample.landmarks, resolution, ft), img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, sample.shape[0], ft), (sample.shape[0],sample.shape[0]), flags=cv2.INTER_CUBIC )
(resolution, resolution), flags=cv2.INTER_CUBIC)
mask = img[...,3:4] if img.shape[2] > 3 else None
img = img[...,0:3]
img = do_transform (img, mask)
img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC )
else:
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)

14
utils/Path_utils.py
View file

@ -11,7 +11,7 @@ def get_image_paths(dir_path, image_extensions=image_extensions):
for x in list(scandir(str(dir_path))): for x in list(scandir(str(dir_path))):
if any([x.name.lower().endswith(ext) for ext in image_extensions]): if any([x.name.lower().endswith(ext) for ext in image_extensions]):
result.append(x.path) result.append(x.path)
return result return sorted(result)
def get_image_unique_filestem_paths(dir_path, verbose_print_func=None): def get_image_unique_filestem_paths(dir_path, verbose_print_func=None):
result = get_image_paths(dir_path) result = get_image_paths(dir_path)
@ -26,7 +26,7 @@ def get_image_unique_filestem_paths(dir_path, verbose_print_func=None):
continue continue
result_dup.add(f_stem) result_dup.add(f_stem)
return result return sorted(result)
def get_file_paths(dir_path): def get_file_paths(dir_path):
dir_path = Path (dir_path) dir_path = Path (dir_path)
@ -34,7 +34,7 @@ def get_file_paths(dir_path):
result = [] result = []
if dir_path.exists(): if dir_path.exists():
return [ x.path for x in list(scandir(str(dir_path))) if x.is_file() ] return [ x.path for x in list(scandir(str(dir_path))) if x.is_file() ]
return result return sorted(result)
def get_all_dir_names (dir_path): def get_all_dir_names (dir_path):
dir_path = Path (dir_path) dir_path = Path (dir_path)
@ -43,7 +43,7 @@ def get_all_dir_names (dir_path):
if dir_path.exists(): if dir_path.exists():
return [ x.name for x in list(scandir(str(dir_path))) if x.is_dir() ] return [ x.name for x in list(scandir(str(dir_path))) if x.is_dir() ]
return result return sorted(result)
def get_all_dir_names_startswith (dir_path, startswith): def get_all_dir_names_startswith (dir_path, startswith):
dir_path = Path (dir_path) dir_path = Path (dir_path)
@ -54,14 +54,14 @@ def get_all_dir_names_startswith (dir_path, startswith):
for x in list(scandir(str(dir_path))): for x in list(scandir(str(dir_path))):
if x.name.lower().startswith(startswith): if x.name.lower().startswith(startswith):
result.append ( x.name[len(startswith):] ) result.append ( x.name[len(startswith):] )
return result return sorted(result)
def get_first_file_by_stem (dir_path, stem, exts=None): def get_first_file_by_stem (dir_path, stem, exts=None):
dir_path = Path (dir_path) dir_path = Path (dir_path)
stem = stem.lower() stem = stem.lower()
if dir_path.exists(): if dir_path.exists():
for x in list(scandir(str(dir_path))): for x in sorted(list(scandir(str(dir_path)))):
if not x.is_file(): if not x.is_file():
continue continue
xp = Path(x.path) xp = Path(x.path)
@ -80,4 +80,4 @@ def delete_all_files (dir_path):
paths = get_file_paths(dir_path) paths = get_file_paths(dir_path)
for p in paths: for p in paths:
p = Path(p) p = Path(p)
p.unlink() p.unlink()