diff --git a/converters/ConvertAvatar.py b/converters/ConvertAvatar.py new file mode 100644 index 0000000..2dcfeaa --- /dev/null +++ b/converters/ConvertAvatar.py @@ -0,0 +1,35 @@ +import cv2 +import numpy as np + +import imagelib +from facelib import FaceType, LandmarksProcessor +from utils.cv2_utils import * + +def process_frame_info(frame_info, inp_sh): + img_uint8 = cv2_imread (frame_info.filename) + img_uint8 = imagelib.normalize_channels (img_uint8, 3) + img = img_uint8.astype(np.float32) / 255.0 + + img_mat = LandmarksProcessor.get_transform_mat (frame_info.landmarks_list[0], inp_sh[0], face_type=FaceType.FULL_NO_ALIGN) + img = cv2.warpAffine( img, img_mat, inp_sh[0:2], flags=cv2.INTER_CUBIC ) + return img + +def ConvertFaceAvatar (cfg, prev_temporal_frame_infos, frame_info, next_temporal_frame_infos): + inp_sh = cfg.predictor_input_shape + + prev_imgs=[] + next_imgs=[] + for i in range(cfg.temporal_face_count): + prev_imgs.append( process_frame_info(prev_temporal_frame_infos[i], inp_sh) ) + next_imgs.append( process_frame_info(next_temporal_frame_infos[i], inp_sh) ) + img = process_frame_info(frame_info, inp_sh) + + prd_f = cfg.predictor_func ( prev_imgs, img, next_imgs ) + + out_img = np.clip(prd_f, 0.0, 1.0) + + if cfg.add_source_image: + out_img = np.concatenate ( [cv2.resize ( img, (prd_f.shape[1], prd_f.shape[0]) ), + out_img], axis=1 ) + + return (out_img*255).astype(np.uint8) diff --git a/converters/ConvertMasked.py b/converters/ConvertMasked.py new file mode 100644 index 0000000..6c33328 --- /dev/null +++ b/converters/ConvertMasked.py @@ -0,0 +1,392 @@ +import traceback + +import cv2 +import numpy as np + +import imagelib +from facelib import FaceType, LandmarksProcessor +from interact import interact as io +from utils.cv2_utils import * + + +def ConvertMaskedFace (cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmarks): + + #if debug: + # debugs = [img_bgr.copy()] + + img_size = img_bgr.shape[1], img_bgr.shape[0] + + img_face_mask_a = LandmarksProcessor.get_image_hull_mask (img_bgr.shape, img_face_landmarks) + + if cfg.mode == 'original': + if cfg.export_mask_alpha: + img_bgr = np.concatenate ( [img_bgr, img_face_mask_a], -1 ) + return img_bgr, img_face_mask_a + + out_img = img_bgr.copy() + out_merging_mask = None + + output_size = cfg.predictor_input_shape[0] + if cfg.super_resolution: + output_size *= 2 + + face_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, output_size, face_type=cfg.face_type) + face_output_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, output_size, face_type=cfg.face_type, scale= 1.0 + 0.01*cfg.output_face_scale ) + + dst_face_bgr = cv2.warpAffine( img_bgr , face_mat, (output_size, output_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 ) + + predictor_input_bgr = cv2.resize (dst_face_bgr, cfg.predictor_input_shape[0:2] ) + + if cfg.predictor_masked: + prd_face_bgr, prd_face_mask_a_0 = cfg.predictor_func (predictor_input_bgr) + + prd_face_bgr = np.clip (prd_face_bgr, 0, 1.0 ) + prd_face_mask_a_0 = np.clip (prd_face_mask_a_0, 0.0, 1.0) + else: + predicted = cfg.predictor_func (predictor_input_bgr) + prd_face_bgr = np.clip (predicted, 0, 1.0 ) + prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, cfg.predictor_input_shape[0:2] ) + + if cfg.super_resolution: + #if debug: + # tmp = cv2.resize (prd_face_bgr, (output_size,output_size), cv2.INTER_CUBIC) + # debugs += [ np.clip( cv2.warpAffine( tmp, face_output_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + + prd_face_bgr = cfg.dcscn_upscale_func(prd_face_bgr) + #if debug: + # debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + + if cfg.predictor_masked: + prd_face_mask_a_0 = cv2.resize (prd_face_mask_a_0, (output_size, output_size), cv2.INTER_CUBIC) + else: + prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (output_size, output_size), cv2.INTER_CUBIC) + + if cfg.mask_mode == 2: #dst + prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC) + elif cfg.mask_mode >= 3 and cfg.mask_mode <= 7: + + if cfg.mask_mode == 3 or cfg.mask_mode == 5 or cfg.mask_mode == 6: + prd_face_fanseg_bgr = cv2.resize (prd_face_bgr, (cfg.fanseg_input_size,)*2 ) + prd_face_fanseg_mask = cfg.fanseg_extract_func(FaceType.FULL, prd_face_fanseg_bgr) + FAN_prd_face_mask_a_0 = cv2.resize ( prd_face_fanseg_mask, (output_size, output_size), cv2.INTER_CUBIC) + + if cfg.mask_mode >= 4 or cfg.mask_mode <= 7: + + full_face_fanseg_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, cfg.fanseg_input_size, face_type=FaceType.FULL) + dst_face_fanseg_bgr = cv2.warpAffine(img_bgr, full_face_fanseg_mat, (cfg.fanseg_input_size,)*2, flags=cv2.INTER_CUBIC ) + dst_face_fanseg_mask = cfg.fanseg_extract_func( FaceType.FULL, dst_face_fanseg_bgr ) + + if cfg.face_type == FaceType.FULL: + FAN_dst_face_mask_a_0 = cv2.resize (dst_face_fanseg_mask, (output_size,output_size), cv2.INTER_CUBIC) + elif cfg.face_type == FaceType.HALF: + half_face_fanseg_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, cfg.fanseg_input_size, face_type=FaceType.HALF) + + fanseg_rect_corner_pts = np.array ( [ [0,0], [cfg.fanseg_input_size-1,0], [0,cfg.fanseg_input_size-1] ], dtype=np.float32 ) + a = LandmarksProcessor.transform_points (fanseg_rect_corner_pts, half_face_fanseg_mat, invert=True ) + b = LandmarksProcessor.transform_points (a, full_face_fanseg_mat ) + m = cv2.getAffineTransform(b, fanseg_rect_corner_pts) + FAN_dst_face_mask_a_0 = cv2.warpAffine(dst_face_fanseg_mask, m, (cfg.fanseg_input_size,)*2, flags=cv2.INTER_CUBIC ) + FAN_dst_face_mask_a_0 = cv2.resize (FAN_dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC) + else: + raise ValueError ("cfg.face_type unsupported") + + if cfg.mask_mode == 3: #FAN-prd + prd_face_mask_a_0 = FAN_prd_face_mask_a_0 + elif cfg.mask_mode == 4: #FAN-dst + prd_face_mask_a_0 = FAN_dst_face_mask_a_0 + elif cfg.mask_mode == 5: + prd_face_mask_a_0 = FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0 + elif cfg.mask_mode == 6: + prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0 + elif cfg.mask_mode == 7: + prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_dst_face_mask_a_0 + + prd_face_mask_a_0[ prd_face_mask_a_0 < 0.001 ] = 0.0 + + prd_face_mask_a = prd_face_mask_a_0[...,np.newaxis] + prd_face_mask_aaa = np.repeat (prd_face_mask_a, (3,), axis=-1) + + img_face_mask_aaa = cv2.warpAffine( prd_face_mask_aaa, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC ) + img_face_mask_aaa = np.clip (img_face_mask_aaa, 0.0, 1.0) + img_face_mask_aaa [ img_face_mask_aaa <= 0.1 ] = 0.0 #get rid of noise + + #if debug: + # debugs += [img_face_mask_aaa.copy()] + + if 'raw' in cfg.mode: + face_corner_pts = np.array ([ [0,0], [output_size-1,0], [output_size-1,output_size-1], [0,output_size-1] ], dtype=np.float32) + square_mask = np.zeros(img_bgr.shape, dtype=np.float32) + cv2.fillConvexPoly(square_mask, \ + LandmarksProcessor.transform_points (face_corner_pts, face_output_mat, invert=True ).astype(np.int), \ + (1,1,1) ) + + if cfg.mode == 'raw-rgb': + out_merging_mask = square_mask + + if cfg.mode == 'raw-rgb' or cfg.mode == 'raw-rgb-mask': + out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ) + + if cfg.mode == 'raw-rgb-mask': + out_img = np.concatenate ( [out_img, np.expand_dims (img_face_mask_aaa[:,:,0],-1)], -1 ) + out_merging_mask = square_mask + + elif cfg.mode == 'raw-mask-only': + out_img = img_face_mask_aaa + out_merging_mask = img_face_mask_aaa + elif cfg.mode == 'raw-predicted-only': + out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ) + out_merging_mask = square_mask + + out_img = np.clip (out_img, 0.0, 1.0 ) + else: + #averaging [lenx, leny, maskx, masky] by grayscale gradients of upscaled mask + ar = [] + for i in range(1, 10): + maxregion = np.argwhere( img_face_mask_aaa > i / 10.0 ) + if maxregion.size != 0: + miny,minx = maxregion.min(axis=0)[:2] + maxy,maxx = maxregion.max(axis=0)[:2] + lenx = maxx - minx + leny = maxy - miny + if min(lenx,leny) >= 4: + ar += [ [ lenx, leny] ] + + if len(ar) > 0: + lenx, leny = np.mean ( ar, axis=0 ) + lowest_len = min (lenx, leny) + #if debug: + # io.log_info ("lenx/leny:(%d/%d) " % (lenx, leny ) ) + # io.log_info ("lowest_len = %f" % (lowest_len) ) + + if cfg.erode_mask_modifier != 0: + ero = int( lowest_len * ( 0.126 - lowest_len * 0.00004551365 ) * 0.01*cfg.erode_mask_modifier ) + #if debug: + # io.log_info ("erode_size = %d" % (ero) ) + if ero > 0: + img_face_mask_aaa = cv2.erode(img_face_mask_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero,ero)), iterations = 1 ) + elif ero < 0: + img_face_mask_aaa = cv2.dilate(img_face_mask_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(-ero,-ero)), iterations = 1 ) + + if cfg.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_border_size = int ( prd_hborder_rect_mask_a.shape[1] * cfg.clip_hborder_mask_per ) + 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 = np.expand_dims(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, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC ) + img_prd_hborder_rect_mask_a = np.expand_dims (img_prd_hborder_rect_mask_a, -1) + img_face_mask_aaa *= img_prd_hborder_rect_mask_a + img_face_mask_aaa = np.clip( img_face_mask_aaa, 0, 1.0 ) + + #if debug: + # debugs += [img_face_mask_aaa.copy()] + + if cfg.blur_mask_modifier > 0: + blur = int( lowest_len * 0.10 * 0.01*cfg.blur_mask_modifier ) + #if debug: + # io.log_info ("blur_size = %d" % (blur) ) + if blur > 0: + img_face_mask_aaa = cv2.blur(img_face_mask_aaa, (blur, blur) ) + + img_face_mask_aaa = np.clip( img_face_mask_aaa, 0, 1.0 ) + + + #if debug: + # debugs += [img_face_mask_aaa.copy()] + + if 'seamless' not in cfg.mode and cfg.color_transfer_mode != 0: + if cfg.color_transfer_mode == 1: + #if debug: + # debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + + prd_face_bgr = imagelib.reinhard_color_transfer ( np.clip( (prd_face_bgr*255).astype(np.uint8), 0, 255), + np.clip( (dst_face_bgr*255).astype(np.uint8), 0, 255), + source_mask=prd_face_mask_a, target_mask=prd_face_mask_a) + prd_face_bgr = np.clip( prd_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0) + + #if debug: + # debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + + + elif cfg.color_transfer_mode == 2: + #if debug: + # debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + + prd_face_bgr = imagelib.linear_color_transfer (prd_face_bgr, dst_face_bgr) + prd_face_bgr = np.clip( prd_face_bgr, 0.0, 1.0) + + #if debug: + # debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + + if cfg.mode == 'hist-match-bw': + prd_face_bgr = cv2.cvtColor(prd_face_bgr, cv2.COLOR_BGR2GRAY) + prd_face_bgr = np.repeat( np.expand_dims (prd_face_bgr, -1), (3,), -1 ) + + if cfg.mode == 'hist-match' or cfg.mode == 'hist-match-bw': + #if debug: + # debugs += [ cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ) ] + + hist_mask_a = np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32) + + if cfg.masked_hist_match: + hist_mask_a *= prd_face_mask_a + + white = (1.0-hist_mask_a)* np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32) + + hist_match_1 = prd_face_bgr*hist_mask_a + white + hist_match_1[ hist_match_1 > 1.0 ] = 1.0 + + hist_match_2 = dst_face_bgr*hist_mask_a + white + hist_match_2[ hist_match_1 > 1.0 ] = 1.0 + + prd_face_bgr = imagelib.color_hist_match(hist_match_1, hist_match_2, cfg.hist_match_threshold ) + + #if cfg.masked_hist_match: + # prd_face_bgr -= white + + if cfg.mode == 'hist-match-bw': + prd_face_bgr = prd_face_bgr.astype(dtype=np.float32) + + out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ) + out_img = np.clip(out_img, 0.0, 1.0) + + #if debug: + # debugs += [out_img.copy()] + + if cfg.mode == 'overlay': + pass + + if 'seamless' in cfg.mode: + #mask used for cv2.seamlessClone + img_face_seamless_mask_a = None + img_face_mask_a = img_face_mask_aaa[...,0:1] + for i in range(1,10): + a = img_face_mask_a > i / 10.0 + if len(np.argwhere(a)) == 0: + continue + img_face_seamless_mask_a = img_face_mask_aaa[...,0:1].copy() + img_face_seamless_mask_a[a] = 1.0 + img_face_seamless_mask_a[img_face_seamless_mask_a <= i / 10.0] = 0.0 + break + + try: + #calc same bounding rect and center point as in cv2.seamlessClone to prevent jittering (not flickering) + l,t,w,h = cv2.boundingRect( (img_face_seamless_mask_a*255).astype(np.uint8) ) + s_maskx, s_masky = int(l+w/2), int(t+h/2) + + out_img = cv2.seamlessClone( (out_img*255).astype(np.uint8), img_bgr_uint8, (img_face_seamless_mask_a*255).astype(np.uint8), (s_maskx,s_masky) , cv2.NORMAL_CLONE ) + out_img = out_img.astype(dtype=np.float32) / 255.0 + except Exception as e: + #seamlessClone may fail in some cases + e_str = traceback.format_exc() + + if 'MemoryError' in e_str: + raise Exception("Seamless fail: " + e_str) #reraise MemoryError in order to reprocess this data by other processes + else: + print ("Seamless fail: " + e_str) + + #if debug: + # debugs += [out_img.copy()] + + out_img = img_bgr*(1-img_face_mask_aaa) + (out_img*img_face_mask_aaa) + + if 'seamless' in cfg.mode and cfg.color_transfer_mode != 0: + out_face_bgr = cv2.warpAffine( out_img, face_mat, (output_size, output_size) ) + + if cfg.color_transfer_mode == 1: + #if debug: + # debugs += [ np.clip( cv2.warpAffine( out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + face_mask_aaa = cv2.warpAffine( img_face_mask_aaa, face_mat, (output_size, output_size) ) + + new_out_face_bgr = imagelib.reinhard_color_transfer ( np.clip( (out_face_bgr*255).astype(np.uint8), 0, 255), + np.clip( (dst_face_bgr*255).astype(np.uint8), 0, 255), + source_mask=face_mask_aaa, target_mask=face_mask_aaa) + new_out_face_bgr = np.clip( new_out_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0) + + #if debug: + # debugs += [ np.clip( cv2.warpAffine( new_out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + + + elif cfg.color_transfer_mode == 2: + #if debug: + # debugs += [ np.clip( cv2.warpAffine( out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + + new_out_face_bgr = imagelib.linear_color_transfer (out_face_bgr, dst_face_bgr) + new_out_face_bgr = np.clip( new_out_face_bgr, 0.0, 1.0) + + #if debug: + # debugs += [ np.clip( cv2.warpAffine( new_out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] + + new_out = cv2.warpAffine( new_out_face_bgr, face_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ) + out_img = np.clip( img_bgr*(1-img_face_mask_aaa) + (new_out*img_face_mask_aaa) , 0, 1.0 ) + + if cfg.mode == 'seamless-hist-match': + out_face_bgr = cv2.warpAffine( out_img, face_mat, (output_size, output_size) ) + new_out_face_bgr = imagelib.color_hist_match(out_face_bgr, dst_face_bgr, cfg.hist_match_threshold) + new_out = cv2.warpAffine( new_out_face_bgr, face_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ) + out_img = np.clip( img_bgr*(1-img_face_mask_aaa) + (new_out*img_face_mask_aaa) , 0, 1.0 ) + + + cfg_mp = cfg.motion_blur_power / 100.0 + if cfg_mp != 0: + k_size = int(frame_info.motion_power*cfg_mp) + if k_size >= 1: + k_size = np.clip (k_size+1, 2, 50) + if cfg.super_resolution: + k_size *= 2 + out_face_bgr = cv2.warpAffine( out_img, face_mat, (output_size, output_size) ) + new_out_face_bgr = imagelib.LinearMotionBlur (out_face_bgr, k_size , frame_info.motion_deg) + new_out = cv2.warpAffine( new_out_face_bgr, face_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT ) + out_img = np.clip( img_bgr*(1-img_face_mask_aaa) + (new_out*img_face_mask_aaa) , 0, 1.0 ) + + if cfg.color_degrade_power != 0: + #if debug: + # debugs += [out_img.copy()] + out_img_reduced = imagelib.reduce_colors(out_img, 256) + if cfg.color_degrade_power == 100: + out_img = out_img_reduced + else: + alpha = cfg.color_degrade_power / 100.0 + out_img = (out_img*(1.0-alpha) + out_img_reduced*alpha) + + if cfg.export_mask_alpha: + out_img = np.concatenate ( [out_img, img_face_mask_aaa[:,:,0:1]], -1 ) + out_merging_mask = img_face_mask_aaa + + + #if debug: + # debugs += [out_img.copy()] + + return out_img, out_merging_mask + + +def ConvertMasked (cfg, frame_info): + img_bgr_uint8 = cv2_imread(frame_info.filename) + img_bgr_uint8 = imagelib.normalize_channels (img_bgr_uint8, 3) + img_bgr = img_bgr_uint8.astype(np.float32) / 255.0 + + outs = [] + for face_num, img_landmarks in enumerate( frame_info.landmarks_list ): + out_img, out_img_merging_mask = ConvertMaskedFace (cfg, frame_info, img_bgr_uint8, img_bgr, img_landmarks) + outs += [ (out_img, out_img_merging_mask) ] + + #Combining multiple face outputs + final_img = None + for img, merging_mask in outs: + h,w,c = img.shape + + if final_img is None: + final_img = img + else: + merging_mask = merging_mask[...,0:1] + if c == 3: + final_img = final_img*(1-merging_mask) + img*merging_mask + elif c == 4: + final_img_bgr = final_img[...,0:3]*(1-merging_mask) + img[...,0:3]*merging_mask + final_img_mask = np.clip ( final_img[...,3:4] + img[...,3:4], 0, 1 ) + final_img = np.concatenate ( [final_img_bgr, final_img_mask], -1 ) + + return (final_img*255).astype(np.uint8) \ No newline at end of file diff --git a/converters/Converter.py b/converters/Converter.py deleted file mode 100644 index b4e4213..0000000 --- a/converters/Converter.py +++ /dev/null @@ -1,50 +0,0 @@ -import copy -''' -You can implement your own Converter, check example ConverterMasked.py -''' - -class Converter(object): - TYPE_FACE = 0 #calls convert_face - TYPE_FACE_AVATAR = 1 #calls convert_face with avatar_operator_face - TYPE_IMAGE = 2 #calls convert_image without landmarks - TYPE_IMAGE_WITH_LANDMARKS = 3 #calls convert_image with landmarks - - #overridable - def __init__(self, predictor_func, type): - self.predictor_func = predictor_func - self.type = type - - #overridable - def on_cli_initialize(self): - #cli initialization - pass - - #overridable - def on_host_tick(self): - pass - - #overridable - def cli_convert_face (self, img_bgr, img_face_landmarks, debug, avaperator_face_bgr=None, **kwargs): - #return float32 image - #if debug , return tuple ( images of any size and channels, ...) - return image - - #overridable - def cli_convert_image (self, img_bgr, img_landmarks, debug): - #img_landmarks not None, if input image is png with embedded data - #return float32 image - #if debug , return tuple ( images of any size and channels, ...) - return image - - #overridable - def dummy_predict(self): - #do dummy predict here - pass - - def copy(self): - return copy.copy(self) - - def copy_and_set_predictor(self, predictor_func): - result = self.copy() - result.predictor_func = predictor_func - return result diff --git a/converters/ConverterAvatar.py b/converters/ConverterAvatar.py deleted file mode 100644 index ca26c63..0000000 --- a/converters/ConverterAvatar.py +++ /dev/null @@ -1,61 +0,0 @@ -import time - -import cv2 -import numpy as np - -from facelib import FaceType, LandmarksProcessor -from joblib import SubprocessFunctionCaller -from utils.pickle_utils import AntiPickler - -from .Converter import Converter - -class ConverterAvatar(Converter): - - #override - def __init__(self, predictor_func, - predictor_input_size=0): - - super().__init__(predictor_func, Converter.TYPE_FACE_AVATAR) - - self.predictor_input_size = predictor_input_size - - #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 ), - 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) - - predictor_func_host, predictor_func = SubprocessFunctionCaller.make_pair(predictor_func) - self.predictor_func_host = AntiPickler(predictor_func_host) - self.predictor_func = predictor_func - - #overridable - def on_host_tick(self): - self.predictor_func_host.obj.process_messages() - - #override - def cli_convert_face (self, f0, f0_lmrk, f1, f1_lmrk, f2, f2_lmrk, debug, **kwargs): - if debug: - debugs = [] - - inp_size = self.predictor_input_size - - f0_mat = LandmarksProcessor.get_transform_mat (f0_lmrk, inp_size, face_type=FaceType.FULL_NO_ALIGN) - 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) - - inp_f0 = cv2.warpAffine( f0, f0_mat, (inp_size, inp_size), flags=cv2.INTER_CUBIC ) - 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 ) - - prd_f = self.predictor_func ( inp_f0, inp_f1, inp_f2 ) - - out_img = np.clip(prd_f, 0.0, 1.0) - - out_img = np.concatenate ( [cv2.resize ( inp_f1, (prd_f.shape[1], prd_f.shape[0]) ), - out_img], axis=1 ) - - if debug: - debugs += [out_img.copy()] - - return debugs if debug else out_img diff --git a/converters/ConverterConfig.py b/converters/ConverterConfig.py new file mode 100644 index 0000000..8c77902 --- /dev/null +++ b/converters/ConverterConfig.py @@ -0,0 +1,317 @@ +import numpy as np +import copy + +from facelib import FaceType +from interact import interact as io + + +class ConverterConfig(object): + TYPE_NONE = 0 + TYPE_MASKED = 1 + TYPE_FACE_AVATAR = 2 + #### + + TYPE_IMAGE = 3 + TYPE_IMAGE_WITH_LANDMARKS = 4 + + def __init__(self, type=0, predictor_func=None, + predictor_input_shape=None): + self.type = type + self.predictor_func = predictor_func + self.predictor_input_shape = predictor_input_shape + + self.dcscn_upscale_func = None + self.fanseg_input_size = None + self.fanseg_extract_func = None + + def copy(self): + return copy.copy(self) + + #overridable + def ask_settings(self): + pass + + #overridable + def __eq__(self, other): + #check equality of changeable params + + if isinstance(other, ConverterConfig): + return True + + return False + + #overridable + def __str__(self): + return "ConverterConfig: ." + +class ConverterConfigMasked(ConverterConfig): + + def __init__(self, predictor_func=None, + predictor_input_shape=None, + predictor_masked=True, + face_type=FaceType.FULL, + default_mode = 4, + base_erode_mask_modifier = 0, + base_blur_mask_modifier = 0, + default_erode_mask_modifier = 0, + default_blur_mask_modifier = 0, + clip_hborder_mask_per = 0, + ): + + super().__init__(type=ConverterConfig.TYPE_MASKED, + predictor_func=predictor_func, + predictor_input_shape=predictor_input_shape, + ) + if len(predictor_input_shape) != 3: + raise ValueError("ConverterConfigMasked: predictor_input_shape must be rank 3.") + + if predictor_input_shape[0] != predictor_input_shape[1]: + raise ValueError("ConverterConfigMasked: predictor_input_shape must be a square.") + + self.predictor_masked = predictor_masked + self.face_type = face_type + if self.face_type not in [FaceType.FULL, FaceType.HALF]: + raise ValueError("ConverterConfigMasked supports only full or half face masks.") + + self.default_mode = default_mode + self.base_erode_mask_modifier = base_erode_mask_modifier + self.base_blur_mask_modifier = base_blur_mask_modifier + self.default_erode_mask_modifier = default_erode_mask_modifier + self.default_blur_mask_modifier = default_blur_mask_modifier + self.clip_hborder_mask_per = clip_hborder_mask_per + + #default changeable params + + + self.mode = 'overlay' + self.masked_hist_match = True + self.hist_match_threshold = 238 + self.mask_mode = 1 + self.erode_mask_modifier = 0 + self.blur_mask_modifier = 0 + self.motion_blur_power = 0 + self.output_face_scale = 0 + self.color_transfer_mode = 0 + self.super_resolution = False + self.color_degrade_power = 0 + self.export_mask_alpha = False + + self.mode_dict = {0:'original', + 1:'overlay', + 2:'hist-match', + 3:'hist-match-bw', + 4:'seamless', + 5:'seamless-hist-match', + 6:'raw-rgb', + 7:'raw-rgb-mask', + 8:'raw-mask-only', + 9:'raw-predicted-only'} + + self.full_face_mask_mode_dict = {1:'learned', + 2:'dst', + 3:'FAN-prd', + 4:'FAN-dst', + 5:'FAN-prd*FAN-dst', + 6:'learned*FAN-prd*FAN-dst'} + + self.half_face_mask_mode_dict = {1:'learned', + 2:'dst', + 4:'FAN-dst', + 7:'learned*FAN-dst'} + + self.ctm_dict = { 0: "None", 1:"rct", 2:"lct" } + self.ctm_str_dict = {None:None, "rct":1, "lct": 2 } + + def copy(self): + return copy.copy(self) + + def set_mode (self, mode): + self.mode = self.mode_dict.get (mode, self.mode_dict[self.default_mode] ) + + def toggle_masked_hist_match(self): + if self.mode == 'hist-match' or self.mode == 'hist-match-bw': + self.masked_hist_match = not self.masked_hist_match + + def add_hist_match_threshold(self, diff): + if self.mode == 'hist-match' or self.mode == 'hist-match-bw' or self.mode == 'seamless-hist-match': + self.hist_match_threshold = np.clip ( self.hist_match_threshold+diff , 0, 255) + + def toggle_mask_mode(self): + if self.face_type == FaceType.FULL: + a = list( self.full_face_mask_mode_dict.keys() ) + else: + a = list( self.half_face_mask_mode_dict.keys() ) + self.mask_mode = a[ (a.index(self.mask_mode)+1) % len(a) ] + + def add_erode_mask_modifier(self, diff): + self.erode_mask_modifier = np.clip ( self.erode_mask_modifier+diff , -200, 200) + + def add_blur_mask_modifier(self, diff): + self.blur_mask_modifier = np.clip ( self.blur_mask_modifier+diff , -200, 200) + + def add_motion_blur_power(self, diff): + self.motion_blur_power = np.clip ( self.motion_blur_power+diff, 0, 100) + + def add_output_face_scale(self, diff): + self.output_face_scale = np.clip ( self.output_face_scale+diff , -50, 50) + + def toggle_color_transfer_mode(self): + self.color_transfer_mode = (self.color_transfer_mode+1) % 3 + + def toggle_super_resolution(self): + self.super_resolution = not self.super_resolution + + def add_color_degrade_power(self, diff): + self.color_degrade_power = np.clip ( self.color_degrade_power+diff , 0, 100) + + def toggle_export_mask_alpha(self): + self.export_mask_alpha = not self.export_mask_alpha + + def ask_settings(self): + + s = """Choose mode: \n""" + for key in self.mode_dict.keys(): + s += f"""({key}) {self.mode_dict[key]}\n""" + s += f"""Default: {self.default_mode} : """ + + mode = io.input_int (s, self.default_mode) + + self.mode = self.mode_dict.get (mode, self.mode_dict[self.default_mode] ) + + if 'raw' not in self.mode: + if self.mode == 'hist-match' or self.mode == 'hist-match-bw': + self.masked_hist_match = io.input_bool("Masked hist match? (y/n skip:y) : ", True) + + if self.mode == 'hist-match' or self.mode == 'hist-match-bw' or self.mode == 'seamless-hist-match': + self.hist_match_threshold = np.clip ( io.input_int("Hist match threshold [0..255] (skip:255) : ", 255), 0, 255) + + if self.face_type == FaceType.FULL: + s = """Choose mask mode: \n""" + for key in self.full_face_mask_mode_dict.keys(): + s += f"""({key}) {self.full_face_mask_mode_dict[key]}\n""" + s += f"""?:help Default: 1 : """ + + self.mask_mode = io.input_int (s, 1, valid_list=self.full_face_mask_mode_dict.keys(), help_message="If you learned the mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images. 'FAN-prd' - using super smooth mask by pretrained FAN-model from predicted face. 'FAN-dst' - using super smooth mask by pretrained FAN-model from dst face. 'FAN-prd*FAN-dst' or 'learned*FAN-prd*FAN-dst' - using multiplied masks.") + else: + s = """Choose mask mode: \n""" + for key in self.half_face_mask_mode_dict.keys(): + s += f"""({key}) {self.half_face_mask_mode_dict[key]}\n""" + s += f"""?:help , Default: 1 : """ + self.mask_mode = io.input_int (s, 1, valid_list=self.half_face_mask_mode_dict.keys(), help_message="If you learned the mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images.") + + if 'raw' not in self.mode: + self.erode_mask_modifier = self.base_erode_mask_modifier + np.clip ( io.input_int ("Choose erode mask modifier [-200..200] (skip:%d) : " % (self.default_erode_mask_modifier), self.default_erode_mask_modifier), -200, 200) + self.blur_mask_modifier = self.base_blur_mask_modifier + np.clip ( io.input_int ("Choose blur mask modifier [-200..200] (skip:%d) : " % (self.default_blur_mask_modifier), self.default_blur_mask_modifier), -200, 200) + self.motion_blur_power = np.clip ( io.input_int ("Choose motion blur power [0..100] (skip:%d) : " % (0), 0), 0, 100) + + self.output_face_scale = np.clip (io.input_int ("Choose output face scale modifier [-50..50] (skip:0) : ", 0), -50, 50) + + if 'raw' not in self.mode: + self.color_transfer_mode = io.input_str ("Apply color transfer to predicted face? Choose mode ( rct/lct skip:None ) : ", None, ['rct','lct']) + self.color_transfer_mode = self.ctm_str_dict[self.color_transfer_mode] + + self.super_resolution = io.input_bool("Apply super resolution? (y/n ?:help skip:n) : ", False, help_message="Enhance details by applying DCSCN network.") + + if 'raw' not in self.mode: + self.color_degrade_power = np.clip ( io.input_int ("Degrade color power of final image [0..100] (skip:0) : ", 0), 0, 100) + self.export_mask_alpha = io.input_bool("Export png with alpha channel of the mask? (y/n skip:n) : ", False) + + io.log_info ("") + + def __eq__(self, other): + #check equality of changeable params + + if isinstance(other, ConverterConfigMasked): + return self.mode == other.mode and \ + self.masked_hist_match == other.masked_hist_match and \ + self.hist_match_threshold == other.hist_match_threshold and \ + self.mask_mode == other.mask_mode and \ + self.erode_mask_modifier == other.erode_mask_modifier and \ + self.blur_mask_modifier == other.blur_mask_modifier and \ + self.motion_blur_power == other.motion_blur_power and \ + self.output_face_scale == other.output_face_scale and \ + self.color_transfer_mode == other.color_transfer_mode and \ + self.super_resolution == other.super_resolution and \ + self.color_degrade_power == other.color_degrade_power and \ + self.export_mask_alpha == other.export_mask_alpha + + return False + + def __str__(self): + r = ( + """ConverterConfig:\n""" + f"""Mode: {self.mode}\n""" + ) + + if self.mode == 'hist-match' or self.mode == 'hist-match-bw': + r += f"""masked_hist_match: {self.masked_hist_match}\n""" + + if self.mode == 'hist-match' or self.mode == 'hist-match-bw' or self.mode == 'seamless-hist-match': + r += f"""hist_match_threshold: {self.hist_match_threshold}\n""" + + if self.face_type == FaceType.FULL: + r += f"""mask_mode: { self.full_face_mask_mode_dict[self.mask_mode] }\n""" + else: + r += f"""mask_mode: { self.half_face_mask_mode_dict[self.mask_mode] }\n""" + + if 'raw' not in self.mode: + r += (f"""erode_mask_modifier: {self.erode_mask_modifier}\n""" + f"""blur_mask_modifier: {self.blur_mask_modifier}\n""" + f"""motion_blur_power: {self.motion_blur_power}\n""") + + r += f"""output_face_scale: {self.output_face_scale}\n""" + + if 'raw' not in self.mode: + r += f"""color_transfer_mode: { self.ctm_dict[self.color_transfer_mode]}\n""" + + r += f"""super_resolution: {self.super_resolution}\n""" + + if 'raw' not in self.mode: + r += (f"""color_degrade_power: {self.color_degrade_power}\n""" + f"""export_mask_alpha: {self.export_mask_alpha}\n""") + + r += "================" + + return r + + +class ConverterConfigFaceAvatar(ConverterConfig): + + def __init__(self, predictor_func=None, + predictor_input_shape=None, + temporal_face_count=0 + ): + super().__init__(type=ConverterConfig.TYPE_FACE_AVATAR, + predictor_func=predictor_func, + predictor_input_shape=predictor_input_shape + ) + self.temporal_face_count = temporal_face_count + + #changeable params + self.add_source_image = False + + def copy(self): + return copy.copy(self) + + #override + def ask_settings(self): + self.add_source_image = io.input_bool("Add source image? (y/n ?:help skip:n) : ", False, help_message="Add source image for comparison.") + + def toggle_add_source_image(self): + self.add_source_image = not self.add_source_image + + #override + def __eq__(self, other): + #check equality of changeable params + + if isinstance(other, ConverterConfigFaceAvatar): + return self.add_source_image == other.add_source_image + + return False + + #override + def __str__(self): + return ("ConverterConfig: \n" + f"add_source_image : {self.add_source_image}\n" + "================" + ) \ No newline at end of file diff --git a/converters/ConverterImage.py b/converters/ConverterImage.py deleted file mode 100644 index 58b1faa..0000000 --- a/converters/ConverterImage.py +++ /dev/null @@ -1,50 +0,0 @@ -import time - -import cv2 -import numpy as np - -from facelib import FaceType, LandmarksProcessor -from joblib import SubprocessFunctionCaller -from utils.pickle_utils import AntiPickler - -from .Converter import Converter - -class ConverterImage(Converter): - - #override - def __init__(self, predictor_func, - predictor_input_size=0): - - super().__init__(predictor_func, Converter.TYPE_IMAGE) - - self.predictor_input_size = predictor_input_size - - #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 ) ) - time.sleep(2) - - predictor_func_host, predictor_func = SubprocessFunctionCaller.make_pair(predictor_func) - self.predictor_func_host = AntiPickler(predictor_func_host) - self.predictor_func = predictor_func - - #overridable - def on_host_tick(self): - self.predictor_func_host.obj.process_messages() - - #override - def cli_convert_image (self, img_bgr, img_landmarks, debug): - img_size = img_bgr.shape[1], img_bgr.shape[0] - - predictor_input_bgr = cv2.resize ( img_bgr, (self.predictor_input_size, self.predictor_input_size), cv2.INTER_LANCZOS4 ) - - if debug: - debugs = [predictor_input_bgr] - - output = self.predictor_func ( predictor_input_bgr ) - - if debug: - return (predictor_input_bgr,output,) - if debug: - debugs += [out_img.copy()] - - return debugs if debug else output diff --git a/converters/ConverterMasked.py b/converters/ConverterMasked.py deleted file mode 100644 index d080e49..0000000 --- a/converters/ConverterMasked.py +++ /dev/null @@ -1,436 +0,0 @@ -import time -import traceback - -import cv2 -import numpy as np - -import imagelib -from facelib import FaceType, FANSegmentator, LandmarksProcessor -from interact import interact as io -from joblib import SubprocessFunctionCaller -from utils.pickle_utils import AntiPickler - -from .Converter import Converter - - -''' -default_mode = {1:'overlay', - 2:'hist-match', - 3:'hist-match-bw', - 4:'seamless', - 5:'seamless-hist-match', - 6:'raw'} -''' -class ConverterMasked(Converter): - - #override - def __init__(self, predictor_func, - predictor_input_size=0, - predictor_masked=True, - face_type=FaceType.FULL, - default_mode = 4, - base_erode_mask_modifier = 0, - base_blur_mask_modifier = 0, - default_erode_mask_modifier = 0, - default_blur_mask_modifier = 0, - clip_hborder_mask_per = 0, - force_mask_mode=-1): - - super().__init__(predictor_func, Converter.TYPE_FACE) - - #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 ) ) - time.sleep(2) - - predictor_func_host, predictor_func = SubprocessFunctionCaller.make_pair(predictor_func) - self.predictor_func_host = AntiPickler(predictor_func_host) - self.predictor_func = predictor_func - - self.predictor_masked = predictor_masked - self.predictor_input_size = predictor_input_size - self.face_type = face_type - self.clip_hborder_mask_per = clip_hborder_mask_per - - mode = io.input_int ("Choose mode: (1) overlay, (2) hist match, (3) hist match bw, (4) seamless, (5) raw. Default - %d : " % (default_mode) , default_mode) - - mode_dict = {1:'overlay', - 2:'hist-match', - 3:'hist-match-bw', - 4:'seamless', - 5:'raw'} - - self.mode = mode_dict.get (mode, mode_dict[default_mode] ) - - if self.mode == 'raw': - mode = io.input_int ("Choose raw mode: (1) rgb, (2) rgb+mask (default), (3) mask only, (4) predicted only : ", 2) - self.raw_mode = {1:'rgb', - 2:'rgb-mask', - 3:'mask-only', - 4:'predicted-only'}.get (mode, 'rgb-mask') - - if self.mode != 'raw': - - if self.mode == 'seamless': - if io.input_bool("Seamless hist match? (y/n skip:n) : ", False): - self.mode = 'seamless-hist-match' - - if self.mode == 'hist-match' or self.mode == 'hist-match-bw': - self.masked_hist_match = io.input_bool("Masked hist match? (y/n skip:y) : ", True) - - if self.mode == 'hist-match' or self.mode == 'hist-match-bw' or self.mode == 'seamless-hist-match': - self.hist_match_threshold = np.clip ( io.input_int("Hist match threshold [0..255] (skip:255) : ", 255), 0, 255) - - if force_mask_mode != -1: - self.mask_mode = force_mask_mode - else: - if face_type == FaceType.FULL: - self.mask_mode = np.clip ( io.input_int ("Mask mode: (1) learned, (2) dst, (3) FAN-prd, (4) FAN-dst , (5) FAN-prd*FAN-dst (6) learned*FAN-prd*FAN-dst (?) help. Default - %d : " % (1) , 1, help_message="If you learned mask, then option 1 should be choosed. 'dst' mask is raw shaky mask from dst aligned images. 'FAN-prd' - using super smooth mask by pretrained FAN-model from predicted face. 'FAN-dst' - using super smooth mask by pretrained FAN-model from dst face. 'FAN-prd*FAN-dst' or 'learned*FAN-prd*FAN-dst' - using multiplied masks."), 1, 6 ) - else: - self.mask_mode = np.clip ( io.input_int ("Mask mode: (1) learned, (2) dst . Default - %d : " % (1) , 1), 1, 2 ) - - if self.mask_mode >= 3 and self.mask_mode <= 6: - self.fan_seg = None - - if self.mode != 'raw': - self.erode_mask_modifier = base_erode_mask_modifier + np.clip ( io.input_int ("Choose erode mask modifier [-200..200] (skip:%d) : " % (default_erode_mask_modifier), default_erode_mask_modifier), -200, 200) - self.blur_mask_modifier = base_blur_mask_modifier + np.clip ( io.input_int ("Choose blur mask modifier [-200..200] (skip:%d) : " % (default_blur_mask_modifier), default_blur_mask_modifier), -200, 200) - - self.output_face_scale = np.clip ( 1.0 + io.input_int ("Choose output face scale modifier [-50..50] (skip:0) : ", 0)*0.01, 0.5, 1.5) - - if self.mode != 'raw': - self.color_transfer_mode = io.input_str ("Apply color transfer to predicted face? Choose mode ( rct/lct skip:None ) : ", None, ['rct','lct']) - - self.super_resolution = io.input_bool("Apply super resolution? (y/n ?:help skip:n) : ", False, help_message="Enhance details by applying DCSCN network.") - - if self.mode != 'raw': - self.final_image_color_degrade_power = np.clip ( io.input_int ("Degrade color power of final image [0..100] (skip:0) : ", 0), 0, 100) - self.alpha = io.input_bool("Export png with alpha channel? (y/n skip:n) : ", False) - - io.log_info ("") - - if self.super_resolution: - host_proc, dc_upscale = SubprocessFunctionCaller.make_pair( imagelib.DCSCN().upscale ) - self.dc_host = AntiPickler(host_proc) - self.dc_upscale = dc_upscale - else: - self.dc_host = None - - #overridable - def on_host_tick(self): - self.predictor_func_host.obj.process_messages() - - if self.dc_host is not None: - self.dc_host.obj.process_messages() - - #overridable - def on_cli_initialize(self): - if (self.mask_mode >= 3 and self.mask_mode <= 6) and self.fan_seg == None: - self.fan_seg = FANSegmentator(256, FaceType.toString( self.face_type ) ) - - #override - def cli_convert_face (self, img_bgr, img_face_landmarks, debug, **kwargs): - if debug: - debugs = [img_bgr.copy()] - - img_size = img_bgr.shape[1], img_bgr.shape[0] - - img_face_mask_a = LandmarksProcessor.get_image_hull_mask (img_bgr.shape, img_face_landmarks) - - output_size = self.predictor_input_size - if self.super_resolution: - output_size *= 2 - - face_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, output_size, face_type=self.face_type) - face_output_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, output_size, face_type=self.face_type, scale=self.output_face_scale) - - dst_face_bgr = cv2.warpAffine( img_bgr , face_mat, (output_size, output_size), flags=cv2.INTER_LANCZOS4 ) - dst_face_mask_a_0 = cv2.warpAffine( img_face_mask_a, face_mat, (output_size, output_size), flags=cv2.INTER_LANCZOS4 ) - - predictor_input_bgr = cv2.resize (dst_face_bgr, (self.predictor_input_size,self.predictor_input_size)) - - if self.predictor_masked: - prd_face_bgr, prd_face_mask_a_0 = self.predictor_func (predictor_input_bgr) - - prd_face_bgr = np.clip (prd_face_bgr, 0, 1.0 ) - prd_face_mask_a_0 = np.clip (prd_face_mask_a_0, 0.0, 1.0) - else: - predicted = self.predictor_func (predictor_input_bgr) - prd_face_bgr = np.clip (predicted, 0, 1.0 ) - prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (self.predictor_input_size,self.predictor_input_size)) - - if self.super_resolution: - if debug: - tmp = cv2.resize (prd_face_bgr, (output_size,output_size), cv2.INTER_CUBIC) - debugs += [ np.clip( cv2.warpAffine( tmp, face_output_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] - - prd_face_bgr = self.dc_upscale(prd_face_bgr) - if debug: - debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] - - if self.predictor_masked: - prd_face_mask_a_0 = cv2.resize (prd_face_mask_a_0, (output_size, output_size), cv2.INTER_CUBIC) - else: - prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (output_size, output_size), cv2.INTER_CUBIC) - - if self.mask_mode == 2: #dst - prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (output_size,output_size), cv2.INTER_CUBIC) - elif self.mask_mode >= 3 and self.mask_mode <= 6: - - if self.mask_mode == 3 or self.mask_mode == 5 or self.mask_mode == 6: - prd_face_bgr_256 = cv2.resize (prd_face_bgr, (256,256) ) - prd_face_bgr_256_mask = self.fan_seg.extract( prd_face_bgr_256 ) - FAN_prd_face_mask_a_0 = cv2.resize (prd_face_bgr_256_mask, (output_size,output_size), cv2.INTER_CUBIC) - - if self.mask_mode == 4 or self.mask_mode == 5 or self.mask_mode == 6: - face_256_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, 256, face_type=FaceType.FULL) - dst_face_256_bgr = cv2.warpAffine(img_bgr, face_256_mat, (256, 256), flags=cv2.INTER_LANCZOS4 ) - dst_face_256_mask = self.fan_seg.extract( dst_face_256_bgr ) - FAN_dst_face_mask_a_0 = cv2.resize (dst_face_256_mask, (output_size,output_size), cv2.INTER_CUBIC) - - if self.mask_mode == 3: #FAN-prd - prd_face_mask_a_0 = FAN_prd_face_mask_a_0 - elif self.mask_mode == 4: #FAN-dst - prd_face_mask_a_0 = FAN_dst_face_mask_a_0 - elif self.mask_mode == 5: - prd_face_mask_a_0 = FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0 - elif self.mask_mode == 6: - prd_face_mask_a_0 = prd_face_mask_a_0 * FAN_prd_face_mask_a_0 * FAN_dst_face_mask_a_0 - - prd_face_mask_a_0[ prd_face_mask_a_0 < 0.001 ] = 0.0 - - prd_face_mask_a = prd_face_mask_a_0[...,np.newaxis] - prd_face_mask_aaa = np.repeat (prd_face_mask_a, (3,), axis=-1) - - img_face_mask_aaa = cv2.warpAffine( prd_face_mask_aaa, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), flags=cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4 ) - img_face_mask_aaa = np.clip (img_face_mask_aaa, 0.0, 1.0) - img_face_mask_aaa [ img_face_mask_aaa <= 0.1 ] = 0.0 #get rid of noise - - if debug: - debugs += [img_face_mask_aaa.copy()] - - - out_img = img_bgr.copy() - - if self.mode == 'raw': - if self.raw_mode == 'rgb' or self.raw_mode == 'rgb-mask': - out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) - - if self.raw_mode == 'rgb-mask': - out_img = np.concatenate ( [out_img, np.expand_dims (img_face_mask_aaa[:,:,0],-1)], -1 ) - - if self.raw_mode == 'mask-only': - out_img = img_face_mask_aaa - - if self.raw_mode == 'predicted-only': - out_img = cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(out_img.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) - - else: - #averaging [lenx, leny, maskx, masky] by grayscale gradients of upscaled mask - ar = [] - for i in range(1, 10): - maxregion = np.argwhere( img_face_mask_aaa > i / 10.0 ) - if maxregion.size != 0: - miny,minx = maxregion.min(axis=0)[:2] - maxy,maxx = maxregion.max(axis=0)[:2] - lenx = maxx - minx - leny = maxy - miny - if min(lenx,leny) >= 4: - ar += [ [ lenx, leny] ] - - if len(ar) > 0: - lenx, leny = np.mean ( ar, axis=0 ) - lowest_len = min (lenx, leny) - if debug: - io.log_info ("lenx/leny:(%d/%d) " % (lenx, leny ) ) - io.log_info ("lowest_len = %f" % (lowest_len) ) - - if self.erode_mask_modifier != 0: - ero = int( lowest_len * ( 0.126 - lowest_len * 0.00004551365 ) * 0.01*self.erode_mask_modifier ) - if debug: - io.log_info ("erode_size = %d" % (ero) ) - if ero > 0: - img_face_mask_aaa = cv2.erode(img_face_mask_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero,ero)), iterations = 1 ) - elif ero < 0: - img_face_mask_aaa = cv2.dilate(img_face_mask_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(-ero,-ero)), iterations = 1 ) - - img_mask_blurry_aaa = img_face_mask_aaa - - 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_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[:,-prd_border_size:,:] = 0 - prd_hborder_rect_mask_a[-prd_border_size:,:,:] = 0 - prd_hborder_rect_mask_a = np.expand_dims(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, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4 ) - img_prd_hborder_rect_mask_a = np.expand_dims (img_prd_hborder_rect_mask_a, -1) - img_mask_blurry_aaa *= img_prd_hborder_rect_mask_a - img_mask_blurry_aaa = np.clip( img_mask_blurry_aaa, 0, 1.0 ) - - if debug: - debugs += [img_mask_blurry_aaa.copy()] - - if self.blur_mask_modifier > 0: - blur = int( lowest_len * 0.10 * 0.01*self.blur_mask_modifier ) - if debug: - io.log_info ("blur_size = %d" % (blur) ) - if blur > 0: - img_mask_blurry_aaa = cv2.blur(img_mask_blurry_aaa, (blur, blur) ) - - img_mask_blurry_aaa = np.clip( img_mask_blurry_aaa, 0, 1.0 ) - face_mask_blurry_aaa = cv2.warpAffine( img_mask_blurry_aaa, face_mat, (output_size, output_size) ) - - if debug: - debugs += [img_mask_blurry_aaa.copy()] - - if 'seamless' not in self.mode and self.color_transfer_mode is not None: - if self.color_transfer_mode == 'rct': - if debug: - debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] - - prd_face_bgr = imagelib.reinhard_color_transfer ( np.clip( (prd_face_bgr*255).astype(np.uint8), 0, 255), - np.clip( (dst_face_bgr*255).astype(np.uint8), 0, 255), - source_mask=prd_face_mask_a, target_mask=prd_face_mask_a) - prd_face_bgr = np.clip( prd_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0) - - if debug: - debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] - - - elif self.color_transfer_mode == 'lct': - if debug: - debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] - - prd_face_bgr = imagelib.linear_color_transfer (prd_face_bgr, dst_face_bgr) - prd_face_bgr = np.clip( prd_face_bgr, 0.0, 1.0) - - if debug: - debugs += [ np.clip( cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] - - if self.mode == 'hist-match-bw': - prd_face_bgr = cv2.cvtColor(prd_face_bgr, cv2.COLOR_BGR2GRAY) - prd_face_bgr = np.repeat( np.expand_dims (prd_face_bgr, -1), (3,), -1 ) - - if self.mode == 'hist-match' or self.mode == 'hist-match-bw': - if debug: - debugs += [ cv2.warpAffine( prd_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) ] - - hist_mask_a = np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32) - - if self.masked_hist_match: - hist_mask_a *= prd_face_mask_a - - white = (1.0-hist_mask_a)* np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32) - - hist_match_1 = prd_face_bgr*hist_mask_a + white - hist_match_1[ hist_match_1 > 1.0 ] = 1.0 - - hist_match_2 = dst_face_bgr*hist_mask_a + white - hist_match_2[ hist_match_1 > 1.0 ] = 1.0 - - prd_face_bgr = imagelib.color_hist_match(hist_match_1, hist_match_2, self.hist_match_threshold ) - - #if self.masked_hist_match: - # prd_face_bgr -= white - - if self.mode == 'hist-match-bw': - prd_face_bgr = prd_face_bgr.astype(dtype=np.float32) - - out_img = cv2.warpAffine( prd_face_bgr, face_output_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: - debugs += [out_img.copy()] - - if self.mode == 'overlay': - pass - - if 'seamless' in self.mode: - #mask used for cv2.seamlessClone - img_face_seamless_mask_a = None - img_face_mask_a = img_mask_blurry_aaa[...,0:1] - for i in range(1,10): - a = img_face_mask_a > i / 10.0 - if len(np.argwhere(a)) == 0: - continue - img_face_seamless_mask_a = img_mask_blurry_aaa[...,0:1].copy() - img_face_seamless_mask_a[a] = 1.0 - img_face_seamless_mask_a[img_face_seamless_mask_a <= i / 10.0] = 0.0 - break - - try: - #calc same bounding rect and center point as in cv2.seamlessClone to prevent jittering - l,t,w,h = cv2.boundingRect( (img_face_seamless_mask_a*255).astype(np.uint8) ) - s_maskx, s_masky = int(l+w/2), int(t+h/2) - - out_img = cv2.seamlessClone( (out_img*255).astype(np.uint8), (img_bgr*255).astype(np.uint8), (img_face_seamless_mask_a*255).astype(np.uint8), (s_maskx,s_masky) , cv2.NORMAL_CLONE ) - out_img = out_img.astype(dtype=np.float32) / 255.0 - except Exception as e: - #seamlessClone may fail in some cases - e_str = traceback.format_exc() - - if 'MemoryError' in e_str: - raise Exception("Seamless fail: " + e_str) #reraise MemoryError in order to reprocess this data by other processes - else: - print ("Seamless fail: " + e_str) - - if debug: - debugs += [out_img.copy()] - - out_img = np.clip( img_bgr*(1-img_mask_blurry_aaa) + (out_img*img_mask_blurry_aaa) , 0, 1.0 ) - - if 'seamless' in self.mode and self.color_transfer_mode is not None: - out_face_bgr = cv2.warpAffine( out_img, face_mat, (output_size, output_size) ) - - if self.color_transfer_mode == 'rct': - if debug: - debugs += [ np.clip( cv2.warpAffine( out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] - - new_out_face_bgr = imagelib.reinhard_color_transfer ( np.clip( (out_face_bgr*255).astype(np.uint8), 0, 255), - np.clip( (dst_face_bgr*255).astype(np.uint8), 0, 255), - source_mask=face_mask_blurry_aaa, target_mask=face_mask_blurry_aaa) - new_out_face_bgr = np.clip( new_out_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0) - - if debug: - debugs += [ np.clip( cv2.warpAffine( new_out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] - - - elif self.color_transfer_mode == 'lct': - if debug: - debugs += [ np.clip( cv2.warpAffine( out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] - - new_out_face_bgr = imagelib.linear_color_transfer (out_face_bgr, dst_face_bgr) - new_out_face_bgr = np.clip( new_out_face_bgr, 0.0, 1.0) - - if debug: - debugs += [ np.clip( cv2.warpAffine( new_out_face_bgr, face_output_mat, img_size, np.zeros(img_bgr.shape, dtype=np.float32), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ), 0, 1.0) ] - - new_out = cv2.warpAffine( new_out_face_bgr, face_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) - out_img = np.clip( img_bgr*(1-img_mask_blurry_aaa) + (new_out*img_mask_blurry_aaa) , 0, 1.0 ) - - if self.mode == 'seamless-hist-match': - out_face_bgr = cv2.warpAffine( out_img, face_mat, (output_size, output_size) ) - new_out_face_bgr = imagelib.color_hist_match(out_face_bgr, dst_face_bgr, self.hist_match_threshold) - new_out = cv2.warpAffine( new_out_face_bgr, face_mat, img_size, img_bgr.copy(), cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT ) - out_img = np.clip( img_bgr*(1-img_mask_blurry_aaa) + (new_out*img_mask_blurry_aaa) , 0, 1.0 ) - - if self.final_image_color_degrade_power != 0: - if debug: - debugs += [out_img.copy()] - out_img_reduced = imagelib.reduce_colors(out_img, 256) - if self.final_image_color_degrade_power == 100: - out_img = out_img_reduced - else: - alpha = self.final_image_color_degrade_power / 100.0 - out_img = (out_img*(1.0-alpha) + out_img_reduced*alpha) - - if self.alpha: - out_img = np.concatenate ( [out_img, np.expand_dims (img_mask_blurry_aaa[:,:,0],-1)], -1 ) - - out_img = np.clip (out_img, 0.0, 1.0 ) - - if debug: - debugs += [out_img.copy()] - - return debugs if debug else out_img diff --git a/converters/FrameInfo.py b/converters/FrameInfo.py new file mode 100644 index 0000000..e48154d --- /dev/null +++ b/converters/FrameInfo.py @@ -0,0 +1,6 @@ +class FrameInfo(object): + def __init__(self, filename=None, landmarks_list=None): + self.filename = filename + self.landmarks_list = landmarks_list or [] + self.motion_deg = 0 + self.motion_power = 0 \ No newline at end of file diff --git a/converters/__init__.py b/converters/__init__.py index ef544cb..8aa1057 100644 --- a/converters/__init__.py +++ b/converters/__init__.py @@ -1,4 +1,4 @@ -from .Converter import Converter -from .ConverterMasked import ConverterMasked -from .ConverterImage import ConverterImage -from .ConverterAvatar import ConverterAvatar +from .FrameInfo import FrameInfo +from .ConverterConfig import ConverterConfig, ConverterConfigMasked, ConverterConfigFaceAvatar +from .ConvertMasked import ConvertMasked +from .ConvertAvatar import ConvertFaceAvatar diff --git a/doc/manual_en_google_translated.docx b/doc/manual_en_google_translated.docx index cd8fd5a..d575ec6 100644 Binary files a/doc/manual_en_google_translated.docx and b/doc/manual_en_google_translated.docx differ diff --git a/doc/manual_en_google_translated.pdf b/doc/manual_en_google_translated.pdf index a03f7dc..e878346 100644 Binary files a/doc/manual_en_google_translated.pdf and b/doc/manual_en_google_translated.pdf differ diff --git a/doc/manual_ru.pdf b/doc/manual_ru.pdf index 8a15598..9567d7e 100644 Binary files a/doc/manual_ru.pdf and b/doc/manual_ru.pdf differ diff --git a/doc/manual_ru_source.docx b/doc/manual_ru_source.docx index c73bdf6..1686ec6 100644 Binary files a/doc/manual_ru_source.docx and b/doc/manual_ru_source.docx differ diff --git a/facelib/FaceType.py b/facelib/FaceType.py index dba999d..8f769fd 100644 --- a/facelib/FaceType.py +++ b/facelib/FaceType.py @@ -4,10 +4,11 @@ class FaceType(IntEnum): HALF = 0, FULL = 1, HEAD = 2, - + FULL_NO_ALIGN = 5, + HEAD_NO_ALIGN = 6, MARK_ONLY = 10, #no align at all, just embedded faceinfo - + @staticmethod def fromString (s): r = from_string_dict.get (s.lower()) @@ -24,10 +25,12 @@ from_string_dict = {'half_face': FaceType.HALF, 'head' : FaceType.HEAD, 'mark_only' : FaceType.MARK_ONLY, 'full_face_no_align' : FaceType.FULL_NO_ALIGN, + 'head_no_align' : FaceType.HEAD_NO_ALIGN, } to_string_dict = { FaceType.HALF : 'half_face', FaceType.FULL : 'full_face', FaceType.HEAD : 'head', FaceType.MARK_ONLY :'mark_only', - FaceType.FULL_NO_ALIGN : 'full_face_no_align' + FaceType.FULL_NO_ALIGN : 'full_face_no_align', + FaceType.HEAD_NO_ALIGN : 'head_no_align' } diff --git a/facelib/LandmarksProcessor.py b/facelib/LandmarksProcessor.py index 7c0b697..3b8b547 100644 --- a/facelib/LandmarksProcessor.py +++ b/facelib/LandmarksProcessor.py @@ -116,7 +116,7 @@ def transform_points(points, mat, invert=False): 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): if not isinstance(image_landmarks, np.ndarray): image_landmarks = np.array (image_landmarks) @@ -142,13 +142,16 @@ def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0): if face_type == FaceType.FULL_NO_ALIGN: face_type = FaceType.FULL remove_align = True - + elif face_type == FaceType.HEAD_NO_ALIGN: + face_type = FaceType.HEAD + remove_align = True + if face_type == FaceType.HALF: padding = 0 elif face_type == FaceType.FULL: padding = (output_size / 64) * 12 elif face_type == FaceType.HEAD: - padding = (output_size / 64) * 24 + padding = (output_size / 64) * 21 else: raise ValueError ('wrong face_type: ', face_type) @@ -157,13 +160,13 @@ def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0): mat[:,2] += padding mat *= (1 / scale) 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) diff --git a/imagelib/__init__.py b/imagelib/__init__.py index 14ed304..1d8ce5e 100644 --- a/imagelib/__init__.py +++ b/imagelib/__init__.py @@ -1,26 +1,21 @@ from .estimate_sharpness import estimate_sharpness from .equalize_and_stack_square import equalize_and_stack_square -from .text import get_text_image -from .text import get_draw_text_lines +from .text import get_text_image, get_draw_text_lines -from .draw import draw_polygon -from .draw import draw_rect +from .draw import draw_polygon, draw_rect from .morph import morph_by_points -from .warp import gen_warp_params -from .warp import warp_by_params +from .warp import gen_warp_params, warp_by_params from .reduce_colors import reduce_colors -from .color_transfer import color_hist_match -from .color_transfer import reinhard_color_transfer -from .color_transfer import linear_color_transfer +from .color_transfer import color_hist_match, reinhard_color_transfer, linear_color_transfer from .DCSCN import DCSCN -from .common import normalize_channels +from .common import normalize_channels, overlay_alpha_image from .IEPolys import IEPolys diff --git a/imagelib/blur.py b/imagelib/blur.py index e12ccfd..4aaa518 100644 --- a/imagelib/blur.py +++ b/imagelib/blur.py @@ -1,143 +1,9 @@ -import math +import cv2 import numpy as np -from PIL import Image -from scipy.signal import convolve2d -from skimage.draw import line -class LineDictionary: - def __init__(self): - self.lines = {} - self.Create3x3Lines() - self.Create5x5Lines() - self.Create7x7Lines() - self.Create9x9Lines() - return - - def Create3x3Lines(self): - lines = {} - lines[0] = [1,0,1,2] - lines[45] = [2,0,0,2] - lines[90] = [0,1,2,1] - lines[135] = [0,0,2,2] - self.lines[3] = lines - return - - def Create5x5Lines(self): - lines = {} - lines[0] = [2,0,2,4] - lines[22.5] = [3,0,1,4] - lines[45] = [0,4,4,0] - lines[67.5] = [0,3,4,1] - lines[90] = [0,2,4,2] - lines[112.5] = [0,1,4,3] - lines[135] = [0,0,4,4] - lines[157.5]= [1,0,3,4] - self.lines[5] = lines - return - - def Create7x7Lines(self): - lines = {} - lines[0] = [3,0,3,6] - lines[15] = [4,0,2,6] - lines[30] = [5,0,1,6] - lines[45] = [6,0,0,6] - lines[60] = [6,1,0,5] - lines[75] = [6,2,0,4] - lines[90] = [0,3,6,3] - lines[105] = [0,2,6,4] - lines[120] = [0,1,6,5] - lines[135] = [0,0,6,6] - lines[150] = [1,0,5,6] - lines[165] = [2,0,4,6] - self.lines[7] = lines - return - - def Create9x9Lines(self): - lines = {} - lines[0] = [4,0,4,8] - lines[11.25] = [5,0,3,8] - lines[22.5] = [6,0,2,8] - lines[33.75] = [7,0,1,8] - lines[45] = [8,0,0,8] - lines[56.25] = [8,1,0,7] - lines[67.5] = [8,2,0,6] - lines[78.75] = [8,3,0,5] - lines[90] = [8,4,0,4] - lines[101.25] = [0,3,8,5] - lines[112.5] = [0,2,8,6] - lines[123.75] = [0,1,8,7] - lines[135] = [0,0,8,8] - lines[146.25] = [1,0,7,8] - lines[157.5] = [2,0,6,8] - lines[168.75] = [3,0,5,8] - self.lines[9] = lines - return - -lineLengths =[3,5,7,9] -lineTypes = ["full", "right", "left"] - -lineDict = LineDictionary() - -def LinearMotionBlur_random(img): - lineLengthIdx = np.random.randint(0, len(lineLengths)) - lineTypeIdx = np.random.randint(0, len(lineTypes)) - lineLength = lineLengths[lineLengthIdx] - lineType = lineTypes[lineTypeIdx] - lineAngle = randomAngle(lineLength) - return LinearMotionBlur(img, lineLength, lineAngle, lineType) - -def LinearMotionBlur(img, dim, angle, linetype='full'): - if len(img.shape) == 2: - h, w = img.shape - c = 1 - img = img[...,np.newaxis] - elif len(img.shape) == 3: - h,w,c = img.shape - else: - raise ValueError('unsupported img.shape') - - kernel = LineKernel(dim, angle, linetype) - - imgs = [] - for i in range(c): - imgs.append ( convolve2d(img[...,i], kernel, mode='same') ) - - img = np.stack(imgs, axis=-1) - img = np.squeeze(img) - return img - -def LineKernel(dim, angle, linetype): - kernelwidth = dim - kernelCenter = int(math.floor(dim/2)) - angle = SanitizeAngleValue(kernelCenter, angle) - kernel = np.zeros((kernelwidth, kernelwidth), dtype=np.float32) - lineAnchors = lineDict.lines[dim][angle] - if(linetype == 'right'): - lineAnchors[0] = kernelCenter - lineAnchors[1] = kernelCenter - if(linetype == 'left'): - lineAnchors[2] = kernelCenter - lineAnchors[3] = kernelCenter - rr,cc = line(lineAnchors[0], lineAnchors[1], lineAnchors[2], lineAnchors[3]) - kernel[rr,cc]=1 - normalizationFactor = np.count_nonzero(kernel) - kernel = kernel / normalizationFactor - return kernel - -def SanitizeAngleValue(kernelCenter, angle): - numDistinctLines = kernelCenter * 4 - angle = math.fmod(angle, 180.0) - validLineAngles = np.linspace(0,180, numDistinctLines, endpoint = False) - angle = nearestValue(angle, validLineAngles) - return angle - -def nearestValue(theta, validAngles): - idx = (np.abs(validAngles-theta)).argmin() - return validAngles[idx] - -def randomAngle(kerneldim): - kernelCenter = int(math.floor(kerneldim/2)) - numDistinctLines = kernelCenter * 4 - validLineAngles = np.linspace(0,180, numDistinctLines, endpoint = False) - angleIdx = np.random.randint(0, len(validLineAngles)) - return int(validLineAngles[angleIdx]) \ No newline at end of file +def LinearMotionBlur(image, size, angle): + k = np.zeros((size, size), dtype=np.float32) + k[ (size-1)// 2 , :] = np.ones(size, dtype=np.float32) + k = cv2.warpAffine(k, cv2.getRotationMatrix2D( (size / 2 -0.5 , size / 2 -0.5 ) , angle, 1.0), (size, size) ) + k = k * ( 1.0 / np.sum(k) ) + return cv2.filter2D(image, -1, k) \ No newline at end of file diff --git a/imagelib/common.py b/imagelib/common.py index 229387f..3838c43 100644 --- a/imagelib/common.py +++ b/imagelib/common.py @@ -9,13 +9,28 @@ def normalize_channels(img, target_channels): h, w, c = img.shape else: raise ValueError("normalize: incorrect image dimensions.") - + if c == 0 and target_channels > 0: - img = img[...,np.newaxis] + img = img[...,np.newaxis] if c == 1 and target_channels > 1: - img = np.repeat (img, target_channels, -1) - if c > target_channels: + img = np.repeat (img, target_channels, -1) + if c > target_channels: img = img[...,0:target_channels] c = target_channels - - return img \ No newline at end of file + + return img + +def overlay_alpha_image(img_target, img_source, xy_offset=(0,0) ): + (h,w,c) = img_source.shape + if c != 4: + raise ValueError("overlay_alpha_image, img_source must have 4 channels") + + x1, x2 = xy_offset[0], xy_offset[0] + w + y1, y2 = xy_offset[1], xy_offset[1] + h + + alpha_s = img_source[:, :, 3] / 255.0 + alpha_l = 1.0 - alpha_s + + for c in range(0, 3): + img_target[y1:y2, x1:x2, c] = (alpha_s * img_source[:, :, c] + + alpha_l * img_target[y1:y2, x1:x2, c]) \ No newline at end of file diff --git a/interact/interact.py b/interact/interact.py index bda74af..b1b4e23 100644 --- a/interact/interact.py +++ b/interact/interact.py @@ -67,20 +67,12 @@ class InteractBase(object): def log_info(self, msg, end='\n'): if self.pg_bar is not None: - try: # Attempt print before the pb - tqdm.write(msg) - return - except: - pass #Fallback to normal print upon failure + print ("\n") print (msg, end=end) def log_err(self, msg, end='\n'): if self.pg_bar is not None: - try: # Attempt print before the pb - tqdm.write(f'{self.error_log_line_prefix}{msg}') - return - except: - pass #Fallback to normal print upon failure + print ("\n") print (f'{self.error_log_line_prefix}{msg}', end=end) def named_window(self, wnd_name): @@ -166,9 +158,9 @@ class InteractBase(object): self.pg_bar = tqdm( data, desc=desc, leave=leave, ascii=True ) for x in self.pg_bar: yield x - self.pg_bar.close() + self.pg_bar.close() self.pg_bar = None - + def process_messages(self, sleep_time=0): self.on_process_messages(sleep_time) diff --git a/joblib/SubprocessFunctionCaller.py b/joblib/SubprocessFunctionCaller.py index 4ea3101..9e326d3 100644 --- a/joblib/SubprocessFunctionCaller.py +++ b/joblib/SubprocessFunctionCaller.py @@ -30,8 +30,15 @@ class SubprocessFunctionCaller(object): result = self.func ( *obj['args'], **obj['kwargs'] ) self.s2c.put (result) + def __getstate__(self): + #disable pickling this class + return dict() + + def __setstate__(self, d): + self.__dict__.update(d) + @staticmethod - def make_pair( func ): + def make_pair(func): s2c = multiprocessing.Queue() c2s = multiprocessing.Queue() lock = multiprocessing.Lock() diff --git a/joblib/SubprocessorBase.py b/joblib/SubprocessorBase.py index dd48424..a9cbc36 100644 --- a/joblib/SubprocessorBase.py +++ b/joblib/SubprocessorBase.py @@ -87,13 +87,15 @@ class Subprocessor(object): c2s.put ( {'op': 'error', 'data' : data} ) #overridable - def __init__(self, name, SubprocessorCli_class, no_response_time_sec = 0): + def __init__(self, name, SubprocessorCli_class, no_response_time_sec = 0, io_loop_sleep_time=0.005, initialize_subprocesses_in_serial=True): if not issubclass(SubprocessorCli_class, Subprocessor.Cli): raise ValueError("SubprocessorCli_class must be subclass of Subprocessor.Cli") self.name = name self.SubprocessorCli_class = SubprocessorCli_class self.no_response_time_sec = no_response_time_sec + self.io_loop_sleep_time = io_loop_sleep_time + self.initialize_subprocesses_in_serial = initialize_subprocesses_in_serial #overridable def process_info_generator(self): @@ -133,7 +135,8 @@ class Subprocessor(object): #overridable def on_tick(self): #tick in main loop - pass + #return True if system can be finalized when no data in get_data, orelse False + return True #overridable def on_check_run(self): @@ -157,23 +160,24 @@ class Subprocessor(object): self.clis.append (cli) - while True: - while not cli.c2s.empty(): - obj = cli.c2s.get() - op = obj.get('op','') - if op == 'init_ok': - cli.state = 0 - elif op == 'log_info': - io.log_info(obj['msg']) - elif op == 'log_err': - io.log_err(obj['msg']) - elif op == 'error': - cli.kill() - self.clis.remove(cli) + if self.initialize_subprocesses_in_serial: + while True: + while not cli.c2s.empty(): + obj = cli.c2s.get() + op = obj.get('op','') + if op == 'init_ok': + cli.state = 0 + elif op == 'log_info': + io.log_info(obj['msg']) + elif op == 'log_err': + io.log_err(obj['msg']) + elif op == 'error': + cli.kill() + self.clis.remove(cli) + break + if cli.state == 0: break - if cli.state == 0: - break - io.process_messages(0.005) + io.process_messages(0.005) except: raise Exception ("Unable to start subprocess %s" % (name)) @@ -232,6 +236,15 @@ class Subprocessor(object): elif op == 'progress_bar_inc': io.progress_bar_inc(obj['c']) + for cli in self.clis[:]: + if cli.state == 1: + if self.no_response_time_sec != 0 and (time.time() - cli.sent_time) > self.no_response_time_sec: + #subprocess busy too long + print ( '%s doesnt response, terminating it.' % (cli.name) ) + self.on_data_return (cli.host_dict, cli.sent_data ) + cli.kill() + self.clis.remove(cli) + for cli in self.clis[:]: if cli.state == 0: #free state of subprocess, get some data from get_data @@ -243,19 +256,14 @@ class Subprocessor(object): cli.sent_data = data cli.state = 1 - elif cli.state == 1: - if self.no_response_time_sec != 0 and (time.time() - cli.sent_time) > self.no_response_time_sec: - #subprocess busy too long - print ( '%s doesnt response, terminating it.' % (cli.name) ) - self.on_data_return (cli.host_dict, cli.sent_data ) - cli.kill() - self.clis.remove(cli) + if self.io_loop_sleep_time != 0: + io.process_messages(self.io_loop_sleep_time) - if all ([cli.state == 0 for cli in self.clis]): + if self.on_tick() and all ([cli.state == 0 for cli in self.clis]): #all subprocesses free and no more data available to process, ending loop break - io.process_messages(0.005) - self.on_tick() + + #gracefully terminating subprocesses for cli in self.clis[:]: diff --git a/main.py b/main.py index 240c044..bd2b537 100644 --- a/main.py +++ b/main.py @@ -48,8 +48,8 @@ if __name__ == "__main__": p.add_argument('--manual-window-size', type=int, dest="manual_window_size", default=1368, help="Manual fix window size. Default: 1368.") p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU. Forces to use MT extractor.") p.set_defaults (func=process_extract) - - + + def process_dev_extract_umd_csv(arguments): os_utils.set_process_lowest_prio() from mainscripts import Extractor @@ -80,7 +80,7 @@ if __name__ == "__main__": p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Extract on CPU.") p.set_defaults (func=process_extract_fanseg) """ - + def process_sort(arguments): os_utils.set_process_lowest_prio() from mainscripts import Sorter @@ -103,17 +103,17 @@ if __name__ == "__main__": if arguments.recover_original_aligned_filename: Util.recover_original_aligned_filename (input_path=arguments.input_dir) - + #if arguments.remove_fanseg: # Util.remove_fanseg_folder (input_path=arguments.input_dir) - + p = subparsers.add_parser( "util", help="Utilities.") 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('--convert-png-to-jpg', action="store_true", dest="convert_png_to_jpg", default=False, help="Convert DeepFaceLAB PNG files to JPEG.") p.add_argument('--add-landmarks-debug-images', action="store_true", dest="add_landmarks_debug_images", default=False, help="Add landmarks debug image for aligned faces.") p.add_argument('--recover-original-aligned-filename', action="store_true", dest="recover_original_aligned_filename", default=False, help="Recover original aligned filename.") #p.add_argument('--remove-fanseg', action="store_true", dest="remove_fanseg", default=False, help="Remove fanseg mask from aligned faces.") - + p.set_defaults (func=process_util) def process_train(arguments): @@ -153,7 +153,6 @@ if __name__ == "__main__": 'aligned_dir' : arguments.aligned_dir, 'model_dir' : arguments.model_dir, 'model_name' : arguments.model_name, - 'debug' : arguments.debug, } device_args = {'cpu_only' : arguments.cpu_only, 'force_gpu_idx' : arguments.force_gpu_idx, @@ -167,7 +166,6 @@ if __name__ == "__main__": 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('--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('--debug', action="store_true", dest="debug", default=False, help="Debug converter.") p.add_argument('--force-gpu-idx', type=int, dest="force_gpu_idx", default=-1, help="Force to choose this GPU idx.") p.add_argument('--cpu-only', action="store_true", dest="cpu_only", default=False, help="Convert on CPU.") p.set_defaults(func=process_convert) @@ -242,7 +240,7 @@ if __name__ == "__main__": p.add_argument('--confirmed-dir', required=True, action=fixPathAction, dest="confirmed_dir", help="This is where the labeled faces will be stored.") p.add_argument('--skipped-dir', required=True, action=fixPathAction, dest="skipped_dir", help="This is where the labeled faces will be stored.") p.set_defaults(func=process_labelingtool_edit_mask) - + def bad_args(arguments): parser.print_help() exit(0) diff --git a/mainscripts/Converter.py b/mainscripts/Converter.py index 53c3e4f..fdba366 100644 --- a/mainscripts/Converter.py +++ b/mainscripts/Converter.py @@ -1,37 +1,77 @@ -import sys +import math import multiprocessing import operator import os import shutil +import sys import time import traceback from pathlib import Path import cv2 import numpy as np +import numpy.linalg as npla -from converters import Converter +import imagelib +from converters import FrameInfo, ConverterConfig, ConvertMasked, ConvertFaceAvatar +from facelib import FaceType, FANSegmentator, LandmarksProcessor from interact import interact as io from joblib import SubprocessFunctionCaller, Subprocessor from utils import Path_utils from utils.cv2_utils import * from utils.DFLJPG import DFLJPG from utils.DFLPNG import DFLPNG -from imagelib import normalize_channels + +from .ConverterScreen import Screen, ScreenManager + +CONVERTER_DEBUG = False class ConvertSubprocessor(Subprocessor): + + class Frame(object): + def __init__(self, prev_temporal_frame_infos=None, + frame_info=None, + next_temporal_frame_infos=None): + self.prev_temporal_frame_infos = prev_temporal_frame_infos + self.frame_info = frame_info + self.next_temporal_frame_infos = next_temporal_frame_infos + self.output_filename = None + + self.idx = None + self.cfg = None + self.is_done = False + self.is_processing = False + self.is_shown = False + self.image = None + + class ProcessingFrame(object): + def __init__(self, idx=None, + cfg=None, + prev_temporal_frame_infos=None, + frame_info=None, + next_temporal_frame_infos=None, + output_filename=None, + need_return_image = False): + self.idx = idx + self.cfg = cfg + self.prev_temporal_frame_infos = prev_temporal_frame_infos + self.frame_info = frame_info + self.next_temporal_frame_infos = next_temporal_frame_infos + self.output_filename = output_filename + + self.need_return_image = need_return_image + if self.need_return_image: + self.image = None + class Cli(Subprocessor.Cli): #override def on_initialize(self, client_dict): - io.log_info ('Running on %s.' % (client_dict['device_name']) ) + self.log_info ('Running on %s.' % (client_dict['device_name']) ) self.device_idx = client_dict['device_idx'] self.device_name = client_dict['device_name'] - 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.alignments = client_dict['alignments'] - self.debug = client_dict['debug'] + self.predictor_func = client_dict['predictor_func'] + self.dcscn_upscale_func = client_dict['dcscn_upscale_func'] #transfer and set stdin in order to work code.interact in debug subprocess stdin_fd = client_dict['stdin_fd'] @@ -44,187 +84,372 @@ class ConvertSubprocessor(Subprocessor): #therefore forcing active_DeviceConfig to CPU only nnlib.active_DeviceConfig = nnlib.DeviceConfig (cpu_only=True) - self.converter.on_cli_initialize() + print(nnlib.backend) + + self.fanseg_by_face_type = {} + self.fanseg_input_size = 256 + + def fanseg_extract(face_type, *args, **kwargs): + fanseg = self.fanseg_by_face_type.get(face_type, None) + if self.fanseg_by_face_type.get(face_type, None) is None: + fanseg = FANSegmentator( self.fanseg_input_size , FaceType.toString( face_type ) ) + self.fanseg_by_face_type[face_type] = fanseg + + return fanseg.extract(*args, **kwargs) + + self.fanseg_extract_func = fanseg_extract return None #override - def process_data(self, data): - files_processed = 1 - faces_processed = 0 - - idx, = data - filename = self.input_data[idx][0] + def process_data(self, pf): #pf=ProcessingFrame + cfg = pf.cfg.copy() + cfg.predictor_func = self.predictor_func + + frame_info = pf.frame_info + + filename = frame_info.filename + landmarks_list = frame_info.landmarks_list + filename_path = Path(filename) + output_filename = pf.output_filename + need_return_image = pf.need_return_image - 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 len(landmarks_list) == 0: + self.log_info ( 'no faces found for %s, copying without faces' % (filename_path.name) ) - if filename_path.suffix == '.png': - shutil.copy ( str(filename_path), str(output_filename_path) ) - else: - image = cv2_imread(str(filename_path)) - cv2_imwrite ( str(output_filename_path), image ) + if filename_path.suffix == '.png': + shutil.copy (filename, output_filename ) else: - image = (cv2_imread(str(filename_path)) / 255.0).astype(np.float32) - image = normalize_channels (image, 3) - - faces = self.alignments[filename_path.stem] + img_bgr = cv2_imread(filename) + cv2_imwrite (output_filename, img_bgr) - if self.debug: - debug_images = [] + if need_return_image: + img_bgr = cv2_imread(filename) + pf.image = img_bgr + else: + if cfg.type == ConverterConfig.TYPE_MASKED: + cfg.dcscn_upscale_func = self.dcscn_upscale_func + cfg.fanseg_input_size = self.fanseg_input_size + cfg.fanseg_extract_func = self.fanseg_extract_func - for face_num, image_landmarks in enumerate(faces): - try: - if self.debug: - self.log_info ( '\nConverting face_num [%d] in file [%s]' % (face_num, filename_path) ) + try: + final_img = ConvertMasked (cfg, frame_info) + except Exception as e: + e_str = traceback.format_exc() + if 'MemoryError' in e_str: + raise Subprocessor.SilenceException + else: + raise Exception( 'Error while converting file [%s]: %s' % (filename, e_str) ) - if self.debug: - debug_images += self.converter.cli_convert_face(image, image_landmarks, self.debug) - else: - image = self.converter.cli_convert_face(image, image_landmarks, self.debug) + elif cfg.type == ConverterConfig.TYPE_FACE_AVATAR: + final_img = ConvertFaceAvatar (cfg, pf.prev_temporal_frame_infos, + pf.frame_info, + pf.next_temporal_frame_infos ) - except Exception as e: - e_str = traceback.format_exc() - if 'MemoryError' in e_str: - raise Subprocessor.SilenceException - else: - raise Exception( 'Error while converting face_num [%d] in file [%s]: %s' % (face_num, filename_path, e_str) ) + if output_filename is not None and final_img is not None: + cv2_imwrite (output_filename, final_img ) - if self.debug: - return (1, debug_images) + if need_return_image: + pf.image = final_img - 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 self.debug: - 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 pf #overridable - def get_data_name (self, data): + def get_data_name (self, pf): #return string identificator of your data - idx, = data - return self.input_data[idx][0] + return pf.frame_info.filename #override - def __init__(self, converter, input_data, output_path, alignments, debug = False): - super().__init__('Converter', ConvertSubprocessor.Cli, 86400 if debug == True else 60) + def __init__(self, is_interactive, converter_config, frames, output_path): + if len (frames) == 0: + raise ValueError ("len (frames) == 0") - self.converter = converter - self.input_data = input_data - self.input_data_idxs = [ *range(len(self.input_data)) ] + super().__init__('Converter', ConvertSubprocessor.Cli, 86400 if CONVERTER_DEBUG else 60, io_loop_sleep_time=0.001, initialize_subprocesses_in_serial=False)# if debug == True else 60) + + self.is_interactive = is_interactive + self.converter_config = converter_config + + #dummy predict and sleep, tensorflow caching kernels. If remove it, sometime conversion speed can be x2 slower + self.converter_config.predictor_func (dummy_predict=True) + time.sleep(2) + + self.predictor_func_host, self.predictor_func = SubprocessFunctionCaller.make_pair(self.converter_config.predictor_func) + self.converter_config.predictor_func = None + + self.dcscn = None + def DCSCN_upscale(*args, **kwargs): + if self.dcscn is None: + self.dcscn = imagelib.DCSCN() + return self.dcscn.upscale(*args, **kwargs) + + self.dcscn_host, self.dcscn_upscale_func = SubprocessFunctionCaller.make_pair(DCSCN_upscale) + + self.frames = frames self.output_path = output_path - self.alignments = alignments - self.debug = debug + self.prefetch_frame_count = self.process_count = min(6,multiprocessing.cpu_count()) - self.files_processed = 0 - self.faces_processed = 0 + self.frames_idxs = [ *range(len(self.frames)) ] + self.frames_done_idxs = [] + + for i in range( len(self.frames) ): + frame = self.frames[i] + frame.idx = i + frame.output_filename = self.output_path / ('%.5d.png' % i) + + frames[0].cfg = self.converter_config.copy() #override def process_info_generator(self): - r = [0] if self.debug else range( min(6,multiprocessing.cpu_count()) ) + r = [0] if CONVERTER_DEBUG else range(self.process_count) for i in r: yield 'CPU%d' % (i), {}, {'device_idx': i, 'device_name': 'CPU%d' % (i), - 'converter' : self.converter, - 'input_data' : self.input_data, - 'output_dir' : str(self.output_path), - 'alignments' : self.alignments, - 'debug': self.debug, - 'stdin_fd': sys.stdin.fileno() if self.debug else None + 'predictor_func': self.predictor_func, + 'dcscn_upscale_func': self.dcscn_upscale_func, + 'stdin_fd': sys.stdin.fileno() if CONVERTER_DEBUG else None } #overridable optional def on_clients_initialized(self): - if self.debug: - io.named_window ("Debug convert") + io.progress_bar ("Converting", len (self.frames_idxs) ) - io.progress_bar ("Converting", len (self.input_data_idxs) ) + self.process_remain_frames = not self.is_interactive + self.is_interactive_quitting = not self.is_interactive + + if self.is_interactive: + help_images = { + ConverterConfig.TYPE_MASKED : cv2_imread ( str(Path(__file__).parent / 'gfx' / 'help_converter_masked.jpg') ), + ConverterConfig.TYPE_FACE_AVATAR : cv2_imread ( str(Path(__file__).parent / 'gfx' / 'help_converter_face_avatar.jpg') ), + } + + self.main_screen = Screen(initial_scale_to_width=1368, image=None, waiting_icon=True) + self.help_screen = Screen(initial_scale_to_height=768, image=help_images[self.converter_config.type], waiting_icon=False) + self.screen_manager = ScreenManager( "Converter", [self.main_screen, self.help_screen], capture_keys=True ) + self.screen_manager.set_current (self.help_screen) + self.screen_manager.show_current() #overridable optional def on_clients_finalized(self): + + if self.is_interactive: + self.screen_manager.finalize() + io.progress_bar_close() - if self.debug: - io.destroy_all_windows() + cfg_change_keys = ['`','1', '2', '3', '4', '5', '6', '7', '8', '9', + 'q', 'a', 'w', 's', 'e', 'd', 'r', 'f', 't', 'g','y','h', + 'z', 'x', 'c', 'v', 'b' ] + #override + def on_tick(self): + self.predictor_func_host.process_messages() + self.dcscn_host.process_messages() + + go_prev_frame = False + go_next_frame = self.process_remain_frames + go_next_frame_overriding_cfg = False + + cur_frame = None + if len(self.frames_idxs) != 0: + cur_frame = self.frames[self.frames_idxs[0]] + + if self.is_interactive: + self.main_screen.set_waiting_icon(False) + + if not self.is_interactive_quitting and not self.process_remain_frames: + if cur_frame is not None: + if not cur_frame.is_shown: + if cur_frame.is_done: + cur_frame.is_shown = True + io.log_info (cur_frame.cfg) + + if cur_frame.image is None: + cur_frame.image = cv2_imread ( cur_frame.output_filename) + if cur_frame.image is None: + cur_frame.is_done = False #unable to read? recompute then + cur_frame.is_shown = False + self.main_screen.set_image(cur_frame.image) + else: + self.main_screen.set_waiting_icon(True) + + else: + self.main_screen.set_image(None) + else: + self.main_screen.set_image(None) + self.main_screen.set_waiting_icon(True) + + self.screen_manager.show_current() + + key_events = self.screen_manager.get_key_events() + key, chr_key, ctrl_pressed, alt_pressed, shift_pressed = key_events[-1] if len(key_events) > 0 else (0,0,False,False,False) + + if key == 9: #tab + self.screen_manager.switch_screens() + else: + if key == 27: #esc + self.is_interactive_quitting = True + elif self.screen_manager.get_current() is self.main_screen: + if chr_key in self.cfg_change_keys: + self.process_remain_frames = False + + if cur_frame is not None: + cfg = cur_frame.cfg + prev_cfg = cfg.copy() + + if cfg.type == ConverterConfig.TYPE_MASKED: + if chr_key == '`': + cfg.set_mode(0) + elif key >= ord('1') and key <= ord('9'): + cfg.set_mode( key - ord('0') ) + elif chr_key == 'q': + cfg.add_hist_match_threshold(1 if not shift_pressed else 5) + elif chr_key == 'a': + cfg.add_hist_match_threshold(-1 if not shift_pressed else -5) + elif chr_key == 'w': + cfg.add_erode_mask_modifier(1 if not shift_pressed else 5) + elif chr_key == 's': + cfg.add_erode_mask_modifier(-1 if not shift_pressed else -5) + elif chr_key == 'e': + cfg.add_blur_mask_modifier(1 if not shift_pressed else 5) + elif chr_key == 'd': + cfg.add_blur_mask_modifier(-1 if not shift_pressed else -5) + elif chr_key == 'r': + cfg.add_motion_blur_power(1 if not shift_pressed else 5) + elif chr_key == 'f': + cfg.add_motion_blur_power(-1 if not shift_pressed else -5) + elif chr_key == 't': + cfg.add_color_degrade_power(1 if not shift_pressed else 5) + elif chr_key == 'g': + cfg.add_color_degrade_power(-1 if not shift_pressed else -5) + elif chr_key == 'y': + cfg.add_output_face_scale(1 if not shift_pressed else 5) + elif chr_key == 'h': + cfg.add_output_face_scale(-1 if not shift_pressed else -5) + + elif chr_key == 'z': + cfg.toggle_masked_hist_match() + elif chr_key == 'x': + cfg.toggle_mask_mode() + elif chr_key == 'c': + cfg.toggle_color_transfer_mode() + elif chr_key == 'v': + cfg.toggle_super_resolution() + elif chr_key == 'b': + cfg.toggle_export_mask_alpha() + else: + if chr_key == 's': + cfg.toggle_add_source_image() + + if prev_cfg != cfg: + io.log_info (cfg) + cur_frame.is_done = False + cur_frame.is_shown = False + else: + if chr_key == ',': + self.process_remain_frames = False + go_prev_frame = True + elif chr_key == '.' or chr_key == '/': + self.process_remain_frames = False + go_next_frame = True + go_next_frame_overriding_cfg = chr_key == '/' + elif chr_key == '\r' or chr_key == '\n': + self.process_remain_frames = not self.process_remain_frames + elif chr_key == '-': + self.screen_manager.get_current().diff_scale(-0.1) + elif chr_key == '=': + self.screen_manager.get_current().diff_scale(0.1) + + + if go_prev_frame: + if cur_frame is not None and cur_frame.is_done: + cur_frame.image = None + + if len(self.frames_done_idxs) > 0: + prev_frame = self.frames[self.frames_done_idxs.pop()] + self.frames_idxs.insert(0, prev_frame.idx) + prev_frame.is_shown = False + io.progress_bar_inc(-1) + + elif go_next_frame: + if cur_frame is not None and cur_frame.is_done: + cur_frame.image = None + cur_frame.is_shown = True + self.frames_done_idxs.append(cur_frame.idx) + self.frames_idxs.pop(0) + io.progress_bar_inc(1) + + if len(self.frames_idxs) != 0: + next_frame = self.frames[ self.frames_idxs[0] ] + + if go_next_frame_overriding_cfg: + f = self.frames + for i in range( next_frame.idx, len(self.frames) ): + f[i].cfg = None + f[i].is_shown = False + + if next_frame.cfg is None or next_frame.is_shown == False: #next frame is never shown or override current cfg to next frames and the prefetches + for i in range( min(len(self.frames_idxs), self.prefetch_frame_count) ): + frame = self.frames[ self.frames_idxs[i] ] + + if frame.cfg is None or frame.cfg != cur_frame.cfg: + frame.cfg = cur_frame.cfg.copy() + frame.is_done = False #initiate solve again + + + next_frame.is_shown = False + + if len(self.frames_idxs) == 0: + self.process_remain_frames = False + + return (self.is_interactive and self.is_interactive_quitting) or \ + (not self.is_interactive and self.process_remain_frames == False) + + + #override + def on_data_return (self, host_dict, pf): + frame = self.frames[pf.idx] + frame.is_done = False + frame.is_processing = False + + #override + def on_result (self, host_dict, pf_sent, pf_result): + frame = self.frames[pf_result.idx] + frame.is_processing = False + if frame.cfg == pf_result.cfg: + frame.is_done = True + frame.image = pf_result.image #override def get_data(self, host_dict): - if len (self.input_data_idxs) > 0: - idx = self.input_data_idxs.pop(0) - return (idx, ) + if self.is_interactive and self.is_interactive_quitting: + return None + + for i in range ( min(len(self.frames_idxs), self.prefetch_frame_count) ): + frame = self.frames[ self.frames_idxs[i] ] + + if not frame.is_done and not frame.is_processing and frame.cfg is not None: + frame.is_processing = True + return ConvertSubprocessor.ProcessingFrame(idx=frame.idx, + cfg=frame.cfg.copy(), + prev_temporal_frame_infos=frame.prev_temporal_frame_infos, + frame_info=frame.frame_info, + next_temporal_frame_infos=frame.next_temporal_frame_infos, + output_filename=frame.output_filename, + need_return_image=True ) + return None - #override - def on_data_return (self, host_dict, data): - idx, = data - self.input_data_idxs.insert(0, idx) - - #override - def on_result (self, host_dict, data, result): - if result[0] == 0: - self.files_processed += result[0] - self.faces_processed += result[1] - elif result[0] == 1: - for img in result[1]: - io.show_image ('Debug convert', (img*255).astype(np.uint8) ) - io.wait_any_key() - io.progress_bar_inc(1) - - #override - def on_tick(self): - self.converter.on_host_tick() - #override def get_result(self): - return self.files_processed, self.faces_processed + return 0 def main (args, device_args): io.log_info ("Running converter.\r\n") aligned_dir = args.get('aligned_dir', None) avaperator_aligned_dir = args.get('avaperator_aligned_dir', None) - + try: input_path = Path(args['input_dir']) output_path = Path(args['output_dir']) @@ -244,14 +469,19 @@ def main (args, device_args): io.log_err('Model directory not found. Please ensure it exists.') return + is_interactive = io.input_bool ("Use interactive converter? (y/n skip:y) : ", True) if not io.is_colab() else False + import models model = models.import_model( args['model_name'] )(model_path, device_args=device_args) - converter = model.get_converter() + + cfg = model.get_ConverterConfig() + + if not is_interactive: + cfg.ask_settings() input_path_image_paths = Path_utils.get_image_paths(input_path) - - alignments = None - if converter.type == Converter.TYPE_FACE: + + if cfg.type == ConverterConfig.TYPE_MASKED: if aligned_dir is None: io.log_err('Aligned directory not found. Please ensure it exists.') return @@ -262,7 +492,7 @@ def main (args, device_args): return alignments = {} - + multiple_faces_detected = False aligned_path_image_paths = Path_utils.get_image_paths(aligned_path) for filepath in io.progress_bar_generator(aligned_path_image_paths, "Collecting alignments"): filepath = Path(filepath) @@ -282,13 +512,52 @@ def main (args, device_args): if source_filename_stem not in alignments.keys(): alignments[ source_filename_stem ] = [] - alignments[ source_filename_stem ].append (dflimg.get_source_landmarks()) - #avatar_alignments += [ ( str(filepath), dflimg.get_source_landmarks(), dflimg.get_source_filename() ) ] - - input_data = [ (p,) for p in input_path_image_paths ] - elif converter.type == Converter.TYPE_FACE_AVATAR: - - input_data = [] + alignments_ar = alignments[ source_filename_stem ] + alignments_ar.append (dflimg.get_source_landmarks()) + if len(alignments_ar) > 1: + multiple_faces_detected = True + + if multiple_faces_detected: + io.log_info ("Warning: multiple faces detected. Strongly recommended to process them separately.") + + frames = [ ConvertSubprocessor.Frame( frame_info=FrameInfo(filename=p, landmarks_list=alignments.get(Path(p).stem, None))) for p in input_path_image_paths ] + + if multiple_faces_detected: + io.log_info ("Warning: multiple faces detected. Motion blur will not be used.") + else: + s = 256 + local_pts = [ (s//2-1, s//2-1), (s//2-1,0) ] #center+up + frames_len = len(frames) + for i in io.progress_bar_generator( range(len(frames)) , "Computing motion vectors"): + fi_prev = frames[max(0, i-1)].frame_info + fi = frames[i].frame_info + fi_next = frames[min(i+1, frames_len-1)].frame_info + if len(fi_prev.landmarks_list) == 0 or \ + len(fi.landmarks_list) == 0 or \ + len(fi_next.landmarks_list) == 0: + continue + + mat_prev = LandmarksProcessor.get_transform_mat ( fi_prev.landmarks_list[0], s, face_type=FaceType.FULL) + mat = LandmarksProcessor.get_transform_mat ( fi.landmarks_list[0] , s, face_type=FaceType.FULL) + mat_next = LandmarksProcessor.get_transform_mat ( fi_next.landmarks_list[0], s, face_type=FaceType.FULL) + + pts_prev = LandmarksProcessor.transform_points (local_pts, mat_prev, True) + pts = LandmarksProcessor.transform_points (local_pts, mat, True) + pts_next = LandmarksProcessor.transform_points (local_pts, mat_next, True) + + prev_vector = pts[0]-pts_prev[0] + next_vector = pts_next[0]-pts[0] + + motion_vector = pts_next[0] - pts_prev[0] + fi.motion_power = npla.norm(motion_vector) + + motion_vector = motion_vector / fi.motion_power if fi.motion_power != 0 else np.array([0,0],dtype=np.float32) + + fi.motion_deg = -math.atan2(motion_vector[1],motion_vector[0])*180 / math.pi + + + elif cfg.type == ConverterConfig.TYPE_FACE_AVATAR: + filesdata = [] for filepath in io.progress_bar_generator(input_path_image_paths, "Collecting info"): filepath = Path(filepath) @@ -302,18 +571,36 @@ def main (args, device_args): if dflimg is None: io.log_err ("%s is not a dfl image file" % (filepath.name) ) continue - input_data += [ ( str(filepath), dflimg.get_landmarks(), dflimg.get_source_filename() ) ] - - input_data = sorted(input_data, key=operator.itemgetter(2)) + filesdata += [ ( FrameInfo(filename=str(filepath), landmarks_list=[dflimg.get_landmarks()] ), dflimg.get_source_filename() ) ] + + filesdata = sorted(filesdata, key=operator.itemgetter(1)) #sort by filename + frames = [] + filesdata_len = len(filesdata) + for i in range(len(filesdata)): + frame_info = filesdata[i][0] + + prev_temporal_frame_infos = [] + next_temporal_frame_infos = [] + + for t in range (cfg.temporal_face_count): + prev_frame_info = filesdata[ max(i -t, 0) ][0] + next_frame_info = filesdata[ min(i +t, filesdata_len-1 )][0] + + prev_temporal_frame_infos.insert (0, prev_frame_info ) + next_temporal_frame_infos.append ( next_frame_info ) + + frames.append ( ConvertSubprocessor.Frame(prev_temporal_frame_infos=prev_temporal_frame_infos, + frame_info=frame_info, + next_temporal_frame_infos=next_temporal_frame_infos) ) + + if len(frames) == 0: + io.log_info ("No frames to convert in input_dir.") else: - input_data = [ (p,) for p in input_path_image_paths ] - - files_processed, faces_processed = ConvertSubprocessor ( - converter = converter, - input_data = input_data, - output_path = output_path, - alignments = alignments, - debug = args.get('debug',False) + ConvertSubprocessor ( + is_interactive = is_interactive, + converter_config = cfg, + frames = frames, + output_path = output_path, ).run() model.finalize() @@ -322,29 +609,6 @@ def main (args, device_args): print ( 'Error: %s' % (str(e))) traceback.print_exc() -''' -if model_name == 'AVATAR': - output_path_image_paths = Path_utils.get_image_paths(output_path) - - last_ok_frame = -1 - for filename in output_path_image_paths: - filename_path = Path(filename) - stem = Path(filename).stem - try: - frame = int(stem) - except: - raise Exception ('Aligned avatars must be created from indexed sequence files.') - - if frame-last_ok_frame > 1: - start = last_ok_frame + 1 - end = frame - 1 - - print ("Filling gaps: [%d...%d]" % (start, end) ) - for i in range (start, end+1): - shutil.copy ( str(filename), str( output_path / ('%.5d%s' % (i, filename_path.suffix )) ) ) - - last_ok_frame = frame -''' #interpolate landmarks #from facelib import LandmarksProcessor #from facelib import FaceType @@ -376,29 +640,3 @@ if model_name == 'AVATAR': # 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) - -""" - 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) ) - -""" \ No newline at end of file diff --git a/mainscripts/ConverterScreen/ConverterScreen.py b/mainscripts/ConverterScreen/ConverterScreen.py new file mode 100644 index 0000000..78ed45d --- /dev/null +++ b/mainscripts/ConverterScreen/ConverterScreen.py @@ -0,0 +1,138 @@ +import math +from pathlib import Path + +import numpy as np + +import imagelib +from interact import interact as io +from utils.cv2_utils import * +from utils.os_utils import get_screen_size + + +class ScreenAssets(object): + waiting_icon_image = cv2_imread ( str(Path(__file__).parent / 'gfx' / 'sand_clock_64.png') ) + + @staticmethod + def build_checkerboard_a( sh, size=5): + h,w = sh[0], sh[1] + tile = np.array([[0,1],[1,0]]).repeat(size, axis=0).repeat(size, axis=1) + grid = np.tile(tile,(int(math.ceil((h+0.0)/(2*size))),int(math.ceil((w+0.0)/(2*size))))) + return grid[:h,:w,None] + +class Screen(object): + def __init__(self, initial_scale_to_width=0, initial_scale_to_height=0, image=None, waiting_icon=False, **kwargs): + self.initial_scale_to_width = initial_scale_to_width + self.initial_scale_to_height = initial_scale_to_height + self.image = image + self.waiting_icon = waiting_icon + + self.state = -1 + self.scale = 1 + self.force_update = True + self.is_first_appear = True + + self.last_screen_shape = (480,640,3) + self.checkerboard_image = None + self.set_image (image) + self.scrn_manager = None + + def set_waiting_icon(self, b): + self.waiting_icon = b + + def set_image(self, img): + if not img is self.image: + self.force_update = True + + self.image = img + + if self.image is not None: + self.last_screen_shape = self.image.shape + + if self.initial_scale_to_width != 0: + if self.last_screen_shape[1] > self.initial_scale_to_width: + self.scale = self.initial_scale_to_width / self.last_screen_shape[1] + self.force_update = True + self.initial_scale_to_width = 0 + + elif self.initial_scale_to_height != 0: + if self.last_screen_shape[0] > self.initial_scale_to_height: + self.scale = self.initial_scale_to_height / self.last_screen_shape[0] + self.force_update = True + self.initial_scale_to_height = 0 + + + def diff_scale(self, diff): + self.scale = np.clip (self.scale + diff, 0.1, 4.0) + self.force_update = True + + def show(self, force=False): + new_state = 0 | self.waiting_icon + + if self.state != new_state or self.force_update or force: + self.state = new_state + self.force_update = False + + if self.image is None: + screen = np.zeros ( self.last_screen_shape, dtype=np.uint8 ) + else: + screen = self.image.copy() + + if self.waiting_icon: + imagelib.overlay_alpha_image (screen, ScreenAssets.waiting_icon_image, (0,0) ) + + h,w,c = screen.shape + if self.scale != 1.0: + screen = cv2.resize ( screen, ( int(w*self.scale), int(h*self.scale) ) ) + + if c == 4: + if self.checkerboard_image is None or self.checkerboard_image.shape[0:2] != screen.shape[0:2]: + self.checkerboard_image = ScreenAssets.build_checkerboard_a(screen.shape) + + screen = screen[...,0:3]*0.75 + 64*self.checkerboard_image*(1- (screen[...,3:4].astype(np.float32)/255.0) ) + screen = screen.astype(np.uint8) + + io.show_image(self.scrn_manager.wnd_name, screen) + + if self.is_first_appear: + self.is_first_appear = False + #center window + desktop_w, desktop_h = get_screen_size() + h,w,c = screen.shape + cv2.moveWindow(self.scrn_manager.wnd_name, max(0,(desktop_w-w) // 2), max(0, (desktop_h-h) // 2) ) + + io.process_messages(0.0001) + +class ScreenManager(object): + def __init__(self, window_name="ScreenManager", screens=None, capture_keys=False ): + self.screens = screens or [] + self.current_screen_id = 0 + + if self.screens is not None: + for screen in self.screens: + screen.scrn_manager = self + + self.wnd_name = window_name + io.named_window(self.wnd_name) + + + if capture_keys: + io.capture_keys(self.wnd_name) + + def finalize(self): + io.destroy_all_windows() + + def get_key_events(self): + return io.get_key_events(self.wnd_name) + + def switch_screens(self): + self.current_screen_id = (self.current_screen_id + 1) % len(self.screens) + self.screens[self.current_screen_id].show(force=True) + + def show_current(self): + self.screens[self.current_screen_id].show() + + def get_current(self): + return self.screens[self.current_screen_id] + + def set_current(self, screen): + self.current_screen_id = self.screens.index(screen) diff --git a/mainscripts/ConverterScreen/__init__.py b/mainscripts/ConverterScreen/__init__.py new file mode 100644 index 0000000..5103fc4 --- /dev/null +++ b/mainscripts/ConverterScreen/__init__.py @@ -0,0 +1 @@ +from .ConverterScreen import Screen, ScreenManager \ No newline at end of file diff --git a/mainscripts/ConverterScreen/gfx/sand_clock_64.png b/mainscripts/ConverterScreen/gfx/sand_clock_64.png new file mode 100644 index 0000000..c8d3486 Binary files /dev/null and b/mainscripts/ConverterScreen/gfx/sand_clock_64.png differ diff --git a/mainscripts/gfx/help_converter_face_avatar.jpg b/mainscripts/gfx/help_converter_face_avatar.jpg new file mode 100644 index 0000000..896f641 Binary files /dev/null and b/mainscripts/gfx/help_converter_face_avatar.jpg differ diff --git a/mainscripts/gfx/help_converter_face_avatar_source.psd b/mainscripts/gfx/help_converter_face_avatar_source.psd new file mode 100644 index 0000000..cd03540 Binary files /dev/null and b/mainscripts/gfx/help_converter_face_avatar_source.psd differ diff --git a/mainscripts/gfx/help_converter_masked.jpg b/mainscripts/gfx/help_converter_masked.jpg new file mode 100644 index 0000000..d66c9cd Binary files /dev/null and b/mainscripts/gfx/help_converter_masked.jpg differ diff --git a/mainscripts/gfx/help_converter_masked_source.psd b/mainscripts/gfx/help_converter_masked_source.psd new file mode 100644 index 0000000..5fc1736 Binary files /dev/null and b/mainscripts/gfx/help_converter_masked_source.psd differ diff --git a/models/ModelBase.py b/models/ModelBase.py index e7ad843..7acb2de 100644 --- a/models/ModelBase.py +++ b/models/ModelBase.py @@ -25,11 +25,11 @@ class ModelBase(object): def __init__(self, model_path, training_data_src_path=None, training_data_dst_path=None, pretraining_data_path=None, debug = False, device_args = None, ask_enable_autobackup=True, - ask_write_preview_history=True, - ask_target_iter=True, - ask_batch_size=True, + ask_write_preview_history=True, + ask_target_iter=True, + ask_batch_size=True, ask_sort_by_yaw=True, - ask_random_flip=True, + ask_random_flip=True, ask_src_scale_mod=True): device_args['force_gpu_idx'] = device_args.get('force_gpu_idx',-1) @@ -55,7 +55,7 @@ class ModelBase(object): self.training_data_src_path = training_data_src_path self.training_data_dst_path = training_data_dst_path self.pretraining_data_path = pretraining_data_path - + self.src_images_paths = None self.dst_images_paths = None self.src_yaw_images_paths = None @@ -106,7 +106,7 @@ class ModelBase(object): choose_preview_history = io.input_bool("Randomly choose new image for preview history? (y/n ?:help skip:%s) : " % (yn_str[False]), False, help_message="Preview image history will stay stuck with old faces if you reuse the same model on different celebs. Choose no unless you are changing src/dst to a new person") else: choose_preview_history = False - + if ask_target_iter: if (self.iter == 0 or ask_override): self.options['target_iter'] = max(0, io.input_int("Target iteration (skip:unlimited/default) : ", 0)) @@ -121,7 +121,7 @@ class ModelBase(object): else: self.options['batch_size'] = self.options.get('batch_size', 0) - if ask_sort_by_yaw: + if ask_sort_by_yaw: if (self.iter == 0 or ask_override): default_sort_by_yaw = self.options.get('sort_by_yaw', False) self.options['sort_by_yaw'] = io.input_bool("Feed faces to network sorted by yaw? (y/n ?:help skip:%s) : " % (yn_str[default_sort_by_yaw]), default_sort_by_yaw, help_message="NN will not learn src face directions that don't match dst face directions. Do not use if the dst face has hair that covers the jaw." ) @@ -139,7 +139,7 @@ class ModelBase(object): self.options['src_scale_mod'] = np.clip( io.input_int("Src face scale modifier % ( -30...30, ?:help skip:0) : ", 0, help_message="If src face shape is wider than dst, try to decrease this value to get a better result."), -30, 30) else: self.options['src_scale_mod'] = self.options.get('src_scale_mod', 0) - + self.autobackup = self.options.get('autobackup', False) if not self.autobackup and 'autobackup' in self.options: self.options.pop('autobackup') @@ -180,10 +180,10 @@ class ModelBase(object): else: self.preview_history_path = self.model_path / ( '%d_%s_history' % (self.device_args['force_gpu_idx'], self.get_model_name()) ) self.autobackups_path = self.model_path / ( '%d_%s_autobackups' % (self.device_args['force_gpu_idx'], self.get_model_name()) ) - + if self.autobackup: self.autobackup_current_hour = time.localtime().tm_hour - + if not self.autobackups_path.exists(): self.autobackups_path.mkdir(exist_ok=True) @@ -202,7 +202,7 @@ class ModelBase(object): if not isinstance(generator, SampleGeneratorBase): raise ValueError('training data generator is not subclass of SampleGeneratorBase') - if self.sample_for_preview is None or choose_preview_history: + if self.sample_for_preview is None or choose_preview_history: if choose_preview_history and io.is_support_windows(): wnd_name = "[p] - next. [enter] - confirm." io.named_window(wnd_name) @@ -221,25 +221,25 @@ class ModelBase(object): break elif key == ord('p'): break - + try: io.process_messages(0.1) except KeyboardInterrupt: choosed = True - + io.destroy_window(wnd_name) - else: - self.sample_for_preview = self.generate_next_sample() + else: + self.sample_for_preview = self.generate_next_sample() self.last_sample = self.sample_for_preview - + ###Generate text summary of model hyperparameters #Find the longest key name and value string. Used as column widths. width_name = max([len(k) for k in self.options.keys()] + [17]) + 1 # Single space buffer to left edge. Minimum of 17, the length of the longest static string used "Current iteration" width_value = max([len(str(x)) for x in self.options.values()] + [len(str(self.iter)), len(self.get_model_name())]) + 1 # Single space buffer to right edge if not self.device_config.cpu_only: #Check length of GPU names - width_value = max([len(nnlib.device.getDeviceName(idx))+1 for idx in self.device_config.gpu_idxs] + [width_value]) + width_value = max([len(nnlib.device.getDeviceName(idx))+1 for idx in self.device_config.gpu_idxs] + [width_value]) width_total = width_name + width_value + 2 #Plus 2 for ": " - + model_summary_text = [] model_summary_text += [f'=={" Model Summary ":=^{width_total}}=='] # Model/status summary model_summary_text += [f'=={" "*width_total}=='] @@ -247,13 +247,13 @@ class ModelBase(object): model_summary_text += [f'=={" "*width_total}=='] model_summary_text += [f'=={"Current iteration": >{width_name}}: {str(self.iter): <{width_value}}=='] # Iter model_summary_text += [f'=={" "*width_total}=='] - + model_summary_text += [f'=={" Model Options ":-^{width_total}}=='] # Model options model_summary_text += [f'=={" "*width_total}=='] for key in self.options.keys(): model_summary_text += [f'=={key: >{width_name}}: {str(self.options[key]): <{width_value}}=='] # self.options key/value pairs model_summary_text += [f'=={" "*width_total}=='] - + model_summary_text += [f'=={" Running On ":-^{width_total}}=='] # Training hardware info model_summary_text += [f'=={" "*width_total}=='] if self.device_config.multi_gpu: @@ -266,10 +266,10 @@ class ModelBase(object): model_summary_text += [f'=={"Device index": >{width_name}}: {idx: <{width_value}}=='] # GPU hardware device index model_summary_text += [f'=={"Name": >{width_name}}: {nnlib.device.getDeviceName(idx): <{width_value}}=='] # GPU name vram_str = f'{nnlib.device.getDeviceVRAMTotalGb(idx):.2f}GB' # GPU VRAM - Formated as #.## (or ##.##) - model_summary_text += [f'=={"VRAM": >{width_name}}: {vram_str: <{width_value}}=='] + model_summary_text += [f'=={"VRAM": >{width_name}}: {vram_str: <{width_value}}=='] model_summary_text += [f'=={" "*width_total}=='] model_summary_text += [f'=={"="*width_total}=='] - + if not self.device_config.cpu_only and self.device_config.gpu_vram_gb[0] <= 2: # Low VRAM warning model_summary_text += ["/!\\"] model_summary_text += ["/!\\ WARNING:"] @@ -277,7 +277,7 @@ class ModelBase(object): model_summary_text += ["/!\\ If training does not start, close all programs and try again."] model_summary_text += ["/!\\ Also you can disable Windows Aero Desktop to increase available VRAM."] model_summary_text += ["/!\\"] - + model_summary_text = "\n".join (model_summary_text) self.model_summary_text = model_summary_text io.log_info(model_summary_text) @@ -323,6 +323,11 @@ class ModelBase(object): def get_model_filename_list(self): return [] + #overridable + def get_ConverterConfig(self): + #return ConverterConfig() for the model + raise NotImplementedError + #overridable def get_converter(self): raise NotImplementedError @@ -372,9 +377,9 @@ class ModelBase(object): } self.model_data_path.write_bytes( pickle.dumps(model_data) ) - bckp_filename_list = [ self.get_strpath_storage_for_file(filename) for _, filename in self.get_model_filename_list() ] - bckp_filename_list += [ str(summary_path), str(self.model_data_path) ] - + bckp_filename_list = [ self.get_strpath_storage_for_file(filename) for _, filename in self.get_model_filename_list() ] + bckp_filename_list += [ str(summary_path), str(self.model_data_path) ] + if self.autobackup: current_hour = time.localtime().tm_hour if self.autobackup_current_hour != current_hour: @@ -383,20 +388,20 @@ class ModelBase(object): for i in range(15,0,-1): idx_str = '%.2d' % i next_idx_str = '%.2d' % (i+1) - + idx_backup_path = self.autobackups_path / idx_str next_idx_packup_path = self.autobackups_path / next_idx_str - + if idx_backup_path.exists(): - if i == 15: + if i == 15: Path_utils.delete_all_files(idx_backup_path) else: next_idx_packup_path.mkdir(exist_ok=True) Path_utils.move_all_files (idx_backup_path, next_idx_packup_path) - + if i == 1: - idx_backup_path.mkdir(exist_ok=True) - for filename in bckp_filename_list: + idx_backup_path.mkdir(exist_ok=True) + for filename in bckp_filename_list: shutil.copy ( str(filename), str(idx_backup_path / Path(filename).name) ) previews = self.get_previews() @@ -440,7 +445,7 @@ class ModelBase(object): model.save_weights( filename + '.tmp' ) rename_list = model_filename_list - + """ #unused , optimizer_filename_list=[] @@ -464,7 +469,7 @@ class ModelBase(object): except Exception as e: print ("Unable to save ", opt_filename) """ - + for _, filename in rename_list: filename = self.get_strpath_storage_for_file(filename) source_filename = Path(filename+'.tmp') @@ -473,7 +478,7 @@ class ModelBase(object): if target_filename.exists(): target_filename.unlink() source_filename.rename ( str(target_filename) ) - + def debug_one_iter(self): images = [] for generator in self.generator_list: @@ -579,8 +584,8 @@ class ModelBase(object): lh_height = 100 lh_img = np.ones ( (lh_height,w,c) ) * 0.1 - - if len(loss_history) != 0: + + if len(loss_history) != 0: loss_count = len(loss_history[0]) lh_len = len(loss_history) diff --git a/models/Model_AVATAR/Model.py b/models/Model_AVATAR/Model.py new file mode 100644 index 0000000..00c8730 --- /dev/null +++ b/models/Model_AVATAR/Model.py @@ -0,0 +1,740 @@ +from functools import partial + +import cv2 +import numpy as np + +from facelib import FaceType +from interact import interact as io +from mathlib import get_power_of_two +from models import ModelBase +from nnlib import nnlib +from samplelib import * + +from facelib import PoseEstimator + +class AVATARModel(ModelBase): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs, + ask_sort_by_yaw=False, + ask_random_flip=False, + ask_src_scale_mod=False) + + #override + def onInitializeOptions(self, is_first_run, ask_override): + if is_first_run: + avatar_type = io.input_int("Avatar type ( 0:source, 1:full_face, 2:head ?:help skip:0) : ", 0, [0,1,2], + help_message="Training target for the model. Source is direct untouched images. Full_face or head are centered nose unaligned faces.") + avatar_type = {0:'source', + 1:'full_face', + 2:'head'}[avatar_type] + + self.options['avatar_type'] = avatar_type + else: + self.options['avatar_type'] = self.options.get('avatar_type', 'source') + + if is_first_run or ask_override: + def_stage = self.options.get('stage', 0) + self.options['stage'] = io.input_int("Stage (0, 1, 2 ?:help skip:%d) : " % def_stage, def_stage, [0,1,2], help_message="Train first stage, then second. Tune batch size to maximum possible for both stages.") + else: + self.options['stage'] = self.options.get('stage', 0) + + + + + #override + def onInitialize(self, batch_size=-1, **in_options): + exec(nnlib.code_import_all, locals(), globals()) + self.set_vram_batch_requirements({6:4}) + AVATARModel.initialize_nn_functions() + + resolution = self.resolution = 224 + avatar_type = self.options['avatar_type'] + stage = self.stage = self.options['stage'] + df_res = self.df_res = 128 + df_bgr_shape = (df_res, df_res, 3) + df_mask_shape = (df_res, df_res, 1) + res_bgr_shape = (resolution, resolution, 3) + res_bgr_t_shape = (resolution, resolution, 9) + + self.enc = modelify(AVATARModel.EncFlow())( [Input(df_bgr_shape),] ) + + self.decA64 = modelify(AVATARModel.DecFlow()) ( [ Input(K.int_shape(self.enc.outputs[0])[1:]) ] ) + self.decB64 = modelify(AVATARModel.DecFlow()) ( [ Input(K.int_shape(self.enc.outputs[0])[1:]) ] ) + self.D = modelify(AVATARModel.Discriminator() ) (Input(df_bgr_shape)) + self.C = modelify(AVATARModel.ResNet (9, use_batch_norm=False, n_blocks=6, ngf=128, use_dropout=False))( Input(res_bgr_t_shape)) + #self.CD = modelify(AVATARModel.CDiscriminator() ) (Input(res_bgr_t_shape)) + + if self.is_first_run(): + conv_weights_list = [] + for model in [self.enc, self.decA64, self.decB64, self.C, self.D, self.CD]: + for layer in model.layers: + if type(layer) == keras.layers.Conv2D: + conv_weights_list += [layer.weights[0]] #Conv2D kernel_weights + CAInitializerMP ( conv_weights_list ) + + if not self.is_first_run(): + self.load_weights_safe( self.get_model_filename_list() ) + + def DLoss(labels,logits): + return K.mean(K.binary_crossentropy(labels,logits)) + + warped_A64 = Input(df_bgr_shape) + real_A64 = Input(df_bgr_shape) + real_A64m = Input(df_mask_shape) + + real_B64_t0 = Input(df_bgr_shape) + real_B64_t1 = Input(df_bgr_shape) + real_B64_t2 = Input(df_bgr_shape) + + real_A64_t0 = Input(df_bgr_shape) + real_A64m_t0 = Input(df_mask_shape) + real_A_t0 = Input(res_bgr_shape) + real_A64_t1 = Input(df_bgr_shape) + real_A64m_t1 = Input(df_mask_shape) + real_A_t1 = Input(res_bgr_shape) + real_A64_t2 = Input(df_bgr_shape) + real_A64m_t2 = Input(df_mask_shape) + real_A_t2 = Input(res_bgr_shape) + + warped_B64 = Input(df_bgr_shape) + real_B64 = Input(df_bgr_shape) + real_B64m = Input(df_mask_shape) + + warped_A_code = self.enc (warped_A64) + warped_B_code = self.enc (warped_B64) + + rec_A64 = self.decA64(warped_A_code) + rec_B64 = self.decB64(warped_B_code) + rec_AB64 = self.decA64(warped_B_code) + + def Lambda_grey_mask (x,m): + return Lambda (lambda x: x[0]*m+(1-m)*0.5, output_shape= K.int_shape(x)[1:3] + (3,)) ([x, m]) + + def Lambda_gray_pad(x): + a = np.ones((resolution,resolution,3))*0.5 + pad = ( resolution - df_res ) // 2 + a[pad:-pad:,pad:-pad:,:] = 0 + + return Lambda ( lambda x: K.spatial_2d_padding(x, padding=((pad, pad), (pad, pad)) ) + K.constant(a, dtype=K.floatx() ), + output_shape=(resolution,resolution,3) ) (x) + + def Lambda_concat ( x ): + c = sum ( [ K.int_shape(l)[-1] for l in x ] ) + return Lambda ( lambda x: K.concatenate (x, axis=-1), output_shape=K.int_shape(x[0])[1:3] + (c,) ) (x) + + def Lambda_Cto3t(x): + return Lambda ( lambda x: x[...,0:3], output_shape= K.int_shape(x)[1:3] + (3,) ) (x), \ + Lambda ( lambda x: x[...,3:6], output_shape= K.int_shape(x)[1:3] + (3,) ) (x), \ + Lambda ( lambda x: x[...,6:9], output_shape= K.int_shape(x)[1:3] + (3,) ) (x) + + real_A64_d = self.D( Lambda_grey_mask(real_A64, real_A64m) ) + + real_A64_d_ones = K.ones_like(real_A64_d) + fake_A64_d = self.D(rec_AB64) + fake_A64_d_ones = K.ones_like(fake_A64_d) + fake_A64_d_zeros = K.zeros_like(fake_A64_d) + + rec_AB_t0 = Lambda_gray_pad( self.decA64 (self.enc (real_B64_t0)) ) + rec_AB_t1 = Lambda_gray_pad( self.decA64 (self.enc (real_B64_t1)) ) + rec_AB_t2 = Lambda_gray_pad( self.decA64 (self.enc (real_B64_t2)) ) + + C_in_A_t0 = Lambda_gray_pad( Lambda_grey_mask (real_A64_t0, real_A64m_t0) ) + C_in_A_t1 = Lambda_gray_pad( Lambda_grey_mask (real_A64_t1, real_A64m_t1) ) + C_in_A_t2 = Lambda_gray_pad( Lambda_grey_mask (real_A64_t2, real_A64m_t2) ) + + rec_C_A_t0, rec_C_A_t1, rec_C_A_t2 = Lambda_Cto3t ( self.C ( Lambda_concat ( [C_in_A_t0, C_in_A_t1, C_in_A_t2]) ) ) + rec_C_AB_t0, rec_C_AB_t1, rec_C_AB_t2 = Lambda_Cto3t( self.C ( Lambda_concat ( [rec_AB_t0, rec_AB_t1, rec_AB_t2]) ) ) + + #real_A_t012_d = self.CD ( K.concatenate ( [real_A_t0, real_A_t1,real_A_t2], axis=-1) ) + #real_A_t012_d_ones = K.ones_like(real_A_t012_d) + #rec_C_AB_t012_d = self.CD ( K.concatenate ( [rec_C_AB_t0,rec_C_AB_t1, rec_C_AB_t2], axis=-1) ) + #rec_C_AB_t012_d_ones = K.ones_like(rec_C_AB_t012_d) + #rec_C_AB_t012_d_zeros = K.zeros_like(rec_C_AB_t012_d) + + self.G64_view = K.function([warped_A64, warped_B64],[rec_A64, rec_B64, rec_AB64]) + self.G_view = K.function([real_A64_t0, real_A64m_t0, real_A64_t1, real_A64m_t1, real_A64_t2, real_A64m_t2, real_B64_t0, real_B64_t1, real_B64_t2], [rec_C_A_t0, rec_C_A_t1, rec_C_A_t2, rec_C_AB_t0, rec_C_AB_t1, rec_C_AB_t2]) + + if self.is_training_mode: + loss_AB64 = K.mean(10 * dssim(kernel_size=int(df_res/11.6),max_value=1.0) ( rec_A64, real_A64*real_A64m + (1-real_A64m)*0.5) ) + \ + K.mean(10 * dssim(kernel_size=int(df_res/11.6),max_value=1.0) ( rec_B64, real_B64*real_B64m + (1-real_B64m)*0.5) ) + 0.1*DLoss(fake_A64_d_ones, fake_A64_d ) + + weights_AB64 = self.enc.trainable_weights + self.decA64.trainable_weights + self.decB64.trainable_weights + + loss_C = K.mean( 10 * dssim(kernel_size=int(resolution/11.6),max_value=1.0) ( real_A_t0, rec_C_A_t0 ) ) + \ + K.mean( 10 * dssim(kernel_size=int(resolution/11.6),max_value=1.0) ( real_A_t1, rec_C_A_t1 ) ) + \ + K.mean( 10 * dssim(kernel_size=int(resolution/11.6),max_value=1.0) ( real_A_t2, rec_C_A_t2 ) ) + #0.1*DLoss(rec_C_AB_t012_d_ones, rec_C_AB_t012_d ) + + weights_C = self.C.trainable_weights + + loss_D = (DLoss(real_A64_d_ones, real_A64_d ) + \ + DLoss(fake_A64_d_zeros, fake_A64_d ) ) * 0.5 + + #loss_CD = ( DLoss(real_A_t012_d_ones, real_A_t012_d) + \ + # DLoss(rec_C_AB_t012_d_zeros, rec_C_AB_t012_d) ) * 0.5 + # + #weights_CD = self.CD.trainable_weights + + def opt(lr=5e-5): + return Adam(lr=lr, beta_1=0.5, beta_2=0.999, tf_cpu_mode=2 if 'tensorflow' in self.device_config.backend else 0 ) + + self.AB64_train = K.function ([warped_A64, real_A64, real_A64m, warped_B64, real_B64, real_B64m], [loss_AB64], opt().get_updates(loss_AB64, weights_AB64) ) + self.C_train = K.function ([real_A64_t0, real_A64m_t0, real_A_t0, + real_A64_t1, real_A64m_t1, real_A_t1, + real_A64_t2, real_A64m_t2, real_A_t2, + real_B64_t0, real_B64_t1, real_B64_t2],[ loss_C ], opt().get_updates(loss_C, weights_C) ) + + self.D_train = K.function ([warped_A64, real_A64, real_A64m, warped_B64, real_B64, real_B64m],[loss_D], opt().get_updates(loss_D, self.D.trainable_weights) ) + + + #self.CD_train = K.function ([real_A64_t0, real_A64m_t0, real_A_t0, + # real_A64_t1, real_A64m_t1, real_A_t1, + # real_A64_t2, real_A64m_t2, real_A_t2, + # real_B64_t0, real_B64_t1, real_B64_t2 ],[ loss_CD ], opt().get_updates(loss_CD, weights_CD) ) + + ########### + t = SampleProcessor.Types + + training_target = {'source' : t.NONE, + 'full_face' : t.FACE_TYPE_FULL_NO_ALIGN, + 'head' : t.FACE_TYPE_HEAD_NO_ALIGN}[avatar_type] + + generators = [ + SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, + sample_process_options=SampleProcessor.Options(random_flip=False), + output_sample_types=[ {'types': (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_BGR), 'resolution':df_res}, + {'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_BGR), 'resolution':df_res}, + {'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_M), 'resolution':df_res} + ] ), + SampleGeneratorFace(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, + sample_process_options=SampleProcessor.Options(random_flip=False), + output_sample_types=[ {'types': (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_BGR), 'resolution':df_res}, + {'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_BGR), 'resolution':df_res}, + {'types': (t.IMG_TRANSFORMED, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_M), 'resolution':df_res} + ] ), + + SampleGeneratorFaceTemporal(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, + temporal_image_count=3, + sample_process_options=SampleProcessor.Options(random_flip=False), + output_sample_types=[{'types': (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_BGR), 'resolution':df_res},#IMG_WARPED_TRANSFORMED + {'types': (t.IMG_WARPED_TRANSFORMED, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_M), 'resolution':df_res}, + {'types': (t.IMG_SOURCE, training_target, t.MODE_BGR), 'resolution':resolution}, + ] ), + + SampleGeneratorFaceTemporal(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, + temporal_image_count=3, + sample_process_options=SampleProcessor.Options(random_flip=False), + output_sample_types=[{'types': (t.IMG_SOURCE, t.FACE_TYPE_FULL_NO_ALIGN, t.MODE_BGR), 'resolution':df_res}, + {'types': (t.IMG_SOURCE, t.NONE, t.MODE_BGR), 'resolution':resolution}, + ] ), + ] + + if self.stage == 1: + generators[2].set_active(False) + generators[3].set_active(False) + elif self.stage == 2: + generators[0].set_active(False) + generators[1].set_active(False) + + self.set_training_data_generators (generators) + else: + self.G_convert = K.function([real_B64_t0, real_B64_t1, real_B64_t2],[rec_C_AB_t1]) + + #override , return [ [model, filename],... ] list + def get_model_filename_list(self): + return [ [self.enc, 'enc.h5'], + [self.decA64, 'decA64.h5'], + [self.decB64, 'decB64.h5'], + [self.C, 'C.h5'], + [self.D, 'D.h5'], + #[self.CD, 'CD.h5'], + ] + + #override + def onSave(self): + self.save_weights_safe( self.get_model_filename_list() ) + + #override + def onTrainOneIter(self, generators_samples, generators_list): + warped_src64, src64, src64m = generators_samples[0] + warped_dst64, dst64, dst64m = generators_samples[1] + + real_A64_t0, real_A64m_t0, real_A_t0, real_A64_t1, real_A64m_t1, real_A_t1, real_A64_t2, real_A64m_t2, real_A_t2 = generators_samples[2] + real_B64_t0, _, real_B64_t1, _, real_B64_t2, _ = generators_samples[3] + + if self.stage == 0 or self.stage == 1: + loss, = self.AB64_train ( [warped_src64, src64, src64m, warped_dst64, dst64, dst64m] ) + loss_D, = self.D_train ( [warped_src64, src64, src64m, warped_dst64, dst64, dst64m] ) + if self.stage != 0: + loss_C = loss_CD = 0 + + if self.stage == 0 or self.stage == 2: + loss_C1, = self.C_train ( [real_A64_t0, real_A64m_t0, real_A_t0, + real_A64_t1, real_A64m_t1, real_A_t1, + real_A64_t2, real_A64m_t2, real_A_t2, + real_B64_t0, real_B64_t1, real_B64_t2] ) + + loss_C2, = self.C_train ( [real_A64_t2, real_A64m_t2, real_A_t2, + real_A64_t1, real_A64m_t1, real_A_t1, + real_A64_t0, real_A64m_t0, real_A_t0, + real_B64_t0, real_B64_t1, real_B64_t2] ) + + #loss_CD1, = self.CD_train ( [real_A64_t0, real_A64m_t0, real_A_t0, + # real_A64_t1, real_A64m_t1, real_A_t1, + # real_A64_t2, real_A64m_t2, real_A_t2, + # real_B64_t0, real_B64_t1, real_B64_t2] ) + # + #loss_CD2, = self.CD_train ( [real_A64_t2, real_A64m_t2, real_A_t2, + # real_A64_t1, real_A64m_t1, real_A_t1, + # real_A64_t0, real_A64m_t0, real_A_t0, + # real_B64_t0, real_B64_t1, real_B64_t2] ) + + loss_C = (loss_C1 + loss_C2) / 2 + #loss_CD = (loss_CD1 + loss_CD2) / 2 + if self.stage != 0: + loss = loss_D = 0 + + return ( ('loss', loss), ('D', loss_D), ('C', loss_C), ) #('CD', loss_CD) ) + + #override + def onGetPreview(self, sample): + test_A064w = sample[0][0][0:4] + test_A064r = sample[0][1][0:4] + test_A064m = sample[0][2][0:4] + + test_B064w = sample[1][0][0:4] + test_B064r = sample[1][1][0:4] + test_B064m = sample[1][2][0:4] + + t_src64_0 = sample[2][0][0:4] + t_src64m_0 = sample[2][1][0:4] + t_src_0 = sample[2][2][0:4] + t_src64_1 = sample[2][3][0:4] + t_src64m_1 = sample[2][4][0:4] + t_src_1 = sample[2][5][0:4] + t_src64_2 = sample[2][6][0:4] + t_src64m_2 = sample[2][7][0:4] + t_src_2 = sample[2][8][0:4] + + t_dst64_0 = sample[3][0][0:4] + t_dst_0 = sample[3][1][0:4] + t_dst64_1 = sample[3][2][0:4] + t_dst_1 = sample[3][3][0:4] + t_dst64_2 = sample[3][4][0:4] + t_dst_2 = sample[3][5][0:4] + + G64_view_result = self.G64_view ([test_A064r, test_B064r]) + test_A064r, test_B064r, rec_A64, rec_B64, rec_AB64 = [ x[0] for x in ([test_A064r, test_B064r] + G64_view_result) ] + + sample64x4 = np.concatenate ([ np.concatenate ( [rec_B64, rec_A64], axis=1 ), + np.concatenate ( [test_B064r, rec_AB64], axis=1) ], axis=0 ) + + sample64x4 = cv2.resize (sample64x4, (self.resolution, self.resolution) ) + + G_view_result = self.G_view([t_src64_0, t_src64m_0, t_src64_1, t_src64m_1, t_src64_2, t_src64m_2, t_dst64_0, t_dst64_1, t_dst64_2 ]) + + t_dst_0, t_dst_1, t_dst_2, rec_C_A_t0, rec_C_A_t1, rec_C_A_t2, rec_C_AB_t0, rec_C_AB_t1, rec_C_AB_t2 = [ x[0] for x in ([t_dst_0, t_dst_1, t_dst_2, ] + G_view_result) ] + + c1 = np.concatenate ( (sample64x4, rec_C_A_t0, t_dst_0, rec_C_AB_t0 ), axis=1 ) + c2 = np.concatenate ( (sample64x4, rec_C_A_t1, t_dst_1, rec_C_AB_t1 ), axis=1 ) + c3 = np.concatenate ( (sample64x4, rec_C_A_t2, t_dst_2, rec_C_AB_t2 ), axis=1 ) + + r = np.concatenate ( [c1,c2,c3], axis=0 ) + + return [ ('AVATAR', r ) ] + + def predictor_func (self, prev_imgs=None, img=None, next_imgs=None, dummy_predict=False): + if dummy_predict: + z = np.zeros ( (1, self.df_res, self.df_res, 3), dtype=np.float32 ) + self.G_convert ([z,z,z]) + else: + feed = [ prev_imgs[-1][np.newaxis,...], img[np.newaxis,...], next_imgs[0][np.newaxis,...] ] + x = self.G_convert (feed)[0] + return np.clip ( x[0], 0, 1) + + #override + def get_ConverterConfig(self): + import converters + return converters.ConverterConfigFaceAvatar(predictor_func=self.predictor_func, + predictor_input_shape=(self.df_res, self.df_res, 3), + temporal_face_count=1 + ) + + @staticmethod + def NLayerDiscriminator(ndf=64, n_layers=3): + exec (nnlib.import_all(), locals(), globals()) + + #use_bias = True + #def XNormalization(x): + # return InstanceNormalization (axis=-1)(x) + use_bias = False + def XNormalization(x): + return BatchNormalization (axis=-1)(x) + + XConv2D = partial(Conv2D, use_bias=use_bias) + + def func(x): + f = ndf + + x = XConv2D( f, 4, strides=2, padding='same', use_bias=True)(x) + f = min( ndf*8, f*2 ) + x = LeakyReLU(0.2)(x) + + for i in range(n_layers): + x = XConv2D( f, 4, strides=2, padding='same')(x) + x = XNormalization(x) + x = LeakyReLU(0.2)(x) + f = min( ndf*8, f*2 ) + + x = XConv2D( f, 4, strides=1, padding='same')(x) + x = XNormalization(x) + x = LeakyReLU(0.2)(x) + + return XConv2D( 1, 4, strides=1, padding='same', use_bias=True, activation='sigmoid')(x)# + return func + + """ + @staticmethod + def Discriminator(ndf=128): + exec (nnlib.import_all(), locals(), globals()) + + #use_bias = True + #def XNormalization(x): + # return InstanceNormalization (axis=-1)(x) + use_bias = False + def XNormalization(x): + return BatchNormalization (axis=-1)(x) + + XConv2D = partial(Conv2D, use_bias=use_bias) + + def func(input): + b,h,w,c = K.int_shape(input) + + x = input + + x = XConv2D( ndf, 4, strides=2, padding='same', use_bias=True)(x) + x = LeakyReLU(0.2)(x) + + x = XConv2D( ndf*2, 4, strides=2, padding='same')(x) + x = XNormalization(x) + x = LeakyReLU(0.2)(x) + + x = XConv2D( ndf*4, 4, strides=2, padding='same')(x) + x = XNormalization(x) + x = LeakyReLU(0.2)(x) + + x = XConv2D( ndf*8, 4, strides=2, padding='same')(x) + x = XNormalization(x) + x = LeakyReLU(0.2)(x) + + return XConv2D( 1, 4, strides=1, padding='same', use_bias=True, activation='sigmoid')(x)# + return func + """ + @staticmethod + def Discriminator(ndf=128): + exec (nnlib.import_all(), locals(), globals()) + + use_bias = True + def XNormalization(x): + return InstanceNormalization (axis=-1)(x) + #use_bias = False + #def XNormalization(x): + # return BatchNormalization (axis=-1)(x) + + XConv2D = partial(Conv2D, use_bias=use_bias) + + def func(input): + b,h,w,c = K.int_shape(input) + + x = input + + x = XConv2D( ndf, 4, strides=2, padding='same', use_bias=True)(x) + x = LeakyReLU(0.2)(x) + + x = XConv2D( ndf*2, 4, strides=2, padding='same')(x) + x = XNormalization(x) + x = LeakyReLU(0.2)(x) + + x = XConv2D( ndf*4, 4, strides=2, padding='same')(x) + x = XNormalization(x) + x = LeakyReLU(0.2)(x) + + x = XConv2D( ndf*8, 4, strides=2, padding='same')(x) + x = XNormalization(x) + x = LeakyReLU(0.2)(x) + + return XConv2D( 1, 4, strides=1, padding='same', use_bias=True, activation='sigmoid')(x)# + return func + + @staticmethod + def CDiscriminator(ndf=256): + exec (nnlib.import_all(), locals(), globals()) + + use_bias = True + def XNormalization(x): + return InstanceNormalization (axis=-1)(x) + #use_bias = False + #def XNormalization(x): + # return BatchNormalization (axis=-1)(x) + + XConv2D = partial(Conv2D, use_bias=use_bias) + + def func(input): + b,h,w,c = K.int_shape(input) + + x = input + + x = XConv2D( ndf, 4, strides=2, padding='same', use_bias=True)(x) + x = LeakyReLU(0.2)(x) + + x = XConv2D( ndf*2, 4, strides=2, padding='same')(x) + x = XNormalization(x) + x = LeakyReLU(0.2)(x) + + x = XConv2D( ndf*4, 4, strides=2, padding='same')(x) + x = XNormalization(x) + x = LeakyReLU(0.2)(x) + + #x = XConv2D( ndf*8, 4, strides=2, padding='same')(x) + #x = XNormalization(x) + #x = LeakyReLU(0.2)(x) + + return XConv2D( 1, 4, strides=1, padding='same', use_bias=True, activation='sigmoid')(x)# + return func + + @staticmethod + def EncFlow(padding='zero', **kwargs): + exec (nnlib.import_all(), locals(), globals()) + + use_bias = False + def XNorm(x): + return BatchNormalization (axis=-1)(x) + XConv2D = partial(Conv2D, padding=padding, use_bias=use_bias) + + def downscale (dim): + def func(x): + return LeakyReLU(0.1)( Conv2D(dim, 5, strides=2, padding='same')(x)) + return func + + def upscale (dim): + def func(x): + return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x))) + return func + + + def func(input): + x, = input + b,h,w,c = K.int_shape(x) + + dim_res = w // 16 + + x = downscale(64)(x) + x = downscale(128)(x) + x = downscale(256)(x) + x = downscale(512)(x) + + x = Dense(512)(Flatten()(x)) + x = Dense(dim_res * dim_res * 512)(x) + x = Reshape((dim_res, dim_res, 512))(x) + x = upscale(512)(x) + return x + + return func + + @staticmethod + def DecFlow(output_nc=3, **kwargs): + exec (nnlib.import_all(), locals(), globals()) + + ResidualBlock = AVATARModel.ResidualBlock + upscale = AVATARModel.upscale + to_bgr = AVATARModel.to_bgr + + def func(input): + x = input[0] + + x = upscale(512)(x) + x = upscale(256)(x) + x = upscale(128)(x) + return to_bgr(output_nc) (x) + + return func + """ + @staticmethod + def CNet(output_nc, use_batch_norm, ngf=64, n_blocks=6, use_dropout=False): + exec (nnlib.import_all(), locals(), globals()) + + if not use_batch_norm: + use_bias = True + def XNormalization(x): + return InstanceNormalization (axis=-1)(x) + else: + use_bias = False + def XNormalization(x): + return BatchNormalization (axis=-1)(x) + + XConv2D = partial(Conv2D, padding='same', use_bias=use_bias) + XConv2DTranspose = partial(Conv2DTranspose, padding='same', use_bias=use_bias) + + def ResnetBlock(dim, use_dropout=False): + def func(input): + x = input + + x = XConv2D(dim, 3, strides=1)(x) + x = XNormalization(x) + x = ReLU()(x) + + if use_dropout: + x = Dropout(0.5)(x) + + x = XConv2D(dim, 3, strides=1)(x) + x = XNormalization(x) + x = ReLU()(x) + return Add()([x,input]) + return func + + def preprocess(target_res): + def func(input): + inp_shape = K.int_shape (input[0]) + t_len = len(input) + total_ch = 0 + for i in range(t_len): + total_ch += K.int_shape (input[i])[-1] + + K.concatenate ( input, axis=-1) ) + import code + c ode.interact(local=dict(globals(), **locals())) + + x_shape = K.int_shape(x)[1:] + + pad = (target_res - x_shape[0]) // 2 + + a = np.ones((target_res,target_res,3))*0.5 + a[pad:-pad:,pad:-pad:,:] = 0 + return K.spatial_2d_padding(x, padding=((pad, pad), (pad, pad)) ) + K.constant(a, dtype=K.floatx() ) + return func + + def func(input): + inp_shape = K.int_shape (input[0]) + t_len = len(input) + total_ch = 0 + for i in range(t_len): + total_ch += K.int_shape (input[i])[-1] + + x = Lambda ( preprocess(128) , output_shape=(inp_shape[1], inp_shape[2], total_ch) ) (input) + + x = ReLU()(XNormalization(XConv2D(ngf, 7, strides=1)(x))) + + x = ReLU()(XNormalization(XConv2D(ngf*2, 3, strides=2)(x))) + x = ReLU()(XNormalization(XConv2D(ngf*4, 3, strides=2)(x))) + + for i in range(n_blocks): + x = ResnetBlock(ngf*4, use_dropout=use_dropout)(x) + + x = ReLU()(XNormalization(XConv2DTranspose(ngf*2, 3, strides=2)(x))) + x = ReLU()(XNormalization(XConv2DTranspose(ngf , 3, strides=2)(x))) + + x = XConv2D(output_nc, 7, strides=1, activation='sigmoid', use_bias=True)(x) + + return x + + return func + """ + @staticmethod + def ResNet(output_nc, use_batch_norm, ngf=64, n_blocks=6, use_dropout=False): + exec (nnlib.import_all(), locals(), globals()) + + if not use_batch_norm: + use_bias = True + def XNormalization(x): + return InstanceNormalization (axis=-1)(x) + else: + use_bias = False + def XNormalization(x): + return BatchNormalization (axis=-1)(x) + + XConv2D = partial(Conv2D, padding='same', use_bias=use_bias) + XConv2DTranspose = partial(Conv2DTranspose, padding='same', use_bias=use_bias) + + def func(input): + + + def ResnetBlock(dim, use_dropout=False): + def func(input): + x = input + + x = XConv2D(dim, 3, strides=1)(x) + x = XNormalization(x) + x = ReLU()(x) + + if use_dropout: + x = Dropout(0.5)(x) + + x = XConv2D(dim, 3, strides=1)(x) + x = XNormalization(x) + x = ReLU()(x) + return Add()([x,input]) + return func + + x = input + + x = ReLU()(XNormalization(XConv2D(ngf, 7, strides=1)(x))) + + x = ReLU()(XNormalization(XConv2D(ngf*2, 3, strides=2)(x))) + x = ReLU()(XNormalization(XConv2D(ngf*4, 3, strides=2)(x))) + + x = ReLU()(XNormalization(XConv2D(ngf*4, 3, strides=2)(x))) + + for i in range(n_blocks): + x = ResnetBlock(ngf*4, use_dropout=use_dropout)(x) + + x = ReLU()(XNormalization(XConv2DTranspose(ngf*4, 3, strides=2)(x))) + + x = ReLU()(XNormalization(XConv2DTranspose(ngf*2, 3, strides=2)(x))) + x = ReLU()(XNormalization(XConv2DTranspose(ngf , 3, strides=2)(x))) + + x = XConv2D(output_nc, 7, strides=1, activation='sigmoid', use_bias=True)(x) + + return x + + return func + + @staticmethod + def initialize_nn_functions(): + exec (nnlib.import_all(), locals(), globals()) + + class ResidualBlock(object): + def __init__(self, filters, kernel_size=3, padding='zero', **kwargs): + self.filters = filters + self.kernel_size = kernel_size + self.padding = padding + + def __call__(self, inp): + x = inp + x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding)(x) + x = LeakyReLU(0.2)(x) + x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding)(x) + x = Add()([x, inp]) + x = LeakyReLU(0.2)(x) + return x + AVATARModel.ResidualBlock = ResidualBlock + + def downscale (dim, padding='zero', act='', **kwargs): + def func(x): + return LeakyReLU(0.2) (Conv2D(dim, kernel_size=5, strides=2, padding=padding)(x)) + return func + AVATARModel.downscale = downscale + + def upscale (dim, padding='zero', norm='', act='', **kwargs): + def func(x): + return SubpixelUpscaler()( LeakyReLU(0.2)(Conv2D(dim * 4, kernel_size=3, strides=1, padding=padding)(x))) + return func + AVATARModel.upscale = upscale + + def to_bgr (output_nc, padding='zero', **kwargs): + def func(x): + return Conv2D(output_nc, kernel_size=5, padding=padding, activation='sigmoid')(x) + return func + AVATARModel.to_bgr = to_bgr + +Model = AVATARModel \ No newline at end of file diff --git a/models/Model_RecycleGAN/__init__.py b/models/Model_AVATAR/__init__.py similarity index 100% rename from models/Model_RecycleGAN/__init__.py rename to models/Model_AVATAR/__init__.py diff --git a/models/Model_DF/Model.py b/models/Model_DF/Model.py index fb98418..0d5af07 100644 --- a/models/Model_DF/Model.py +++ b/models/Model_DF/Model.py @@ -59,13 +59,13 @@ class Model(ModelBase): sample_process_options=SampleProcessor.Options(random_flip=self.random_flip), output_sample_types=output_sample_types) ]) - + #override def get_model_filename_list(self): return [[self.encoder, 'encoder.h5'], [self.decoder_src, 'decoder_src.h5'], [self.decoder_dst, 'decoder_dst.h5']] - + #override def onSave(self): self.save_weights_safe( self.get_model_filename_list() ) @@ -111,18 +111,26 @@ class Model(ModelBase): return [ ('DF', np.concatenate ( st, axis=0 ) ) ] - def predictor_func (self, face): - x, mx = self.convert ( [ face[np.newaxis,...] ] ) - return x[0], mx[0][...,0] + def predictor_func (self, face=None, dummy_predict=False): + if dummy_predict: + self.AE_convert ([ np.zeros ( (1, 128, 128, 3) ), dtype=np.float32 ) ]) + else: + x, mx = self.convert ( [ face[np.newaxis,...] ] ) + return x[0], mx[0][...,0] #override - def get_converter(self): - from converters import ConverterMasked - return ConverterMasked(self.predictor_func, - predictor_input_size=128, - face_type=FaceType.FULL, - base_erode_mask_modifier=30, - base_blur_mask_modifier=0) + def get_ConverterConfig(self): + import converters + return converters.ConverterConfigMasked(predictor_func=self.predictor_func, + predictor_input_shape=(128,128,3), + predictor_masked=True, + face_type=FaceType.FULL, + default_mode=4, + base_erode_mask_modifier=30, + base_blur_mask_modifier=0, + default_erode_mask_modifier=0, + default_blur_mask_modifier=0, + ) def Build(self, input_layer): exec(nnlib.code_import_all, locals(), globals()) diff --git a/models/Model_H128/Model.py b/models/Model_H128/Model.py index 903014b..6be936c 100644 --- a/models/Model_H128/Model.py +++ b/models/Model_H128/Model.py @@ -75,7 +75,7 @@ class Model(ModelBase): return [[self.encoder, 'encoder.h5'], [self.decoder_src, 'decoder_src.h5'], [self.decoder_dst, 'decoder_dst.h5']] - + #override def onSave(self): self.save_weights_safe( self.get_model_filename_list() ) @@ -119,18 +119,26 @@ class Model(ModelBase): return [ ('H128', np.concatenate ( st, axis=0 ) ) ] - def predictor_func (self, face): - x, mx = self.src_view ( [ face[np.newaxis,...] ] ) - return x[0], mx[0][...,0] + def predictor_func (self, face=None, dummy_predict=False): + if dummy_predict: + self.AE_convert ([ np.zeros ( (1, 128, 128, 3) ), dtype=np.float32 ) ]) + else: + x, mx = self.src_view ( [ face[np.newaxis,...] ] ) + return x[0], mx[0][...,0] #override - def get_converter(self): - from converters import ConverterMasked - return ConverterMasked(self.predictor_func, - predictor_input_size=128, - face_type=FaceType.HALF, - base_erode_mask_modifier=100, - base_blur_mask_modifier=100) + def get_ConverterConfig(self): + import converters + return converters.ConverterConfigMasked(predictor_func=self.predictor_func, + predictor_input_shape=(128,128,3), + predictor_masked=True, + face_type=FaceType.HALF, + default_mode=4, + base_erode_mask_modifier=100, + base_blur_mask_modifier=100, + default_erode_mask_modifier=0, + default_blur_mask_modifier=0, + ) def Build(self, lighter_ae): exec(nnlib.code_import_all, locals(), globals()) diff --git a/models/Model_H64/Model.py b/models/Model_H64/Model.py index 76fa5ff..dde4f40 100644 --- a/models/Model_H64/Model.py +++ b/models/Model_H64/Model.py @@ -76,7 +76,7 @@ class Model(ModelBase): return [[self.encoder, 'encoder.h5'], [self.decoder_src, 'decoder_src.h5'], [self.decoder_dst, 'decoder_dst.h5']] - + #override def onSave(self): self.save_weights_safe( self.get_model_filename_list() ) @@ -120,18 +120,26 @@ class Model(ModelBase): return [ ('H64', np.concatenate ( st, axis=0 ) ) ] - def predictor_func (self, face): - x, mx = self.src_view ( [ face[np.newaxis,...] ] ) - return x[0], mx[0][...,0] + def predictor_func (self, face=None, dummy_predict=False): + if dummy_predict: + self.AE_convert ([ np.zeros ( (1, 64, 64, 3) ), dtype=np.float32 ) ]) + else: + x, mx = self.src_view ( [ face[np.newaxis,...] ] ) + return x[0], mx[0][...,0] #override - def get_converter(self): - from converters import ConverterMasked - return ConverterMasked(self.predictor_func, - predictor_input_size=64, - face_type=FaceType.HALF, - base_erode_mask_modifier=100, - base_blur_mask_modifier=100) + def get_ConverterConfig(self): + import converters + return converters.ConverterConfigMasked(predictor_func=self.predictor_func, + predictor_input_shape=(64,64,3), + predictor_masked=True, + face_type=FaceType.HALF, + default_mode=4, + base_erode_mask_modifier=100, + base_blur_mask_modifier=100, + default_erode_mask_modifier=0, + default_blur_mask_modifier=0, + ) def Build(self, lighter_ae): exec(nnlib.code_import_all, locals(), globals()) diff --git a/models/Model_LIAEF128/Model.py b/models/Model_LIAEF128/Model.py index 5fc03af..18f1641 100644 --- a/models/Model_LIAEF128/Model.py +++ b/models/Model_LIAEF128/Model.py @@ -70,8 +70,8 @@ class Model(ModelBase): return [[self.encoder, 'encoder.h5'], [self.decoder, 'decoder.h5'], [self.inter_B, 'inter_B.h5'], - [self.inter_AB, 'inter_AB.h5']] - + [self.inter_AB, 'inter_AB.h5']] + #override def onSave(self): self.save_weights_safe( self.get_model_filename_list() ) @@ -117,18 +117,26 @@ class Model(ModelBase): return [ ('LIAEF128', np.concatenate ( st, axis=0 ) ) ] - def predictor_func (self, face): - x, mx = self.convert ( [ face[np.newaxis,...] ] ) - return x[0], mx[0][...,0] + def predictor_func (self, face=None, dummy_predict=False): + if dummy_predict: + self.AE_convert ([ np.zeros ( (1, 128, 128, 3) ), dtype=np.float32 ) ]) + else: + x, mx = self.convert ( [ face[np.newaxis,...] ] ) + return x[0], mx[0][...,0] #override - def get_converter(self): - from converters import ConverterMasked - return ConverterMasked(self.predictor_func, - predictor_input_size=128, - face_type=FaceType.FULL, - base_erode_mask_modifier=30, - base_blur_mask_modifier=0) + def get_ConverterConfig(self): + import converters + return converters.ConverterConfigMasked(predictor_func=self.predictor_func, + predictor_input_shape=(128,128,3), + predictor_masked=True, + face_type=FaceType.FULL, + default_mode=4, + base_erode_mask_modifier=30, + base_blur_mask_modifier=0, + default_erode_mask_modifier=0, + default_blur_mask_modifier=0, + ) def Build(self, input_layer): exec(nnlib.code_import_all, locals(), globals()) diff --git a/models/Model_RecycleGAN/Model.py b/models/Model_RecycleGAN/Model.py deleted file mode 100644 index 9533ab1..0000000 --- a/models/Model_RecycleGAN/Model.py +++ /dev/null @@ -1,482 +0,0 @@ -from functools import partial - -import cv2 -import numpy as np - -from facelib import FaceType -from interact import interact as io -from mathlib import get_power_of_two -from models import ModelBase -from nnlib import nnlib -from samplelib import * - -class RecycleGANModel(ModelBase): - - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs, - ask_sort_by_yaw=False, - ask_random_flip=False, - ask_src_scale_mod=False) - - #override - def onInitializeOptions(self, is_first_run, ask_override): - if is_first_run: - self.options['resolution'] = io.input_int("Resolution ( 128,256 ?:help skip:128) : ", 128, [128,256], help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16.") - else: - self.options['resolution'] = self.options.get('resolution', 128) - - #override - def onInitialize(self, batch_size=-1, **in_options): - exec(nnlib.code_import_all, locals(), globals()) - self.set_vram_batch_requirements({6:16}) - - resolution = self.options['resolution'] - bgr_shape = (resolution, resolution, 3) - ngf = 64 - npf = 32 - ndf = 64 - lambda_A = 10 - lambda_B = 10 - - use_batch_norm = True #created_batch_size > 1 - self.GA = modelify(RecycleGANModel.ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=True))(Input(bgr_shape)) - self.GB = modelify(RecycleGANModel.ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=True))(Input(bgr_shape)) - - #self.GA = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape)) - #self.GB = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape)) - - self.PA = modelify(RecycleGANModel.UNetTemporalPredictor(bgr_shape[2], use_batch_norm, ngf=npf))([Input(bgr_shape), Input(bgr_shape)]) - self.PB = modelify(RecycleGANModel.UNetTemporalPredictor(bgr_shape[2], use_batch_norm, ngf=npf))([Input(bgr_shape), Input(bgr_shape)]) - - self.DA = modelify(RecycleGANModel.PatchDiscriminator(ndf=ndf) ) (Input(bgr_shape)) - self.DB = modelify(RecycleGANModel.PatchDiscriminator(ndf=ndf) ) (Input(bgr_shape)) - - if not self.is_first_run(): - weights_to_load = [ - (self.GA, 'GA.h5'), - (self.DA, 'DA.h5'), - (self.PA, 'PA.h5'), - (self.GB, 'GB.h5'), - (self.DB, 'DB.h5'), - (self.PB, 'PB.h5'), - ] - self.load_weights_safe(weights_to_load) - - real_A0 = Input(bgr_shape, name="real_A0") - real_A1 = Input(bgr_shape, name="real_A1") - real_A2 = Input(bgr_shape, name="real_A2") - - real_B0 = Input(bgr_shape, name="real_B0") - real_B1 = Input(bgr_shape, name="real_B1") - real_B2 = Input(bgr_shape, name="real_B2") - - DA_ones = K.ones_like ( K.shape(self.DA.outputs[0]) ) - DA_zeros = K.zeros_like ( K.shape(self.DA.outputs[0] )) - DB_ones = K.ones_like ( K.shape(self.DB.outputs[0] )) - DB_zeros = K.zeros_like ( K.shape(self.DB.outputs[0] )) - - def DLoss(labels,logits): - return K.mean(K.binary_crossentropy(labels,logits)) - - def CycleLoss (t1,t2): - return K.mean(K.abs(t1 - t2)) - - def RecurrentLOSS(t1,t2): - return K.mean(K.abs(t1 - t2)) - - def RecycleLOSS(t1,t2): - return K.mean(K.abs(t1 - t2)) - - fake_B0 = self.GA(real_A0) - fake_B1 = self.GA(real_A1) - - fake_A0 = self.GB(real_B0) - fake_A1 = self.GB(real_B1) - - real_A0_d = self.DA(real_A0) - real_A0_d_ones = K.ones_like(real_A0_d) - real_A1_d = self.DA(real_A1) - real_A1_d_ones = K.ones_like(real_A1_d) - - fake_A0_d = self.DA(fake_A0) - fake_A0_d_ones = K.ones_like(fake_A0_d) - fake_A0_d_zeros = K.zeros_like(fake_A0_d) - - fake_A1_d = self.DA(fake_A1) - fake_A1_d_ones = K.ones_like(fake_A1_d) - fake_A1_d_zeros = K.zeros_like(fake_A1_d) - - real_B0_d = self.DB(real_B0) - real_B0_d_ones = K.ones_like(real_B0_d) - - real_B1_d = self.DB(real_B1) - real_B1_d_ones = K.ones_like(real_B1_d) - - fake_B0_d = self.DB(fake_B0) - fake_B0_d_ones = K.ones_like(fake_B0_d) - fake_B0_d_zeros = K.zeros_like(fake_B0_d) - - fake_B1_d = self.DB(fake_B1) - fake_B1_d_ones = K.ones_like(fake_B1_d) - fake_B1_d_zeros = K.zeros_like(fake_B1_d) - - pred_A2 = self.PA ( [real_A0, real_A1]) - pred_B2 = self.PB ( [real_B0, real_B1]) - rec_A2 = self.GB ( self.PB ( [fake_B0, fake_B1]) ) - rec_B2 = self.GA ( self.PA ( [fake_A0, fake_A1])) - - - loss_GA = DLoss(fake_B0_d_ones, fake_B0_d ) + \ - DLoss(fake_B1_d_ones, fake_B1_d ) + \ - lambda_A * (RecurrentLOSS(pred_A2, real_A2) + \ - RecycleLOSS(rec_B2, real_B2) ) - - - weights_GA = self.GA.trainable_weights + self.PA.trainable_weights - - loss_GB = DLoss(fake_A0_d_ones, fake_A0_d ) + \ - DLoss(fake_A1_d_ones, fake_A1_d ) + \ - lambda_B * (RecurrentLOSS(pred_B2, real_B2) + \ - RecycleLOSS(rec_A2, real_A2) ) - - weights_GB = self.GB.trainable_weights + self.PB.trainable_weights - - def opt(): - return Adam(lr=2e-4, beta_1=0.5, beta_2=0.999, tf_cpu_mode=2)#, clipnorm=1) - - self.GA_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_GA], - opt().get_updates(loss_GA, weights_GA) ) - - self.GB_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_GB], - opt().get_updates(loss_GB, weights_GB) ) - - ########### - - loss_D_A0 = ( DLoss(real_A0_d_ones, real_A0_d ) + \ - DLoss(fake_A0_d_zeros, fake_A0_d ) ) * 0.5 - - loss_D_A1 = ( DLoss(real_A1_d_ones, real_A1_d ) + \ - DLoss(fake_A1_d_zeros, fake_A1_d ) ) * 0.5 - - loss_D_A = loss_D_A0 + loss_D_A1 - - self.DA_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_D_A], - opt().get_updates(loss_D_A, self.DA.trainable_weights) ) - - ############ - - loss_D_B0 = ( DLoss(real_B0_d_ones, real_B0_d ) + \ - DLoss(fake_B0_d_zeros, fake_B0_d ) ) * 0.5 - - loss_D_B1 = ( DLoss(real_B1_d_ones, real_B1_d ) + \ - DLoss(fake_B1_d_zeros, fake_B1_d ) ) * 0.5 - - loss_D_B = loss_D_B0 + loss_D_B1 - - self.DB_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_D_B], - opt().get_updates(loss_D_B, self.DB.trainable_weights) ) - - ############ - - - self.G_view = K.function([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[fake_A0, fake_A1, pred_A2, rec_A2, fake_B0, fake_B1, pred_B2, rec_B2 ]) - - - - if self.is_training_mode: - t = SampleProcessor.Types - output_sample_types=[ { 'types': (t.IMG_SOURCE, t.MODE_BGR), 'resolution':resolution, 'normalize_tanh' : True} ] - - self.set_training_data_generators ([ - SampleGeneratorImageTemporal(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size, - temporal_image_count=3, - sample_process_options=SampleProcessor.Options(random_flip = False), - output_sample_types=output_sample_types ), - - SampleGeneratorImageTemporal(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size, - temporal_image_count=3, - sample_process_options=SampleProcessor.Options(random_flip = False), - output_sample_types=output_sample_types ), - ]) - else: - self.G_convert = K.function([real_B0],[fake_A0]) - - #override - def get_model_filename_list(self): - return [ [self.GA, 'GA.h5'], - [self.GB, 'GB.h5'], - [self.DA, 'DA.h5'], - [self.DB, 'DB.h5'], - [self.PA, 'PA.h5'], - [self.PB, 'PB.h5'] ] - - #override - def onSave(self): - self.save_weights_safe( self.get_model_filename_list() ) - - #override - def onTrainOneIter(self, generators_samples, generators_list): - source_src_0, source_src_1, source_src_2, = generators_samples[0] - source_dst_0, source_dst_1, source_dst_2, = generators_samples[1] - - feed = [source_src_0, source_src_1, source_src_2, source_dst_0, source_dst_1, source_dst_2] - - loss_GA, = self.GA_train ( feed ) - loss_GB, = self.GB_train ( feed ) - loss_DA, = self.DA_train( feed ) - loss_DB, = self.DB_train( feed ) - - return ( ('GA', loss_GA), ('GB', loss_GB), ('DA', loss_DA), ('DB', loss_DB) ) - - #override - def onGetPreview(self, sample): - test_A0 = sample[0][0] - test_A1 = sample[0][1] - test_A2 = sample[0][2] - - test_B0 = sample[1][0] - test_B1 = sample[1][1] - test_B2 = sample[1][2] - - G_view_result = self.G_view([test_A0, test_A1, test_A2, test_B0, test_B1, test_B2]) - - fake_A0, fake_A1, pred_A2, rec_A2, fake_B0, fake_B1, pred_B2, rec_B2 = [ x[0] / 2 + 0.5 for x in G_view_result] - test_A0, test_A1, test_A2, test_B0, test_B1, test_B2 = [ x[0] / 2 + 0.5 for x in [test_A0, test_A1, test_A2, test_B0, test_B1, test_B2] ] - - r = np.concatenate ((np.concatenate ( (test_A0, test_A1, test_A2, pred_A2, fake_B0, fake_B1, rec_A2), axis=1), - np.concatenate ( (test_B0, test_B1, test_B2, pred_B2, fake_A0, fake_A1, rec_B2), axis=1) - ), axis=0) - - return [ ('RecycleGAN', r ) ] - - def predictor_func (self, face): - x = self.G_convert ( [ face[np.newaxis,...]*2-1 ] )[0] - return np.clip ( x[0] / 2 + 0.5 , 0, 1) - - #override - def get_converter(self, **in_options): - from converters import ConverterImage - return ConverterImage(self.predictor_func, - predictor_input_size=self.options['resolution'], - **in_options) - - @staticmethod - def ResNet(output_nc, use_batch_norm, ngf=64, n_blocks=6, use_dropout=False): - exec (nnlib.import_all(), locals(), globals()) - - if not use_batch_norm: - use_bias = True - def XNormalization(x): - return InstanceNormalization (axis=-1)(x) - else: - use_bias = False - def XNormalization(x): - return BatchNormalization (axis=-1)(x) - - XConv2D = partial(Conv2D, padding='same', use_bias=use_bias) - XConv2DTranspose = partial(Conv2DTranspose, padding='same', use_bias=use_bias) - - def func(input): - - - def ResnetBlock(dim, use_dropout=False): - def func(input): - x = input - - x = XConv2D(dim, 3, strides=1)(x) - x = XNormalization(x) - x = ReLU()(x) - - if use_dropout: - x = Dropout(0.5)(x) - - x = XConv2D(dim, 3, strides=1)(x) - x = XNormalization(x) - x = ReLU()(x) - return Add()([x,input]) - return func - - x = input - - x = ReLU()(XNormalization(XConv2D(ngf, 7, strides=1)(x))) - - x = ReLU()(XNormalization(XConv2D(ngf*2, 3, strides=2)(x))) - x = ReLU()(XNormalization(XConv2D(ngf*4, 3, strides=2)(x))) - - for i in range(n_blocks): - x = ResnetBlock(ngf*4, use_dropout=use_dropout)(x) - - x = ReLU()(XNormalization(XConv2DTranspose(ngf*2, 3, strides=2)(x))) - x = ReLU()(XNormalization(XConv2DTranspose(ngf , 3, strides=2)(x))) - - x = XConv2D(output_nc, 7, strides=1, activation='tanh', use_bias=True)(x) - - return x - - return func - - @staticmethod - def UNet(output_nc, use_batch_norm, ngf=64, use_dropout=False): - exec (nnlib.import_all(), locals(), globals()) - - if not use_batch_norm: - use_bias = True - def XNormalizationL(): - return InstanceNormalization (axis=-1) - else: - use_bias = False - def XNormalizationL(): - return BatchNormalization (axis=-1) - - def XNormalization(x): - return XNormalizationL()(x) - - XConv2D = partial(Conv2D, padding='same', use_bias=use_bias) - XConv2DTranspose = partial(Conv2DTranspose, padding='same', use_bias=use_bias) - - def func(input): - - b,h,w,c = K.int_shape(input) - - n_downs = get_power_of_two(w) - 4 - - Norm = XNormalizationL() - Norm2 = XNormalizationL() - Norm4 = XNormalizationL() - Norm8 = XNormalizationL() - - x = input - - x = e1 = XConv2D( ngf, 4, strides=2, use_bias=True ) (x) - - x = e2 = Norm2( XConv2D( ngf*2, 4, strides=2 )( LeakyReLU(0.2)(x) ) ) - x = e3 = Norm4( XConv2D( ngf*4, 4, strides=2 )( LeakyReLU(0.2)(x) ) ) - - l = [] - for i in range(n_downs): - x = Norm8( XConv2D( ngf*8, 4, strides=2 )( LeakyReLU(0.2)(x) ) ) - l += [x] - - x = XConv2D( ngf*8, 4, strides=2, use_bias=True )( LeakyReLU(0.2)(x) ) - - for i in range(n_downs): - x = Norm8( XConv2DTranspose( ngf*8, 4, strides=2 )( ReLU()(x) ) ) - if i <= n_downs-2: - x = Dropout(0.5)(x) - x = Concatenate(axis=-1)([x, l[-i-1] ]) - - x = Norm4( XConv2DTranspose( ngf*4, 4, strides=2 )( ReLU()(x) ) ) - x = Concatenate(axis=-1)([x, e3]) - - x = Norm2( XConv2DTranspose( ngf*2, 4, strides=2 )( ReLU()(x) ) ) - x = Concatenate(axis=-1)([x, e2]) - - x = Norm( XConv2DTranspose( ngf, 4, strides=2 )( ReLU()(x) ) ) - x = Concatenate(axis=-1)([x, e1]) - - x = XConv2DTranspose(output_nc, 4, strides=2, activation='tanh', use_bias=True)( ReLU()(x) ) - - return x - return func - - @staticmethod - def UNetTemporalPredictor(output_nc, use_batch_norm, ngf=64, use_dropout=False): - exec (nnlib.import_all(), locals(), globals()) - def func(inputs): - past_2_image_tensor, past_1_image_tensor = inputs - - x = Concatenate(axis=-1)([ past_2_image_tensor, past_1_image_tensor ]) - x = UNet(3, use_batch_norm, ngf=ngf, use_dropout=use_dropout) (x) - - return x - - return func - - @staticmethod - def PatchDiscriminator(ndf=64): - exec (nnlib.import_all(), locals(), globals()) - - #use_bias = True - #def XNormalization(x): - # return InstanceNormalization (axis=-1)(x) - use_bias = False - def XNormalization(x): - return BatchNormalization (axis=-1)(x) - - XConv2D = partial(Conv2D, use_bias=use_bias) - - def func(input): - b,h,w,c = K.int_shape(input) - - x = input - - x = ZeroPadding2D((1,1))(x) - x = XConv2D( ndf, 4, strides=2, padding='valid', use_bias=True)(x) - x = LeakyReLU(0.2)(x) - - x = ZeroPadding2D((1,1))(x) - x = XConv2D( ndf*2, 4, strides=2, padding='valid')(x) - x = XNormalization(x) - x = LeakyReLU(0.2)(x) - - x = ZeroPadding2D((1,1))(x) - x = XConv2D( ndf*4, 4, strides=2, padding='valid')(x) - x = XNormalization(x) - x = LeakyReLU(0.2)(x) - - x = ZeroPadding2D((1,1))(x) - x = XConv2D( ndf*8, 4, strides=2, padding='valid')(x) - x = XNormalization(x) - x = LeakyReLU(0.2)(x) - - x = ZeroPadding2D((1,1))(x) - x = XConv2D( ndf*8, 4, strides=2, padding='valid')(x) - x = XNormalization(x) - x = LeakyReLU(0.2)(x) - - x = ZeroPadding2D((1,1))(x) - return XConv2D( 1, 4, strides=1, padding='valid', use_bias=True, activation='sigmoid')(x)# - return func - - @staticmethod - def NLayerDiscriminator(ndf=64, n_layers=3): - exec (nnlib.import_all(), locals(), globals()) - - #use_bias = True - #def XNormalization(x): - # return InstanceNormalization (axis=-1)(x) - use_bias = False - def XNormalization(x): - return BatchNormalization (axis=-1)(x) - - XConv2D = partial(Conv2D, use_bias=use_bias) - - def func(input): - b,h,w,c = K.int_shape(input) - - x = input - - f = ndf - - x = ZeroPadding2D((1,1))(x) - x = XConv2D( f, 4, strides=2, padding='valid', use_bias=True)(x) - f = min( ndf*8, f*2 ) - x = LeakyReLU(0.2)(x) - - for i in range(n_layers): - x = ZeroPadding2D((1,1))(x) - x = XConv2D( f, 4, strides=2, padding='valid')(x) - f = min( ndf*8, f*2 ) - x = XNormalization(x) - x = LeakyReLU(0.2)(x) - - x = ZeroPadding2D((1,1))(x) - x = XConv2D( f, 4, strides=1, padding='valid')(x) - x = XNormalization(x) - x = LeakyReLU(0.2)(x) - - x = ZeroPadding2D((1,1))(x) - return XConv2D( 1, 4, strides=1, padding='valid', use_bias=True, activation='sigmoid')(x)# - return func - -Model = RecycleGANModel diff --git a/models/Model_SAE/Model.py b/models/Model_SAE/Model.py index dc2f0a2..555880d 100644 --- a/models/Model_SAE/Model.py +++ b/models/Model_SAE/Model.py @@ -24,7 +24,7 @@ class SAEModel(ModelBase): #override def onInitializeOptions(self, is_first_run, ask_override): yn_str = {True:'y',False:'n'} - + default_resolution = 128 default_archi = 'df' default_face_type = 'f' @@ -90,20 +90,20 @@ class SAEModel(ModelBase): default_apply_random_ct = False if is_first_run else self.options.get('apply_random_ct', False) self.options['apply_random_ct'] = io.input_bool ("Apply random color transfer to src faceset? (y/n, ?:help skip:%s) : " % (yn_str[default_apply_random_ct]), default_apply_random_ct, help_message="Increase variativity of src samples by apply LCT color transfer from random dst samples. It is like 'face_style' learning, but more precise color transfer and without risk of model collapse, also it does not require additional GPU resources, but the training time may be longer, due to the src faceset is becoming more diverse.") - + if nnlib.device.backend != 'plaidML': # todo https://github.com/plaidml/plaidml/issues/301 default_clipgrad = False if is_first_run else self.options.get('clipgrad', False) self.options['clipgrad'] = io.input_bool ("Enable gradient clipping? (y/n, ?:help skip:%s) : " % (yn_str[default_clipgrad]), default_clipgrad, help_message="Gradient clipping reduces chance of model collapse, sacrificing speed of training.") else: self.options['clipgrad'] = False - + else: self.options['pixel_loss'] = self.options.get('pixel_loss', False) self.options['face_style_power'] = self.options.get('face_style_power', default_face_style_power) self.options['bg_style_power'] = self.options.get('bg_style_power', default_bg_style_power) self.options['apply_random_ct'] = self.options.get('apply_random_ct', False) self.options['clipgrad'] = self.options.get('clipgrad', False) - + if is_first_run: self.options['pretrain'] = io.input_bool ("Pretrain the model? (y/n, ?:help skip:n) : ", False, help_message="Pretrain the model with large amount of various faces. This technique may help to train the fake with overly different face shapes and light conditions of src/dst data. Face will be look more like a morphed. To reduce the morph effect, some model files will be initialized but not be updated after pretrain: LIAE: inter_AB.h5 DF: encoder.h5. The longer you pretrain the model the more morphed face will look. After that, save and run the training again.") else: @@ -383,7 +383,7 @@ class SAEModel(ModelBase): [ {'types' : (t.IMG_TRANSFORMED, face_type, t_mode_bgr), 'resolution': resolution // (2**i)} for i in range(ms_count)] + \ [ {'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution // (2**i) } for i in range(ms_count)]) ]) - + #override def get_model_filename_list(self): ar = [] @@ -413,7 +413,7 @@ class SAEModel(ModelBase): ar += [ [self.decoder_srcm, 'decoder_srcm.h5'], [self.decoder_dstm, 'decoder_dstm.h5'] ] return ar - + #override def onSave(self): self.save_weights_safe( self.get_model_filename_list() ) @@ -469,17 +469,20 @@ class SAEModel(ModelBase): return result - def predictor_func (self, face): - if self.options['learn_mask']: - bgr, mask_dst_dstm, mask_src_dstm = self.AE_convert ([face[np.newaxis,...]]) - mask = mask_dst_dstm[0] * mask_src_dstm[0] - return bgr[0], mask[...,0] + def predictor_func (self, face=None, dummy_predict=False): + if dummy_predict: + self.AE_convert ([ np.zeros ( (1, self.options['resolution'], self.options['resolution'], 3), dtype=np.float32 ) ]) else: - bgr, = self.AE_convert ([face[np.newaxis,...]]) - return bgr[0] + if self.options['learn_mask']: + bgr, mask_dst_dstm, mask_src_dstm = self.AE_convert ([face[np.newaxis,...]]) + mask = mask_dst_dstm[0] * mask_src_dstm[0] + return bgr[0], mask[...,0] + else: + bgr, = self.AE_convert ([face[np.newaxis,...]]) + return bgr[0] #override - def get_converter(self): + def get_ConverterConfig(self): base_erode_mask_modifier = 30 if self.options['face_type'] == 'f' else 100 base_blur_mask_modifier = 0 if self.options['face_type'] == 'f' else 100 @@ -489,17 +492,18 @@ class SAEModel(ModelBase): face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF - from converters import ConverterMasked - return ConverterMasked(self.predictor_func, - predictor_input_size=self.options['resolution'], - predictor_masked=self.options['learn_mask'], - face_type=face_type, - default_mode = 1 if self.options['apply_random_ct'] or self.options['face_style_power'] or self.options['bg_style_power'] else 4, - base_erode_mask_modifier=base_erode_mask_modifier, - base_blur_mask_modifier=base_blur_mask_modifier, - default_erode_mask_modifier=default_erode_mask_modifier, - default_blur_mask_modifier=default_blur_mask_modifier, - clip_hborder_mask_per=0.0625 if (self.options['face_type'] == 'f') else 0) + import converters + return converters.ConverterConfigMasked(predictor_func=self.predictor_func, + predictor_input_shape=(self.options['resolution'], self.options['resolution'], 3), + predictor_masked=self.options['learn_mask'], + face_type=face_type, + default_mode = 1 if self.options['apply_random_ct'] or self.options['face_style_power'] or self.options['bg_style_power'] else 4, + base_erode_mask_modifier=base_erode_mask_modifier, + base_blur_mask_modifier=base_blur_mask_modifier, + default_erode_mask_modifier=default_erode_mask_modifier, + default_blur_mask_modifier=default_blur_mask_modifier, + clip_hborder_mask_per=0.0625 if (self.options['face_type'] == 'f') else 0, + ) @staticmethod def initialize_nn_functions(): @@ -545,7 +549,7 @@ class SAEModel(ModelBase): return Norm(norm)( Act(act) (Conv2D(dim, kernel_size=5, strides=2, padding=padding)(x)) ) return func 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)) ) ) diff --git a/nnlib/nnlib.py b/nnlib/nnlib.py index 97e98b9..2d3b5e4 100644 --- a/nnlib/nnlib.py +++ b/nnlib/nnlib.py @@ -133,6 +133,10 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator os.environ['TF_MIN_GPU_MULTIPROCESSOR_COUNT'] = '2' os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' #tf log errors only + + import warnings + warnings.simplefilter(action='ignore', category=FutureWarning) + import tensorflow as tf nnlib.tf = tf @@ -457,7 +461,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator nnlib.PixelShuffler = PixelShuffler nnlib.SubpixelUpscaler = PixelShuffler - + class BlurPool(KL.Layer): """ https://arxiv.org/abs/1904.11486 https://github.com/adobe/antialiased-cnns @@ -472,36 +476,36 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator self.a = np.array([1., 1.]) elif(self.filt_size==3): self.a = np.array([1., 2., 1.]) - elif(self.filt_size==4): + elif(self.filt_size==4): self.a = np.array([1., 3., 3., 1.]) - elif(self.filt_size==5): + elif(self.filt_size==5): self.a = np.array([1., 4., 6., 4., 1.]) - elif(self.filt_size==6): + elif(self.filt_size==6): self.a = np.array([1., 5., 10., 10., 5., 1.]) - elif(self.filt_size==7): + 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] + 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 = 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): """ @@ -537,40 +541,40 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator 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) + 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 + return out nnlib.SelfAttention = SelfAttention - + class Adam(keras.optimizers.Optimizer): """Adam optimizer. diff --git a/requirements-colab.txt b/requirements-colab.txt index 7c69191..6ed6edc 100644 --- a/requirements-colab.txt +++ b/requirements-colab.txt @@ -1,7 +1,7 @@ -numpy==1.16.3 +numpy==1.17.0 h5py==2.9.0 Keras==2.2.4 -opencv-python==4.0.0.21 +opencv-python==4.1.0.25 tensorflow-gpu==1.13.1 plaidml-keras==0.5.0 scikit-image diff --git a/requirements-cpu.txt b/requirements-cpu.txt index 6f5f27f..1152b18 100644 --- a/requirements-cpu.txt +++ b/requirements-cpu.txt @@ -1,7 +1,7 @@ -numpy==1.16.3 +numpy==1.17.0 h5py==2.9.0 Keras==2.2.4 -opencv-python==4.0.0.21 +opencv-python==4.1.0.25 tensorflow==1.12.0 scikit-image tqdm diff --git a/requirements-cuda.txt b/requirements-cuda.txt index 5e50c09..0d6ab03 100644 --- a/requirements-cuda.txt +++ b/requirements-cuda.txt @@ -1,7 +1,7 @@ -numpy==1.16.3 +numpy==1.17.0 h5py==2.9.0 Keras==2.2.4 -opencv-python==4.0.0.21 +opencv-python==4.1.0.25 tensorflow-gpu==1.12.0 plaidml==0.6.0 plaidml-keras==0.5.0 diff --git a/requirements-opencl.txt b/requirements-opencl.txt index 99a13c1..2ac2908 100644 --- a/requirements-opencl.txt +++ b/requirements-opencl.txt @@ -1,7 +1,7 @@ -numpy==1.16.3 +numpy==1.17.0 h5py==2.9.0 Keras==2.2.4 -opencv-python==4.0.0.21 +opencv-python==4.1.0.25 tensorflow==1.12.0 plaidml==0.6.0 plaidml-keras==0.5.0 diff --git a/samplelib/SampleProcessor.py b/samplelib/SampleProcessor.py index 338dc65..59799b2 100644 --- a/samplelib/SampleProcessor.py +++ b/samplelib/SampleProcessor.py @@ -62,6 +62,7 @@ class SampleProcessor(object): FACE_TYPE_HEAD = 12 #currently unused FACE_TYPE_AVATAR = 13 #currently unused FACE_TYPE_FULL_NO_ALIGN = 14 + FACE_TYPE_HEAD_NO_ALIGN = 15 FACE_TYPE_END = 20 MODE_BEGIN = 40 @@ -73,7 +74,6 @@ class SampleProcessor(object): 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] ): self.random_flip = random_flip self.rotation_range = rotation_range @@ -81,6 +81,13 @@ class SampleProcessor(object): self.tx_range = tx_range self.ty_range = ty_range + SPTF_FACETYPE_TO_FACETYPE = { Types.FACE_TYPE_HALF : FaceType.HALF, + Types.FACE_TYPE_FULL : FaceType.FULL, + Types.FACE_TYPE_HEAD : FaceType.HEAD, + Types.FACE_TYPE_FULL_NO_ALIGN : FaceType.FULL_NO_ALIGN, + Types.FACE_TYPE_HEAD_NO_ALIGN : FaceType.HEAD_NO_ALIGN, + } + @staticmethod def process (sample, sample_process_options, output_sample_types, debug, ct_sample=None): SPTF = SampleProcessor.Types @@ -101,10 +108,7 @@ class SampleProcessor(object): 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_FULL_NO_ALIGN : FaceType.FULL_NO_ALIGN } + outputs = [] for opts in output_sample_types: @@ -171,7 +175,7 @@ class SampleProcessor(object): img = np.concatenate( (img, mask ), -1 ) return img - + img = cached_images.get(img_type, None) if img is None: @@ -196,7 +200,7 @@ class SampleProcessor(object): if cur_sample.ie_polys is not None: cur_sample.ie_polys.overlay_mask(mask) - + if sample.face_type == FaceType.MARK_ONLY: if mask is not None: img = np.concatenate( (img, mask), -1 ) @@ -206,20 +210,20 @@ class SampleProcessor(object): cached_images[img_type] = img if is_face_sample and target_face_type != SPTF.NONE: - ft = SPTF_FACETYPE_TO_FACETYPE[target_face_type] + ft = SampleProcessor.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) ) - - if sample.face_type == FaceType.MARK_ONLY: + + if sample.face_type == FaceType.MARK_ONLY: img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, sample.shape[0], ft), (sample.shape[0],sample.shape[0]), 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: img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC ) diff --git a/utils/os_utils.py b/utils/os_utils.py index ff1bd98..9297ba9 100644 --- a/utils/os_utils.py +++ b/utils/os_utils.py @@ -23,3 +23,14 @@ def set_process_lowest_prio(): def set_process_dpi_aware(): if sys.platform[0:3] == 'win': windll.user32.SetProcessDPIAware(True) + +def get_screen_size(): + if sys.platform[0:3] == 'win': + user32 = windll.user32 + return user32.GetSystemMetrics(0), user32.GetSystemMetrics(1) + elif 'darwin' in sys.platform: + pass + elif 'linux' in sys.platform: + pass + + return (1366, 768) \ No newline at end of file