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

Added interactive converter.

Browse source 
With interactive converter you can change any parameter of any frame and see the result in real time.

Converter: added motion_blur_power param.
Motion blur is applied by precomputed motion vectors.
So the moving face will look more realistic.

RecycleGAN model is removed.

Added experimental AVATAR model. Minimum required VRAM is 6GB (NVIDIA), 12GB (AMD)
Usage:
1) place data_src.mp4 10-20min square resolution video of news reporter sitting at the table with static background,
   other faces should not appear in frames.
2) process "extract images from video data_src.bat" with FULL fps
3) place data_dst.mp4 video of face who will control the src face
4) process "extract images from video data_dst FULL FPS.bat"
5) process "data_src mark faces S3FD best GPU.bat"
6) process "data_dst extract unaligned faces S3FD best GPU.bat"
7) train AVATAR.bat stage 1, tune batch size to maximum for your card (32 for 6GB), train to 50k+ iters.
8) train AVATAR.bat stage 2, tune batch size to maximum for your card (4 for 6GB), train to decent sharpness.
9) convert AVATAR.bat
10) converted to mp4.bat

updated versions of modules
This commit is contained in:
iperov 2019-08-24 12:57:29 +04:00
parent 3f0bf2e994
commit 407ce3b1ca
46 changed files with 2394 additions and 1659 deletions
Download patch file Download diff file Expand all files Collapse all files

35
converters/ConvertAvatar.py Normal file
View file

@ -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)

392
converters/ConvertMasked.py Normal file
View file

@ -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)

50
converters/Converter.py
View file

@ -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

61
converters/ConverterAvatar.py
View file

@ -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

317
converters/ConverterConfig.py Normal file
View file

@ -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"
"================"
)

50
converters/ConverterImage.py
View file

@ -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

436
converters/ConverterMasked.py
View file

@ -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

6
converters/FrameInfo.py Normal file
View file

@ -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

8
converters/__init__.py
View file

@ -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

BIN
doc/manual_en_google_translated.docx
View file

Binary file not shown.

BIN
doc/manual_en_google_translated.pdf
View file

Binary file not shown.

BIN
doc/manual_ru.pdf
View file

Binary file not shown.

BIN
doc/manual_ru_source.docx
View file

Binary file not shown.

9
facelib/FaceType.py
View file

@ -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'
}

13
facelib/LandmarksProcessor.py
View file

@ -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)

15
imagelib/__init__.py
View file

@ -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

148
imagelib/blur.py
View file

@ -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])
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)

27
imagelib/common.py
View file

@ -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
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])

16
interact/interact.py
View file

@ -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)

9
joblib/SubprocessFunctionCaller.py
View file

@ -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()

64
joblib/SubprocessorBase.py
View file

@ -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[:]:

16
main.py
View file

@ -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)

658
mainscripts/Converter.py
View file

@ -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) )
"""

138
mainscripts/ConverterScreen/ConverterScreen.py Normal file
View file

@ -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)

1
mainscripts/ConverterScreen/__init__.py Normal file
View file

@ -0,0 +1 @@
from .ConverterScreen import Screen, ScreenManager

BIN
mainscripts/ConverterScreen/gfx/sand_clock_64.png Normal file
View file

Binary file not shown.
Side by side

After

Width:  |  Height:  |  Size: 3.2 KiB

BIN
mainscripts/gfx/help_converter_face_avatar.jpg Normal file
View file

Binary file not shown.
Side by side

After

Width:  |  Height:  |  Size: 115 KiB

BIN
mainscripts/gfx/help_converter_face_avatar_source.psd Normal file
View file

Binary file not shown.

BIN
mainscripts/gfx/help_converter_masked.jpg Normal file
View file

Binary file not shown.
Side by side

After

Width:  |  Height:  |  Size: 259 KiB

BIN
mainscripts/gfx/help_converter_masked_source.psd Normal file
View file

Binary file not shown.

79
models/ModelBase.py
View file

@ -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)

740
models/Model_AVATAR/Model.py Normal file
View file

@ -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

0
models/Model_RecycleGAN/__init__.py → models/Model_AVATAR/__init__.py
View file

32
models/Model_DF/Model.py
View file

@ -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())

30
models/Model_H128/Model.py
View file

@ -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())

30
models/Model_H64/Model.py
View file

@ -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())

32
models/Model_LIAEF128/Model.py
View file

@ -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())

482
models/Model_RecycleGAN/Model.py
View file

@ -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

56
models/Model_SAE/Model.py
View file

@ -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)) ) )

48
nnlib/nnlib.py
View file

@ -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.

4
requirements-colab.txt
View file

@ -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

4
requirements-cpu.txt
View file

@ -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

4
requirements-cuda.txt
View file

@ -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

4
requirements-opencl.txt
View file

@ -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

28
samplelib/SampleProcessor.py
View file

@ -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 )

11
utils/os_utils.py
View file

@ -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)