github/DeepFaceLab
1
0
Fork 0
mirror of https://github.com/iperov/DeepFaceLab.git synced 2025-08-22 14:24:40 -07:00
Code Issues Projects Releases Packages Wiki Activity Actions

Revert "1"

Browse source
This commit is contained in:
NiteSoft 2019-09-13 10:52:12 +03:00 • committed by GitHub
commit 2e501ab8ce
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
82 changed files with 2556 additions and 4244 deletions
Download patch file Download diff file Expand all files Collapse all files

21
README.md
View file

@ -1,22 +1,18 @@
![](doc/DFL_welcome.jpg)
![](doc/DFL_welcome.jpg)
![](doc/logo_cuda.jpg)![](doc/logo_opencl.jpg)![](doc/logo_keras.jpg)![](doc/logo_tensorflow.jpg)![](doc/logo_plaidml.jpg)
#deepfakes #faceswap #face-swap #deep-learning #deeplearning #deep-neural-networks #deepface #deep-face-swap #fakeapp #fake-app #neural-networks #neural-nets
## **DeepFaceLab** is a tool that utilizes machine learning to replace faces in videos.
If you like this software, please consider a donation.
GOAL: next DeepFacelab update.
[Donate via Yandex.Money](https://money.yandex.ru/to/41001142318065)
[Donate via Paypal](https://www.paypal.com/cgi-bin/webscr?cmd=_donations&business=lepersorium@gmail.com&lc=US&no_note=0&item_name=Support+DeepFaceLab&cn=&curency_code=USD&bn=PP-DonationsBF:btn_donateCC_LG.gif:NonHosted)
bitcoin:31mPd6DxPCzbpCMZk4k1koWAbErSyqkAXr
[Donate via Paypal](https://www.paypal.com/cgi-bin/webscr?cmd=_s-xclick&hosted_button_id=KK5ZCH4JXWMQS&source=url)
- ### [Gallery](doc/gallery/doc_gallery.md)
[Donate via Alipay](https://i.loli.net/2019/01/13/5c3ae3829809f.jpg)
bitcoin:31mPd6DxPCzbpCMZk4k1koWAbErSyqkAXr
- ### Manuals:
@ -38,15 +34,12 @@ bitcoin:31mPd6DxPCzbpCMZk4k1koWAbErSyqkAXr
- ### Communication groups:
[mrdeepfakes (English)](https://mrdeepfakes.com/forums/) - the biggest SFW and NSFW community
(Chinese) QQ group 951138799 for ML/AI experts
[deepfakes (Chinese)](https://deepfakescn.com)
[deepfakes (Chinese) (outdated) ](https://deepfakes.com.cn/)
[reddit r/GifFakes/ (English)](https://www.reddit.com/r/GifFakes/new/)
[reddit r/SFWdeepfakes/ (English)](https://www.reddit.com/r/SFWdeepfakes/new/)
[mrdeepfakes (English)](https://mrdeepfakes.com/forums/)
[reddit (English)](https://www.reddit.com/r/GifFakes/new/)

41
converters/ConvertAvatar.py
View file

@ -1,41 +0,0 @@
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], borderMode=cv2.BORDER_REPLICATE, flags=cv2.INTER_CUBIC )
return img
def ConvertFaceAvatar (predictor_func, predictor_input_shape, cfg, prev_temporal_frame_infos, frame_info, next_temporal_frame_infos):
inp_sh = 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 = predictor_func ( prev_imgs, img, next_imgs )
if cfg.super_resolution_mode != 0:
prd_f = cfg.superres_func(cfg.super_resolution_mode, prd_f)
if cfg.sharpen_mode != 0 and cfg.sharpen_amount != 0:
prd_f = cfg.sharpen_func ( prd_f, cfg.sharpen_mode, 3, cfg.sharpen_amount)
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)

352
converters/ConvertMasked.py
View file

@ -1,352 +0,0 @@
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 (predictor_func, predictor_input_shape, cfg, frame_info, img_bgr_uint8, img_bgr, img_face_landmarks):
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 = predictor_input_shape[0]
if cfg.super_resolution_mode != 0:
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, predictor_input_shape[0:2] )
predicted = predictor_func (predictor_input_bgr)
if isinstance(predicted, tuple):
#converter return bgr,mask
prd_face_bgr = np.clip (predicted[0], 0, 1.0)
prd_face_mask_a_0 = np.clip (predicted[1], 0, 1.0)
predictor_masked = True
else:
#converter return bgr only, using dst mask
prd_face_bgr = np.clip (predicted, 0, 1.0 )
prd_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, predictor_input_shape[0:2] )
predictor_masked = False
if cfg.super_resolution_mode:
prd_face_bgr = cfg.superres_func(cfg.super_resolution_mode, prd_face_bgr)
if 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 '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 cfg.erode_mask_modifier != 0:
ero = int( lowest_len * ( 0.126 - lowest_len * 0.00004551365 ) * 0.01*cfg.erode_mask_modifier )
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 cfg.blur_mask_modifier > 0:
blur = int( lowest_len * 0.10 * 0.01*cfg.blur_mask_modifier )
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 'seamless' not in cfg.mode and cfg.color_transfer_mode != 0:
if cfg.color_transfer_mode == 1: #rct
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)
elif cfg.color_transfer_mode == 2: #lct
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)
elif cfg.color_transfer_mode == 3: #mkl
prd_face_bgr = imagelib.color_transfer_mkl (prd_face_bgr, dst_face_bgr)
elif cfg.color_transfer_mode == 4: #mkl-m
prd_face_bgr = imagelib.color_transfer_mkl (prd_face_bgr*prd_face_mask_a, dst_face_bgr*prd_face_mask_a)
elif cfg.color_transfer_mode == 5: #idt
prd_face_bgr = imagelib.color_transfer_idt (prd_face_bgr, dst_face_bgr)
elif cfg.color_transfer_mode == 6: #idt-m
prd_face_bgr = imagelib.color_transfer_idt (prd_face_bgr*prd_face_mask_a, dst_face_bgr*prd_face_mask_a)
elif cfg.color_transfer_mode == 7: #ebs
prd_face_bgr = cfg.ebs_ct_func ( np.clip( (dst_face_bgr*255), 0, 255).astype(np.uint8),
np.clip( (prd_face_bgr*255), 0, 255).astype(np.uint8), )#prd_face_mask_a
prd_face_bgr = np.clip( prd_face_bgr.astype(np.float32) / 255.0, 0.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':
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.mode == 'hist-match-bw':
prd_face_bgr = prd_face_bgr.astype(dtype=np.float32)
if 'seamless' in cfg.mode:
#mask used for cv2.seamlessClone
img_face_mask_a = img_face_mask_aaa[...,0:1]
if cfg.mode == 'seamless2':
img_face_mask_a = cv2.warpAffine( img_face_mask_a, face_output_mat, (output_size, output_size), flags=cv2.INTER_CUBIC )
img_face_seamless_mask_a = None
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_a.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
if cfg.mode == 'seamless2':
face_seamless = imagelib.seamless_clone ( prd_face_bgr, dst_face_bgr, img_face_seamless_mask_a )
out_img = cv2.warpAffine( face_seamless, face_output_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_CUBIC, cv2.BORDER_TRANSPARENT )
else:
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 'seamless' in cfg.mode and cfg.mode != 'seamless2':
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)
out_img = img_bgr*(1-img_face_mask_aaa) + (out_img*img_face_mask_aaa)
out_face_bgr = cv2.warpAffine( out_img, face_mat, (output_size, output_size) )
if 'seamless' in cfg.mode and cfg.color_transfer_mode != 0:
if cfg.color_transfer_mode == 1:
face_mask_aaa = cv2.warpAffine( img_face_mask_aaa, face_mat, (output_size, output_size) )
out_face_bgr = imagelib.reinhard_color_transfer ( np.clip( (out_face_bgr*255), 0, 255).astype(np.uint8),
np.clip( (dst_face_bgr*255), 0, 255).astype(np.uint8),
source_mask=face_mask_aaa, target_mask=face_mask_aaa)
out_face_bgr = np.clip( out_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0)
elif cfg.color_transfer_mode == 2: #lct
out_face_bgr = imagelib.linear_color_transfer (out_face_bgr, dst_face_bgr)
out_face_bgr = np.clip( out_face_bgr, 0.0, 1.0)
elif cfg.color_transfer_mode == 3: #mkl
out_face_bgr = imagelib.color_transfer_mkl (out_face_bgr, dst_face_bgr)
elif cfg.color_transfer_mode == 4: #mkl-m
out_face_bgr = imagelib.color_transfer_mkl (out_face_bgr*prd_face_mask_a, dst_face_bgr*prd_face_mask_a)
elif cfg.color_transfer_mode == 5: #idt
out_face_bgr = imagelib.color_transfer_idt (out_face_bgr, dst_face_bgr)
elif cfg.color_transfer_mode == 6: #idt-m
out_face_bgr = imagelib.color_transfer_idt (out_face_bgr*prd_face_mask_a, dst_face_bgr*prd_face_mask_a)
elif cfg.color_transfer_mode == 7: #ebs
out_face_bgr = cfg.ebs_ct_func ( np.clip( (dst_face_bgr*255), 0, 255).astype(np.uint8),
np.clip( (out_face_bgr*255), 0, 255).astype(np.uint8), )
out_face_bgr = np.clip( out_face_bgr.astype(np.float32) / 255.0, 0.0, 1.0)
if cfg.mode == 'seamless-hist-match':
out_face_bgr = imagelib.color_hist_match(out_face_bgr, dst_face_bgr, cfg.hist_match_threshold)
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_mode:
k_size *= 2
out_face_bgr = imagelib.LinearMotionBlur (out_face_bgr, k_size , frame_info.motion_deg)
if cfg.sharpen_mode != 0 and cfg.sharpen_amount != 0:
out_face_bgr = cfg.sharpen_func ( out_face_bgr, cfg.sharpen_mode, 3, cfg.sharpen_amount)
new_out = cv2.warpAffine( 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:
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
return out_img, out_merging_mask
def ConvertMasked (predictor_func, predictor_input_shape, 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 (predictor_func, predictor_input_shape, 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 Normal file
View file

@ -0,0 +1,50 @@
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

70
converters/ConverterAvatar.py Normal file
View file

@ -0,0 +1,70 @@
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,1), 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, img_bgr, img_face_landmarks, debug, avaperator_face_bgr=None, **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)
img_face_mask_aaa = np.repeat(img_face_mask_a, 3, -1)
output_size = self.predictor_input_size
face_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, output_size, face_type=FaceType.FULL)
dst_face_mask_a_0 = cv2.warpAffine( img_face_mask_a, face_mat, (output_size, output_size), flags=cv2.INTER_CUBIC )
predictor_input_dst_face_mask_a_0 = cv2.resize (dst_face_mask_a_0, (self.predictor_input_size,self.predictor_input_size), cv2.INTER_CUBIC )
prd_inp_dst_face_mask_a = predictor_input_dst_face_mask_a_0[...,np.newaxis]
prd_inp_avaperator_face_bgr = cv2.resize (avaperator_face_bgr, (self.predictor_input_size,self.predictor_input_size), cv2.INTER_CUBIC )
prd_face_bgr = self.predictor_func ( prd_inp_avaperator_face_bgr, prd_inp_dst_face_mask_a )
out_img = img_bgr.copy()
out_img = cv2.warpAffine( prd_face_bgr, face_mat, img_size, out_img, cv2.WARP_INVERSE_MAP | cv2.INTER_LANCZOS4, cv2.BORDER_TRANSPARENT )
out_img = np.clip(out_img, 0.0, 1.0)
if debug:
debugs += [out_img.copy()]
out_img = np.clip( img_bgr*(1-img_face_mask_aaa) + (out_img*img_face_mask_aaa) , 0, 1.0 )
if debug:
debugs += [out_img.copy()]
return debugs if debug else out_img

325
converters/ConverterConfig.py
View file

@ -1,325 +0,0 @@
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):
self.type = type
self.superres_func = None
self.sharpen_func = None
self.fanseg_input_size = None
self.fanseg_extract_func = None
self.ebs_ct_func = None
self.super_res_dict = {0:"None", 1:'RankSRGAN'}
self.sharpen_dict = {0:"None", 1:'box', 2:'gaussian'}
#default changeable params
self.super_resolution_mode = 0
self.sharpen_mode = 0
self.sharpen_amount = 0
def copy(self):
return copy.copy(self)
#overridable
def ask_settings(self):
s = """Choose sharpen mode: \n"""
for key in self.sharpen_dict.keys():
s += f"""({key}) {self.sharpen_dict[key]}\n"""
s += f"""?:help Default: {list(self.sharpen_dict.keys())[0]} : """
self.sharpen_mode = io.input_int (s, 0, valid_list=self.sharpen_dict.keys(), help_message="Enhance details by applying sharpen filter.")
if self.sharpen_mode != 0:
self.sharpen_amount = np.clip ( io.input_int ("Choose sharpen amount [0..100] (skip:%d) : " % 10, 10), 0, 100 )
s = """Choose super resolution mode: \n"""
for key in self.super_res_dict.keys():
s += f"""({key}) {self.super_res_dict[key]}\n"""
s += f"""?:help Default: {list(self.super_res_dict.keys())[0]} : """
self.super_resolution_mode = io.input_int (s, 0, valid_list=self.super_res_dict.keys(), help_message="Enhance details by applying superresolution network.")
def toggle_sharpen_mode(self):
a = list( self.sharpen_dict.keys() )
self.sharpen_mode = a[ (a.index(self.sharpen_mode)+1) % len(a) ]
def add_sharpen_amount(self, diff):
self.sharpen_amount = np.clip ( self.sharpen_amount+diff, 0, 100)
def toggle_super_resolution_mode(self):
a = list( self.super_res_dict.keys() )
self.super_resolution_mode = a[ (a.index(self.super_resolution_mode)+1) % len(a) ]
#overridable
def __eq__(self, other):
#check equality of changeable params
if isinstance(other, ConverterConfig):
return self.sharpen_mode == other.sharpen_mode and \
(self.sharpen_mode == 0 or ((self.sharpen_mode == other.sharpen_mode) and (self.sharpen_amount == other.sharpen_amount) )) and \
self.super_resolution_mode == other.super_resolution_mode
return False
#overridable
def to_string(self, filename):
r = ""
r += f"sharpen_mode : {self.sharpen_dict[self.sharpen_mode]}\n"
if self.sharpen_mode != 0:
r += f"sharpen_amount : {self.sharpen_amount}\n"
r += f"super_resolution_mode : {self.super_res_dict[self.super_resolution_mode]}\n"
return r
mode_dict = {0:'original',
1:'overlay',
2:'hist-match',
3:'seamless2',
4:'seamless',
5:'seamless-hist-match',
6:'raw-rgb',
7:'raw-rgb-mask',
8:'raw-mask-only',
9:'raw-predicted-only'}
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'}
half_face_mask_mode_dict = {1:'learned',
2:'dst',
4:'FAN-dst',
7:'learned*FAN-dst'}
ctm_dict = { 0: "None", 1:"rct", 2:"lct", 3:"mkl", 4:"mkl-m", 5:"idt", 6:"idt-m", 7:"ebs" }
ctm_str_dict = {None:0, "rct":1, "lct":2, "mkl":3, "mkl-m":4, "idt":5, "idt-m":6, "ebs":7 }
class ConverterConfigMasked(ConverterConfig):
def __init__(self, face_type=FaceType.FULL,
default_mode = 4,
clip_hborder_mask_per = 0,
):
super().__init__(type=ConverterConfig.TYPE_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.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.color_degrade_power = 0
self.export_mask_alpha = False
def copy(self):
return copy.copy(self)
def set_mode (self, mode):
self.mode = mode_dict.get (mode, 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( full_face_mask_mode_dict.keys() )
else:
a = list( 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 , -400, 400)
def add_blur_mask_modifier(self, diff):
self.blur_mask_modifier = np.clip ( self.blur_mask_modifier+diff , -400, 400)
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) % ( max(ctm_dict.keys())+1 )
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 mode_dict.keys():
s += f"""({key}) {mode_dict[key]}\n"""
s += f"""Default: {self.default_mode} : """
mode = io.input_int (s, self.default_mode)
self.mode = mode_dict.get (mode, 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 full_face_mask_mode_dict.keys():
s += f"""({key}) {full_face_mask_mode_dict[key]}\n"""
s += f"""?:help Default: 1 : """
self.mask_mode = io.input_int (s, 1, valid_list=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 half_face_mask_mode_dict.keys():
s += f"""({key}) {half_face_mask_mode_dict[key]}\n"""
s += f"""?:help , Default: 1 : """
self.mask_mode = io.input_int (s, 1, valid_list=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 = np.clip ( io.input_int ("Choose erode mask modifier [-400..400] (skip:%d) : " % 0, 0), -400, 400)
self.blur_mask_modifier = np.clip ( io.input_int ("Choose blur mask modifier [-400..400] (skip:%d) : " % 0, 0), -400, 400)
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/ebs skip:None ) : ", None, ctm_str_dict.keys() )
self.color_transfer_mode = ctm_str_dict[self.color_transfer_mode]
super().ask_settings()
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 super().__eq__(other) and \
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.color_degrade_power == other.color_degrade_power and \
self.export_mask_alpha == other.export_mask_alpha
return False
def to_string(self, filename):
r = (
f"""ConverterConfig {filename}:\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: { full_face_mask_mode_dict[self.mask_mode] }\n"""
else:
r += f"""mask_mode: { 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: { ctm_dict[self.color_transfer_mode]}\n"""
r += super().to_string(filename)
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, temporal_face_count=0):
super().__init__(type=ConverterConfig.TYPE_FACE_AVATAR)
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.")
super().ask_settings()
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 super().__eq__(other) and \
self.add_source_image == other.add_source_image
return False
#override
def to_string(self, filename):
return (f"ConverterConfig {filename}:\n"
f"add_source_image : {self.add_source_image}\n") + \
super().to_string(filename) + "================"

50
converters/ConverterImage.py Normal file
View file

@ -0,0 +1,50 @@
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

435
converters/ConverterMasked.py Normal file
View file

@ -0,0 +1,435 @@
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 = 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

9
converters/FrameInfo.py
View file

@ -1,9 +0,0 @@
from pathlib import Path
class FrameInfo(object):
def __init__(self, filename=None, landmarks_list=None):
self.filename = filename
self.filename_short = str(Path(filename).name)
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 .FrameInfo import FrameInfo
from .ConverterConfig import ConverterConfig, ConverterConfigMasked, ConverterConfigFaceAvatar
from .ConvertMasked import ConvertMasked
from .ConvertAvatar import ConvertFaceAvatar
from .Converter import Converter
from .ConverterMasked import ConverterMasked
from .ConverterImage import ConverterImage
from .ConverterAvatar import ConverterAvatar

2
doc/doc_prebuilt_windows_app.md
View file

@ -2,8 +2,6 @@
Windows builds with all dependencies included are released regularly. Only the NVIDIA GeForce display driver needs to be installed. Prebuilt DeepFaceLab, including GPU and CPU versions, can be downloaded from
[Google drive](https://drive.google.com/open?id=1BCFK_L7lPNwMbEQ_kFPqPpDdFEOd_Dci)
[Mega](https://mega.nz/#F!b9MzCK4B!zEAG9txu7uaRUjXz9PtBqg)
Available builds:

BIN
doc/gallery/1.jpg
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 182 KiB

BIN
doc/gallery/2.jpg
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 53 KiB

3
doc/gallery/doc_gallery.md
View file

@ -1,3 +0,0 @@
![](1.jpg)
![](2.jpg)

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.

1
ebsynth/__init__.py
View file

@ -1 +0,0 @@
from .ebsynth import color_transfer

BIN
ebsynth/ebsynth.dll
View file

Binary file not shown.

201
ebsynth/ebsynth.py
View file

@ -1,201 +0,0 @@
import os
import sys
from ctypes import *
from pathlib import Path
import cv2
import numpy as np
libebsynth = None
cached_buffer = {}
EBSYNTH_BACKEND_CPU = 0x0001
EBSYNTH_BACKEND_CUDA = 0x0002
EBSYNTH_BACKEND_AUTO = 0x0000
EBSYNTH_MAX_STYLE_CHANNELS = 8
EBSYNTH_MAX_GUIDE_CHANNELS = 24
EBSYNTH_VOTEMODE_PLAIN = 0x0001 # weight = 1
EBSYNTH_VOTEMODE_WEIGHTED = 0x0002 # weight = 1/(1+error)
def _normalize_img_shape (img):
img_len = len(img.shape)
if img_len == 2:
sh, sw = img.shape
sc = 0
elif img_len == 3:
sh, sw, sc = img.shape
if sc == 0:
sc = 1
img = img [...,np.newaxis]
return img
def run (img_style, guides,
patch_size=5,
num_pyramid_levels=-1,
num_search_vote_iters = 6,
num_patch_match_iters = 4,
stop_threshold = 5,
uniformity_weight = 3500.0,
extraPass3x3 = False,
):
if patch_size < 3:
raise ValueError ("patch_size is too small")
if patch_size % 2 == 0:
raise ValueError ("patch_size must be an odd number")
if len(guides) == 0:
raise ValueError ("at least one guide must be specified")
global libebsynth
if libebsynth is None:
if sys.platform[0:3] == 'win':
libebsynth_path = str ( Path(__file__).parent / 'ebsynth.dll' )
libebsynth = CDLL(libebsynth_path)
else:
#todo: implement for linux
pass
if libebsynth is not None:
libebsynth.ebsynthRun.argtypes = ( \
c_int,
c_int,
c_int,
c_int,
c_int,
c_void_p,
c_void_p,
c_int,
c_int,
c_void_p,
c_void_p,
POINTER(c_float),
POINTER(c_float),
c_float,
c_int,
c_int,
c_int,
POINTER(c_int),
POINTER(c_int),
POINTER(c_int),
c_int,
c_void_p,
c_void_p
)
if libebsynth is None:
return img_style
img_style = _normalize_img_shape (img_style)
sh, sw, sc = img_style.shape
t_h, t_w, t_c = 0,0,0
if sc > EBSYNTH_MAX_STYLE_CHANNELS:
raise ValueError (f"error: too many style channels {sc}, maximum number is {EBSYNTH_MAX_STYLE_CHANNELS}")
guides_source = []
guides_target = []
guides_weights = []
for i in range(len(guides)):
source_guide, target_guide, guide_weight = guides[i]
source_guide = _normalize_img_shape(source_guide)
target_guide = _normalize_img_shape(target_guide)
s_h, s_w, s_c = source_guide.shape
nt_h, nt_w, nt_c = target_guide.shape
if s_h != sh or s_w != sw:
raise ValueError ("guide source and style resolution must match style resolution.")
if t_c == 0:
t_h, t_w, t_c = nt_h, nt_w, nt_c
elif nt_h != t_h or nt_w != t_w:
raise ValueError ("guides target resolutions must be equal")
if s_c != nt_c:
raise ValueError ("guide source and target channels must match exactly.")
guides_source.append (source_guide)
guides_target.append (target_guide)
guides_weights += [ guide_weight / s_c ] * s_c
guides_source = np.concatenate ( guides_source, axis=-1)
guides_target = np.concatenate ( guides_target, axis=-1)
guides_weights = (c_float*len(guides_weights) ) ( *guides_weights )
styleWeight = 1.0
style_weights = [ styleWeight / sc for i in range(sc) ]
style_weights = (c_float*sc) ( *style_weights )
maxPyramidLevels = 0
for level in range(32,-1,-1):
if min( min(sh, t_h)*pow(2.0, -level), \
min(sw, t_w)*pow(2.0, -level) ) >= (2*patch_size+1):
maxPyramidLevels = level+1
break
if num_pyramid_levels == -1:
num_pyramid_levels = maxPyramidLevels
num_pyramid_levels = min(num_pyramid_levels, maxPyramidLevels)
num_search_vote_iters_per_level = (c_int*num_pyramid_levels) ( *[num_search_vote_iters]*num_pyramid_levels )
num_patch_match_iters_per_level = (c_int*num_pyramid_levels) ( *[num_patch_match_iters]*num_pyramid_levels )
stop_threshold_per_level = (c_int*num_pyramid_levels) ( *[stop_threshold]*num_pyramid_levels )
buffer = cached_buffer.get ( (t_h,t_w,sc), None )
if buffer is None:
buffer = create_string_buffer (t_h*t_w*sc)
cached_buffer[(t_h,t_w,sc)] = buffer
libebsynth.ebsynthRun (EBSYNTH_BACKEND_CPU, #backend
sc, #numStyleChannels
guides_source.shape[-1], #numGuideChannels
sw, #sourceWidth
sh, #sourceHeight
img_style.tobytes(), #sourceStyleData (width * height * numStyleChannels) bytes, scan-line order
guides_source.tobytes(), #sourceGuideData (width * height * numGuideChannels) bytes, scan-line order
t_w, #targetWidth
t_h, #targetHeight
guides_target.tobytes(), #targetGuideData (width * height * numGuideChannels) bytes, scan-line order
None, #targetModulationData (width * height * numGuideChannels) bytes, scan-line order; pass NULL to switch off the modulation
style_weights, #styleWeights (numStyleChannels) floats
guides_weights, #guideWeights (numGuideChannels) floats
uniformity_weight, #uniformityWeight reasonable values are between 500-15000, 3500 is a good default
patch_size, #patchSize odd sizes only, use 5 for 5x5 patch, 7 for 7x7, etc.
EBSYNTH_VOTEMODE_PLAIN, #voteMode use VOTEMODE_WEIGHTED for sharper result
num_pyramid_levels, #numPyramidLevels
num_search_vote_iters_per_level, #numSearchVoteItersPerLevel how many search/vote iters to perform at each level (array of ints, coarse first, fine last)
num_patch_match_iters_per_level, #numPatchMatchItersPerLevel how many Patch-Match iters to perform at each level (array of ints, coarse first, fine last)
stop_threshold_per_level, #stopThresholdPerLevel stop improving pixel when its change since last iteration falls under this threshold
1 if extraPass3x3 else 0, #extraPass3x3 perform additional polishing pass with 3x3 patches at the finest level, use 0 to disable
None, #outputNnfData (width * height * 2) ints, scan-line order; pass NULL to ignore
buffer #outputImageData (width * height * numStyleChannels) bytes, scan-line order
)
return np.frombuffer(buffer, dtype=np.uint8).reshape ( (t_h,t_w,sc) ).copy()
#transfer color from source to target
def color_transfer(img_source, img_target):
guides = [( cv2.cvtColor(img_source, cv2.COLOR_BGR2GRAY),
cv2.cvtColor(img_target, cv2.COLOR_BGR2GRAY),
1 ) ]
h,w,c = img_source.shape
result = []
for i in range(c):
result += [
run( img_source[...,i:i+1] , guides=guides,
patch_size=11,
num_pyramid_levels=40,
num_search_vote_iters = 6,
num_patch_match_iters = 4,
stop_threshold = 5,
uniformity_weight=500.0,
extraPass3x3=True,
)
]
return np.concatenate( result, axis=-1 )

BIN
facelib/FANSeg_256_full_face.h5
View file

Binary file not shown.

69
facelib/FANSegmentator.py
View file

@ -47,30 +47,10 @@ class FANSegmentator(object):
self.model.get_layer (s).set_weights ( d[s] )
except:
io.log_err("Unable to load VGG11 pretrained weights from vgg11_enc_weights.npy")
conv_weights_list = []
for layer in self.model.layers:
if 'CA.' in layer.name:
conv_weights_list += [layer.weights[0]] #Conv2D kernel_weights
CAInitializerMP ( conv_weights_list )
if training:
inp_t = Input ( (resolution, resolution, 3) )
real_t = Input ( (resolution, resolution, 1) )
out_t = self.model(inp_t)
#loss = K.mean(10*K.square(out_t-real_t))
loss = K.mean(10*K.binary_crossentropy(real_t,out_t) )
out_t_diff1 = out_t[:, 1:, :, :] - out_t[:, :-1, :, :]
out_t_diff2 = out_t[:, :, 1:, :] - out_t[:, :, :-1, :]
total_var_loss = K.mean( 0.1*K.abs(out_t_diff1), axis=[1, 2, 3] ) + K.mean( 0.1*K.abs(out_t_diff2), axis=[1, 2, 3] )
opt = Adam(lr=0.0001, beta_1=0.5, beta_2=0.999, tf_cpu_mode=2)
self.train_func = K.function ( [inp_t, real_t], [K.mean(loss)], opt.get_updates( [loss,total_var_loss], self.model.trainable_weights) )
#self.model.compile(loss='mse', optimizer=Adam(tf_cpu_mode=2))
self.model.compile(loss='binary_crossentropy', optimizer=Adam(tf_cpu_mode=2) )
def __enter__(self):
return self
@ -81,9 +61,8 @@ class FANSegmentator(object):
def save_weights(self):
self.model.save_weights (str(self.weights_path))
def train(self, inp, real):
loss, = self.train_func ([inp, real])
return loss
def train_on_batch(self, inp, outp):
return self.model.train_on_batch(inp, outp)
def extract (self, input_image, is_input_tanh=False):
input_shape_len = len(input_image.shape)
@ -99,62 +78,62 @@ class FANSegmentator(object):
return result
@staticmethod
def BuildModel ( resolution, ngf=64):
def BuildModel ( resolution, ngf=64, norm='', act='lrelu'):
exec( nnlib.import_all(), locals(), globals() )
inp = Input ( (resolution,resolution,3) )
x = inp
x = FANSegmentator.Flow(ngf=ngf)(x)
x = FANSegmentator.Flow(ngf=ngf, norm=norm, act=act)(x)
model = Model(inp,x)
return model
@staticmethod
def Flow(ngf=64):
def Flow(ngf=64, num_downs=4, norm='', act='lrelu'):
exec( nnlib.import_all(), locals(), globals() )
def func(input):
x = input
x0 = x = Conv2D(ngf, kernel_size=3, strides=1, padding='same', activation='relu', name='features.0')(x)
x = BlurPool(filt_size=3)(x) #x = MaxPooling2D()(x)
x = MaxPooling2D()(x)
x1 = x = Conv2D(ngf*2, kernel_size=3, strides=1, padding='same', activation='relu', name='features.3')(x)
x = BlurPool(filt_size=3)(x)
x = MaxPooling2D()(x)
x = Conv2D(ngf*4, kernel_size=3, strides=1, padding='same', activation='relu', name='features.6')(x)
x2 = x = Conv2D(ngf*4, kernel_size=3, strides=1, padding='same', activation='relu', name='features.8')(x)
x = BlurPool(filt_size=3)(x)
x = MaxPooling2D()(x)
x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.11')(x)
x3 = x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.13')(x)
x = BlurPool(filt_size=3)(x)
x = MaxPooling2D()(x)
x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.16')(x)
x4 = x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', activation='relu', name='features.18')(x)
x = BlurPool(filt_size=3)(x)
x = MaxPooling2D()(x)
x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same', name='CA.1')(x)
x = Conv2D(ngf*8, kernel_size=3, strides=1, padding='same')(x)
x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu', name='CA.2') (x)
x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu') (x)
x = Concatenate(axis=3)([ x, x4])
x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu', name='CA.3') (x)
x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu') (x)
x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu', name='CA.4') (x)
x = Conv2DTranspose (ngf*4, 3, strides=2, padding='same', activation='relu') (x)
x = Concatenate(axis=3)([ x, x3])
x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu', name='CA.5') (x)
x = Conv2D (ngf*8, 3, strides=1, padding='same', activation='relu') (x)
x = Conv2DTranspose (ngf*2, 3, strides=2, padding='same', activation='relu', name='CA.6') (x)
x = Conv2DTranspose (ngf*2, 3, strides=2, padding='same', activation='relu') (x)
x = Concatenate(axis=3)([ x, x2])
x = Conv2D (ngf*4, 3, strides=1, padding='same', activation='relu', name='CA.7') (x)
x = Conv2D (ngf*4, 3, strides=1, padding='same', activation='relu') (x)
x = Conv2DTranspose (ngf, 3, strides=2, padding='same', activation='relu', name='CA.8') (x)
x = Conv2DTranspose (ngf, 3, strides=2, padding='same', activation='relu') (x)
x = Concatenate(axis=3)([ x, x1])
x = Conv2D (ngf*2, 3, strides=1, padding='same', activation='relu', name='CA.9') (x)
x = Conv2D (ngf*2, 3, strides=1, padding='same', activation='relu') (x)
x = Conv2DTranspose (ngf // 2, 3, strides=2, padding='same', activation='relu', name='CA.10') (x)
x = Conv2DTranspose (ngf // 2, 3, strides=2, padding='same', activation='relu') (x)
x = Concatenate(axis=3)([ x, x0])
x = Conv2D (ngf, 3, strides=1, padding='same', activation='relu', name='CA.11') (x)
x = Conv2D (ngf, 3, strides=1, padding='same', activation='relu') (x)
return Conv2D(1, 3, strides=1, padding='same', activation='sigmoid', name='CA.12')(x)
return Conv2D(1, 3, strides=1, padding='same', activation='sigmoid')(x)
return func

28
facelib/FaceType.py
View file

@ -1,14 +1,12 @@
from enum import IntEnum
class FaceType(IntEnum):
#enumerating in order "next contains prev"
HALF = 0,
FULL = 1,
FULL_NO_ALIGN = 3,
HEAD = 4,
HEAD_NO_ALIGN = 5,
MARK_ONLY = 10, #no align at all, just embedded faceinfo
HEAD = 2,
AVATAR = 3, #centered nose only
MARK_ONLY = 4, #no align at all, just embedded faceinfo
QTY = 5
@staticmethod
def fromString (s):
@ -19,19 +17,17 @@ class FaceType(IntEnum):
@staticmethod
def toString (face_type):
return to_string_dict[face_type]
return to_string_list[face_type]
from_string_dict = {'half_face': FaceType.HALF,
'full_face': FaceType.FULL,
'head' : FaceType.HEAD,
'avatar' : FaceType.AVATAR,
'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.HEAD_NO_ALIGN : 'head_no_align'
}
to_string_list = [ 'half_face',
'full_face',
'head',
'avatar',
'mark_only'
]

427
facelib/LandmarksProcessor.py
View file

@ -1,107 +1,33 @@
import colorsys
import math
from enum import IntEnum
import cv2
import numpy as np
import numpy.linalg as npla
import imagelib
from enum import IntEnum
import mathlib
from facelib import FaceType
import imagelib
from imagelib import IEPolys
from mathlib.umeyama import umeyama
from facelib import FaceType
import math
landmarks_2D = np.array([
[ 0.000213256, 0.106454 ], #17
[ 0.0752622, 0.038915 ], #18
[ 0.18113, 0.0187482 ], #19
[ 0.29077, 0.0344891 ], #20
[ 0.393397, 0.0773906 ], #21
[ 0.586856, 0.0773906 ], #22
[ 0.689483, 0.0344891 ], #23
[ 0.799124, 0.0187482 ], #24
[ 0.904991, 0.038915 ], #25
[ 0.98004, 0.106454 ], #26
[ 0.490127, 0.203352 ], #27
[ 0.490127, 0.307009 ], #28
[ 0.490127, 0.409805 ], #29
[ 0.490127, 0.515625 ], #30
[ 0.36688, 0.587326 ], #31
[ 0.426036, 0.609345 ], #32
[ 0.490127, 0.628106 ], #33
[ 0.554217, 0.609345 ], #34
[ 0.613373, 0.587326 ], #35
[ 0.121737, 0.216423 ], #36
[ 0.187122, 0.178758 ], #37
[ 0.265825, 0.179852 ], #38
[ 0.334606, 0.231733 ], #39
[ 0.260918, 0.245099 ], #40
[ 0.182743, 0.244077 ], #41
[ 0.645647, 0.231733 ], #42
[ 0.714428, 0.179852 ], #43
[ 0.793132, 0.178758 ], #44
[ 0.858516, 0.216423 ], #45
[ 0.79751, 0.244077 ], #46
[ 0.719335, 0.245099 ], #47
[ 0.254149, 0.780233 ], #48
[ 0.340985, 0.745405 ], #49
[ 0.428858, 0.727388 ], #50
[ 0.490127, 0.742578 ], #51
[ 0.551395, 0.727388 ], #52
[ 0.639268, 0.745405 ], #53
[ 0.726104, 0.780233 ], #54
[ 0.642159, 0.864805 ], #55
[ 0.556721, 0.902192 ], #56
[ 0.490127, 0.909281 ], #57
[ 0.423532, 0.902192 ], #58
[ 0.338094, 0.864805 ], #59
[ 0.290379, 0.784792 ], #60
[ 0.428096, 0.778746 ], #61
[ 0.490127, 0.785343 ], #62
[ 0.552157, 0.778746 ], #63
[ 0.689874, 0.784792 ], #64
[ 0.553364, 0.824182 ], #65
[ 0.490127, 0.831803 ], #66
[ 0.42689 , 0.824182 ] #67
], dtype=np.float32)
mean_face_x = np.array([
0.000213256, 0.0752622, 0.18113, 0.29077, 0.393397, 0.586856, 0.689483, 0.799124,
0.904991, 0.98004, 0.490127, 0.490127, 0.490127, 0.490127, 0.36688, 0.426036,
0.490127, 0.554217, 0.613373, 0.121737, 0.187122, 0.265825, 0.334606, 0.260918,
0.182743, 0.645647, 0.714428, 0.793132, 0.858516, 0.79751, 0.719335, 0.254149,
0.340985, 0.428858, 0.490127, 0.551395, 0.639268, 0.726104, 0.642159, 0.556721,
0.490127, 0.423532, 0.338094, 0.290379, 0.428096, 0.490127, 0.552157, 0.689874,
0.553364, 0.490127, 0.42689 ])
mean_face_y = np.array([
0.106454, 0.038915, 0.0187482, 0.0344891, 0.0773906, 0.0773906, 0.0344891,
0.0187482, 0.038915, 0.106454, 0.203352, 0.307009, 0.409805, 0.515625, 0.587326,
0.609345, 0.628106, 0.609345, 0.587326, 0.216423, 0.178758, 0.179852, 0.231733,
0.245099, 0.244077, 0.231733, 0.179852, 0.178758, 0.216423, 0.244077, 0.245099,
0.780233, 0.745405, 0.727388, 0.742578, 0.727388, 0.745405, 0.780233, 0.864805,
0.902192, 0.909281, 0.902192, 0.864805, 0.784792, 0.778746, 0.785343, 0.778746,
0.784792, 0.824182, 0.831803, 0.824182 ])
landmarks_2D_new = np.array([
[ 0.000213256, 0.106454 ], #17
[ 0.0752622, 0.038915 ], #18
[ 0.18113, 0.0187482 ], #19
[ 0.29077, 0.0344891 ], #20
[ 0.393397, 0.0773906 ], #21
[ 0.586856, 0.0773906 ], #22
[ 0.689483, 0.0344891 ], #23
[ 0.799124, 0.0187482 ], #24
[ 0.904991, 0.038915 ], #25
[ 0.98004, 0.106454 ], #26
[ 0.490127, 0.203352 ], #27
[ 0.490127, 0.307009 ], #28
[ 0.490127, 0.409805 ], #29
[ 0.490127, 0.515625 ], #30
[ 0.36688, 0.587326 ], #31
[ 0.426036, 0.609345 ], #32
[ 0.490127, 0.628106 ], #33
[ 0.554217, 0.609345 ], #34
[ 0.613373, 0.587326 ], #35
[ 0.121737, 0.216423 ], #36
[ 0.187122, 0.178758 ], #37
[ 0.265825, 0.179852 ], #38
[ 0.334606, 0.231733 ], #39
[ 0.260918, 0.245099 ], #40
[ 0.182743, 0.244077 ], #41
[ 0.645647, 0.231733 ], #42
[ 0.714428, 0.179852 ], #43
[ 0.793132, 0.178758 ], #44
[ 0.858516, 0.216423 ], #45
[ 0.79751, 0.244077 ], #46
[ 0.719335, 0.245099 ], #47
[ 0.254149, 0.780233 ], #48
[ 0.726104, 0.780233 ], #54
], dtype=np.float32)
landmarks_2D = np.stack( [ mean_face_x, mean_face_y ], axis=1 )
# 68 point landmark definitions
landmarks_68_pt = { "mouth": (48,68),
@ -183,19 +109,10 @@ landmarks_68_3D = np.array( [
[0.205322 , 31.408738 , -21.903670 ],
[-7.198266 , 30.844876 , -20.328022 ] ], dtype=np.float32)
def transform_points(points, mat, invert=False):
if invert:
mat = cv2.invertAffineTransform (mat)
points = np.expand_dims(points, axis=1)
points = cv2.transform(points, mat, points.shape)
points = np.squeeze(points)
return points
def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0):
if not isinstance(image_landmarks, np.ndarray):
image_landmarks = np.array (image_landmarks)
"""
if face_type == FaceType.AVATAR:
centroid = np.mean (image_landmarks, axis=0)
@ -211,262 +128,82 @@ def get_transform_mat (image_landmarks, output_size, face_type, scale=1.0):
mat = mat * scale * (output_size / 3)
mat[:,2] += output_size / 2
else:
"""
remove_align = False
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
else:
raise ValueError ('wrong face_type: ', face_type)
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) * 21
else:
raise ValueError ('wrong face_type: ', face_type)
#mat = umeyama(image_landmarks[17:], landmarks_2D, True)[0:2]
mat = umeyama( np.concatenate ( [ image_landmarks[17:49] , image_landmarks[54:55] ] ) , landmarks_2D_new, True)[0:2]
mat = mat * (output_size - 2 * padding)
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)
mat = umeyama(image_landmarks[17:], landmarks_2D, True)[0:2]
mat = mat * (output_size - 2 * padding)
mat[:,2] += padding
mat *= (1 / scale)
mat[:,2] += -output_size*( ( (1 / scale) - 1.0 ) / 2 )
return mat
def expand_eyebrows(lmrks, eyebrows_expand_mod=1.0):
if len(lmrks) != 68:
raise Exception('works only with 68 landmarks')
lmrks = np.array( lmrks.copy(), dtype=np.int )
# #nose
ml_pnt = (lmrks[36] + lmrks[0]) // 2
mr_pnt = (lmrks[16] + lmrks[45]) // 2
# mid points between the mid points and eye
ql_pnt = (lmrks[36] + ml_pnt) // 2
qr_pnt = (lmrks[45] + mr_pnt) // 2
# Top of the eye arrays
bot_l = np.array((ql_pnt, lmrks[36], lmrks[37], lmrks[38], lmrks[39]))
bot_r = np.array((lmrks[42], lmrks[43], lmrks[44], lmrks[45], qr_pnt))
# Eyebrow arrays
top_l = lmrks[17:22]
top_r = lmrks[22:27]
# Adjust eyebrow arrays
lmrks[17:22] = top_l + eyebrows_expand_mod * 0.5 * (top_l - bot_l)
lmrks[22:27] = top_r + eyebrows_expand_mod * 0.5 * (top_r - bot_r)
return lmrks
def transform_points(points, mat, invert=False):
if invert:
mat = cv2.invertAffineTransform (mat)
points = np.expand_dims(points, axis=1)
points = cv2.transform(points, mat, points.shape)
points = np.squeeze(points)
return points
def get_image_hull_mask (image_shape, image_landmarks, ie_polys=None):
if len(image_landmarks) != 68:
raise Exception('get_image_hull_mask works only with 68 landmarks')
int_lmrks = np.array(image_landmarks, dtype=np.int)
hull_mask = np.zeros(image_shape[0:2]+(1,),dtype=np.float32)
def get_image_hull_mask (image_shape, image_landmarks, eyebrows_expand_mod=1.0, ie_polys=None, color=(1,) ):
hull_mask = np.zeros(image_shape[0:2]+( len(color),),dtype=np.float32)
cv2.fillConvexPoly( hull_mask, cv2.convexHull(
np.concatenate ( (int_lmrks[0:9],
int_lmrks[17:18]))) , (1,) )
lmrks = expand_eyebrows(image_landmarks, eyebrows_expand_mod)
cv2.fillConvexPoly( hull_mask, cv2.convexHull(
np.concatenate ( (int_lmrks[8:17],
int_lmrks[26:27]))) , (1,) )
r_jaw = (lmrks[0:9], lmrks[17:18])
l_jaw = (lmrks[8:17], lmrks[26:27])
r_cheek = (lmrks[17:20], lmrks[8:9])
l_cheek = (lmrks[24:27], lmrks[8:9])
nose_ridge = (lmrks[19:25], lmrks[8:9],)
r_eye = (lmrks[17:22], lmrks[27:28], lmrks[31:36], lmrks[8:9])
l_eye = (lmrks[22:27], lmrks[27:28], lmrks[31:36], lmrks[8:9])
nose = (lmrks[27:31], lmrks[31:36])
parts = [r_jaw, l_jaw, r_cheek, l_cheek, nose_ridge, r_eye, l_eye, nose]
cv2.fillConvexPoly( hull_mask, cv2.convexHull(
np.concatenate ( (int_lmrks[17:20],
int_lmrks[8:9]))) , (1,) )
for item in parts:
merged = np.concatenate(item)
cv2.fillConvexPoly(hull_mask, cv2.convexHull(merged), color )
cv2.fillConvexPoly( hull_mask, cv2.convexHull(
np.concatenate ( (int_lmrks[24:27],
int_lmrks[8:9]))) , (1,) )
cv2.fillConvexPoly( hull_mask, cv2.convexHull(
np.concatenate ( (int_lmrks[19:25],
int_lmrks[8:9],
))) , (1,) )
cv2.fillConvexPoly( hull_mask, cv2.convexHull(
np.concatenate ( (int_lmrks[17:22],
int_lmrks[27:28],
int_lmrks[31:36],
int_lmrks[8:9]
))) , (1,) )
cv2.fillConvexPoly( hull_mask, cv2.convexHull(
np.concatenate ( (int_lmrks[22:27],
int_lmrks[27:28],
int_lmrks[31:36],
int_lmrks[8:9]
))) , (1,) )
#nose
cv2.fillConvexPoly( hull_mask, cv2.convexHull(int_lmrks[27:36]), (1,) )
if ie_polys is not None:
ie_polys.overlay_mask(hull_mask)
return hull_mask
def alpha_to_color (img_alpha, color):
if len(img_alpha.shape) == 2:
img_alpha = img_alpha[...,None]
h,w,c = img_alpha.shape
result = np.zeros( (h,w, len(color) ), dtype=np.float32 )
result[:,:] = color
return result * img_alpha
def get_cmask (image_shape, lmrks, eyebrows_expand_mod=1.0):
h,w,c = image_shape
hull = get_image_hull_mask (image_shape, lmrks, eyebrows_expand_mod, color=(1,) )
result = np.zeros( (h,w,3), dtype=np.float32 )
def process(w,h, data ):
d = {}
cur_lc = 0
all_lines = []
for s, pts_loop_ar in data:
lines = []
for pts, loop in pts_loop_ar:
pts_len = len(pts)
lines.append ( [ [ pts[i], pts[(i+1) % pts_len ] ] for i in range(pts_len - (0 if loop else 1) ) ] )
lines = np.concatenate (lines)
lc = lines.shape[0]
all_lines.append(lines)
d[s] = cur_lc, cur_lc+lc
cur_lc += lc
all_lines = np.concatenate (all_lines, 0)
#calculate signed distance for all points and lines
line_count = all_lines.shape[0]
pts_count = w*h
all_lines = np.repeat ( all_lines[None,...], pts_count, axis=0 ).reshape ( (pts_count*line_count,2,2) )
pts = np.empty( (h,w,line_count,2), dtype=np.float32 )
pts[...,1] = np.arange(h)[:,None,None]
pts[...,0] = np.arange(w)[:,None]
pts = pts.reshape ( (h*w*line_count, -1) )
a = all_lines[:,0,:]
b = all_lines[:,1,:]
pa = pts-a
ba = b-a
ph = np.clip ( np.einsum('ij,ij->i', pa, ba) / np.einsum('ij,ij->i', ba, ba), 0, 1 )
dists = npla.norm ( pa - ba*ph[...,None], axis=1).reshape ( (h,w,line_count) )
def get_dists(name, thickness=0):
s,e = d[name]
result = dists[...,s:e]
if thickness != 0:
result = np.abs(result)-thickness
return np.min (result, axis=-1)
return get_dists
l_eye = lmrks[42:48]
r_eye = lmrks[36:42]
l_brow = lmrks[22:27]
r_brow = lmrks[17:22]
mouth = lmrks[48:60]
up_nose = np.concatenate( (lmrks[27:31], lmrks[33:34]) )
down_nose = lmrks[31:36]
nose = np.concatenate ( (up_nose, down_nose) )
gdf = process ( w,h,
(
('eyes', ((l_eye, True), (r_eye, True)) ),
('brows', ((l_brow, False), (r_brow,False)) ),
('up_nose', ((up_nose, False),) ),
('down_nose', ((down_nose, False),) ),
('mouth', ((mouth, True),) ),
)
)
#import code
#code.interact(local=dict(globals(), **locals()))
eyes_fall_dist = w // 32
eyes_thickness = max( w // 64, 1 )
brows_fall_dist = w // 32
brows_thickness = max( w // 256, 1 )
nose_fall_dist = w / 12
nose_thickness = max( w // 96, 1 )
mouth_fall_dist = w // 32
mouth_thickness = max( w // 64, 1 )
eyes_mask = gdf('eyes',eyes_thickness)
eyes_mask = 1-np.clip( eyes_mask/ eyes_fall_dist, 0, 1)
#eyes_mask = np.clip ( 1- ( np.sqrt( np.maximum(eyes_mask,0) ) / eyes_fall_dist ), 0, 1)
#eyes_mask = np.clip ( 1- ( np.cbrt( np.maximum(eyes_mask,0) ) / eyes_fall_dist ), 0, 1)
brows_mask = gdf('brows', brows_thickness)
brows_mask = 1-np.clip( brows_mask / brows_fall_dist, 0, 1)
#brows_mask = np.clip ( 1- ( np.sqrt( np.maximum(brows_mask,0) ) / brows_fall_dist ), 0, 1)
mouth_mask = gdf('mouth', mouth_thickness)
mouth_mask = 1-np.clip( mouth_mask / mouth_fall_dist, 0, 1)
#mouth_mask = np.clip ( 1- ( np.sqrt( np.maximum(mouth_mask,0) ) / mouth_fall_dist ), 0, 1)
def blend(a,b,k):
x = np.clip ( 0.5+0.5*(b-a)/k, 0.0, 1.0 )
return (a-b)*x+b - k*x*(1.0-x)
#nose_mask = (a-b)*x+b - k*x*(1.0-x)
#nose_mask = np.minimum (up_nose_mask , down_nose_mask )
#nose_mask = 1-np.clip( nose_mask / nose_fall_dist, 0, 1)
nose_mask = blend ( gdf('up_nose', nose_thickness), gdf('down_nose', nose_thickness), nose_thickness*3 )
nose_mask = 1-np.clip( nose_mask / nose_fall_dist, 0, 1)
up_nose_mask = gdf('up_nose', nose_thickness)
up_nose_mask = 1-np.clip( up_nose_mask / nose_fall_dist, 0, 1)
#up_nose_mask = np.clip ( 1- ( np.cbrt( np.maximum(up_nose_mask,0) ) / nose_fall_dist ), 0, 1)
down_nose_mask = gdf('down_nose', nose_thickness)
down_nose_mask = 1-np.clip( down_nose_mask / nose_fall_dist, 0, 1)
#down_nose_mask = np.clip ( 1- ( np.cbrt( np.maximum(down_nose_mask,0) ) / nose_fall_dist ), 0, 1)
#nose_mask = np.clip( up_nose_mask + down_nose_mask, 0, 1 )
#nose_mask /= np.max(nose_mask)
#nose_mask = np.maximum (up_nose_mask , down_nose_mask )
#nose_mask = down_nose_mask
#nose_mask = np.zeros_like(nose_mask)
eyes_mask = eyes_mask * (1-mouth_mask)
nose_mask = nose_mask * (1-eyes_mask)
hull_mask = hull[...,0].copy()
hull_mask = hull_mask * (1-eyes_mask) * (1-brows_mask) * (1-nose_mask) * (1-mouth_mask)
#eyes_mask = eyes_mask * (1-nose_mask)
mouth_mask= mouth_mask * (1-nose_mask)
brows_mask = brows_mask * (1-nose_mask)* (1-eyes_mask )
hull_mask = alpha_to_color(hull_mask, (0,1,0) )
eyes_mask = alpha_to_color(eyes_mask, (1,0,0) )
brows_mask = alpha_to_color(brows_mask, (0,0,1) )
nose_mask = alpha_to_color(nose_mask, (0,1,1) )
mouth_mask = alpha_to_color(mouth_mask, (0,0,1) )
#nose_mask = np.maximum( up_nose_mask, down_nose_mask )
result = hull_mask + mouth_mask+ nose_mask + brows_mask + eyes_mask
result *= hull
#result = np.clip (result, 0, 1)
return result
def get_image_eye_mask (image_shape, image_landmarks):
if len(image_landmarks) != 68:
raise Exception('get_image_eye_mask works only with 68 landmarks')
@ -569,10 +306,10 @@ def draw_landmarks (image, image_landmarks, color=(0,255,0), transparent_mask=Fa
cv2.circle(image, (x, y), 2, color, lineType=cv2.LINE_AA)
if transparent_mask:
mask = get_image_hull_mask (image.shape, image_landmarks, ie_polys=ie_polys)
mask = get_image_hull_mask (image.shape, image_landmarks, ie_polys)
image[...] = ( image * (1-mask) + image * mask / 2 )[...]
def draw_rect_landmarks (image, rect, image_landmarks, face_size, face_type, transparent_mask=False, ie_polys=None, landmarks_color=(0,255,0)):
def draw_rect_landmarks (image, rect, image_landmarks, face_size, face_type, transparent_mask=False, ie_polys=None, landmarks_color=(0,255,0) ):
draw_landmarks(image, image_landmarks, color=landmarks_color, transparent_mask=transparent_mask, ie_polys=ie_polys)
imagelib.draw_rect (image, rect, (255,0,0), 2 )

2
facelib/MTCExtractor.py
View file

@ -291,7 +291,7 @@ def nms(boxes, threshold, method):
w = np.maximum(0.0, xx2-xx1+1)
h = np.maximum(0.0, yy2-yy1+1)
inter = w * h
if method == 'Min':
if method is 'Min':
o = inter / np.minimum(area[i], area[idx])
else:
o = inter / (area[i] + area[idx] - inter)

BIN
imagelib/DCSCN.h5 Normal file
View file

Binary file not shown.

164
imagelib/DCSCN.py Normal file
View file

@ -0,0 +1,164 @@
import numpy as np
import cv2
from pathlib import Path
from nnlib import nnlib
from interact import interact as io
class DCSCN():
def __init__(self):
exec( nnlib.import_all(), locals(), globals() )
inp_x = KL.Input([None, None, 1])
inp_x2 = KL.Input([None, None, 1])
x = inp_x
layers_count = 12
layers = []
for i in range(1,layers_count+1):
if i == 1:
output_feature_num = 196
else:
x1 = (i-1) / float(layers_count - 1)
y1 = x1 ** (1.0 / 1.5)
output_feature_num = int((196 - 48) * (1 - y1) + 48)
x = Conv2D(output_feature_num, kernel_size=3, strides=1, padding='same', name='CNN%d' % (i) ) (x)
x = PReLU(shared_axes=[1,2], name='CNN%d_prelu' % (i) ) (x)
layers.append(x)
x_concat = KL.Concatenate()(layers)
A1 = Conv2D(64, kernel_size=1, strides=1, padding='same', name='A1' ) (x_concat)
A1 = PReLU(shared_axes=[1,2], name='A1_prelu') (A1)
B1 = Conv2D(32, kernel_size=1, strides=1, padding='same', name='B1' ) (x_concat)
B1 = PReLU(shared_axes=[1,2], name='B1_prelu') (B1)
B2 = Conv2D(32, kernel_size=3, strides=1, padding='same', name='B2' ) (B1)
B2 = PReLU(shared_axes=[1,2], name='B2_prelu') (B2)
x = KL.Concatenate()([B2,A1])
x = Conv2D(96*4, kernel_size=3, strides=1, padding='same', name='Up_PS' )(x)
x = PixelShuffler()(x)
x = Conv2D(1, kernel_size=3, strides=1, padding='same', name='R_CNN1', use_bias=False )(x)
x = KL.Add()([x, inp_x2])
self.model = keras.models.Model ([inp_x, inp_x2], [x])
self.model.load_weights ( Path(__file__).parent / 'DCSCN.h5' )
def upscale(self, img, is_bgr=True, is_float=True):
if is_bgr:
img = img[...,::-1]
if is_float:
img = np.clip (img*255, 0, 255)
img_shape_len = len(img.shape)
h, w = img.shape[:2]
ch = img.shape[2] if len(img.shape) >= 3 else 1
nh, nw = h*2, w*2
img_x = self.convert_rgb_to_y(img)
img_bx = cv2.resize(img_x, (nh, nw), cv2.INTER_CUBIC)
ensemble = 8
output = np.zeros([nh,nw,1], dtype=np.float32)
for i in range(ensemble):
x = np.reshape( self.flip(img_x, i), (1,h,w,1) )
bx = np.reshape( self.flip(img_bx, i), (1,nh,nw,1) )
y = self.model.predict([x,bx])[0]
y = self.flip(y, i, invert=True)
output += y
output /= ensemble
bimg = cv2.resize(img, (nh, nw), cv2.INTER_CUBIC)
bimg_ycbcr = self.convert_rgb_to_ycbcr(bimg)
if ch > 1:
output = self.convert_y_and_cbcr_to_rgb(output, bimg_ycbcr[:, :, 1:3])
if is_float:
output = np.clip (output/255.0, 0, 1.0)
if is_bgr:
output = output[...,::-1]
return output
def convert_rgb_to_y(self, image):
if len(image.shape) <= 2 or image.shape[2] == 1:
return image
xform = np.array([[65.738 / 256.0, 129.057 / 256.0, 25.064 / 256.0]], dtype=np.float32)
y_image = image.dot(xform.T) + 16.0
return y_image
def convert_rgb_to_ycbcr(self, image):
if len(image.shape) <= 2 or image.shape[2] == 1:
return image
xform = np.array(
[[65.738 / 256.0, 129.057 / 256.0, 25.064 / 256.0],
[- 37.945 / 256.0, - 74.494 / 256.0, 112.439 / 256.0],
[112.439 / 256.0, - 94.154 / 256.0, - 18.285 / 256.0]], dtype=np.float32)
ycbcr_image = image.dot(xform.T)
ycbcr_image[:, :, 0] += 16.0
ycbcr_image[:, :, [1, 2]] += 128.0
return ycbcr_image
def convert_ycbcr_to_rgb(self,ycbcr_image):
rgb_image = np.zeros([ycbcr_image.shape[0], ycbcr_image.shape[1], 3], dtype=np.float32)
rgb_image[:, :, 0] = ycbcr_image[:, :, 0] - 16.0
rgb_image[:, :, [1, 2]] = ycbcr_image[:, :, [1, 2]] - 128.0
xform = np.array(
[[298.082 / 256.0, 0, 408.583 / 256.0],
[298.082 / 256.0, -100.291 / 256.0, -208.120 / 256.0],
[298.082 / 256.0, 516.412 / 256.0, 0]], dtype=np.float32)
rgb_image = rgb_image.dot(xform.T)
return rgb_image
def convert_y_and_cbcr_to_rgb(self,y_image, cbcr_image):
if len(y_image.shape) <= 2:
y_image = y_image.reshape[y_image.shape[0], y_image.shape[1], 1]
if len(y_image.shape) == 3 and y_image.shape[2] == 3:
y_image = y_image[:, :, 0:1]
ycbcr_image = np.zeros([y_image.shape[0], y_image.shape[1], 3], dtype=np.float32)
ycbcr_image[:, :, 0] = y_image[:, :, 0]
ycbcr_image[:, :, 1:3] = cbcr_image[:, :, 0:2]
return self.convert_ycbcr_to_rgb(ycbcr_image)
def flip(self, image, flip_type, invert=False):
if flip_type == 0:
return image
elif flip_type == 1:
return np.flipud(image)
elif flip_type == 2:
return np.fliplr(image)
elif flip_type == 3:
return np.flipud(np.fliplr(image))
elif flip_type == 4:
return np.rot90(image, 1 if invert is False else -1)
elif flip_type == 5:
return np.rot90(image, -1 if invert is False else 1)
elif flip_type == 6:
if invert is False:
return np.flipud(np.rot90(image))
else:
return np.rot90(np.flipud(image), -1)
elif flip_type == 7:
if invert is False:
return np.flipud(np.rot90(image, -1))
else:
return np.rot90(np.flipud(image), 1)

BIN
imagelib/RankSRGAN.h5
View file

Binary file not shown.

109
imagelib/RankSRGAN.py
View file

@ -1,109 +0,0 @@
import numpy as np
import cv2
from pathlib import Path
from nnlib import nnlib
from interact import interact as io
class RankSRGAN():
def __init__(self):
exec( nnlib.import_all(), locals(), globals() )
class PixelShufflerTorch(KL.Layer):
def __init__(self, size=(2, 2), data_format='channels_last', **kwargs):
super(PixelShufflerTorch, self).__init__(**kwargs)
self.data_format = data_format
self.size = size
def call(self, inputs):
input_shape = K.shape(inputs)
if K.int_shape(input_shape)[0] != 4:
raise ValueError('Inputs should have rank 4; Received input shape:', str(K.int_shape(inputs)))
batch_size, h, w, c = input_shape[0], input_shape[1], input_shape[2], K.int_shape(inputs)[-1]
rh, rw = self.size
oh, ow = h * rh, w * rw
oc = c // (rh * rw)
out = inputs
out = K.permute_dimensions(out, (0, 3, 1, 2)) #NCHW
out = K.reshape(out, (batch_size, oc, rh, rw, h, w))
out = K.permute_dimensions(out, (0, 1, 4, 2, 5, 3))
out = K.reshape(out, (batch_size, oc, oh, ow))
out = K.permute_dimensions(out, (0, 2, 3, 1))
return out
def compute_output_shape(self, input_shape):
if len(input_shape) != 4:
raise ValueError('Inputs should have rank ' + str(4) + '; Received input shape:', str(input_shape))
height = input_shape[1] * self.size[0] if input_shape[1] is not None else None
width = input_shape[2] * self.size[1] if input_shape[2] is not None else None
channels = input_shape[3] // self.size[0] // self.size[1]
if channels * self.size[0] * self.size[1] != input_shape[3]:
raise ValueError('channels of input and size are incompatible')
return (input_shape[0],
height,
width,
channels)
def get_config(self):
config = {'size': self.size,
'data_format': self.data_format}
base_config = super(PixelShufflerTorch, self).get_config()
return dict(list(base_config.items()) + list(config.items()))
def res_block(inp, name_prefix):
x = inp
x = Conv2D (ndf, kernel_size=3, strides=1, padding='same', activation="relu", name=name_prefix+"0")(x)
x = Conv2D (ndf, kernel_size=3, strides=1, padding='same', name=name_prefix+"2")(x)
return Add()([inp,x])
ndf = 64
nb = 16
inp = Input ( (None, None,3) )
x = inp
x = x0 = Conv2D (ndf, kernel_size=3, strides=1, padding='same', name="model0")(x)
for i in range(nb):
x = res_block(x, "model1%.2d" %i )
x = Conv2D (ndf, kernel_size=3, strides=1, padding='same', name="model1160")(x)
x = Add()([x0,x])
x = ReLU() ( PixelShufflerTorch() ( Conv2D (ndf*4, kernel_size=3, strides=1, padding='same', name="model2")(x) ) )
x = ReLU() ( PixelShufflerTorch() ( Conv2D (ndf*4, kernel_size=3, strides=1, padding='same', name="model5")(x) ) )
x = Conv2D (ndf, kernel_size=3, strides=1, padding='same', activation="relu", name="model8")(x)
x = Conv2D (3, kernel_size=3, strides=1, padding='same', name="model10")(x)
self.model = Model(inp, x )
self.model.load_weights ( Path(__file__).parent / 'RankSRGAN.h5')
def upscale(self, img, scale=2, is_bgr=True, is_float=True):
if scale not in [2,4]:
raise ValueError ("RankSRGAN: supported scale are 2 or 4.")
if not is_bgr:
img = img[...,::-1]
if not is_float:
img /= 255.0
h, w = img.shape[:2]
ch = img.shape[2] if len(img.shape) >= 3 else 1
output = self.model.predict([img[None,...]])[0]
if scale == 2:
output = cv2.resize (output, (w*scale, h*scale), cv2.INTER_CUBIC)
if not is_float:
output = np.clip (output * 255.0, 0, 255.0)
if not is_bgr:
output = output[...,::-1]
return output

17
imagelib/__init__.py
View file

@ -1,21 +1,26 @@
from .estimate_sharpness import estimate_sharpness
from .equalize_and_stack_square import equalize_and_stack_square
from .text import get_text_image, get_draw_text_lines
from .text import get_text_image
from .text import get_draw_text_lines
from .draw import draw_polygon, draw_rect
from .draw import draw_polygon
from .draw import draw_rect
from .morph import morph_by_points
from .warp import gen_warp_params, warp_by_params
from .warp import gen_warp_params
from .warp import warp_by_params
from .reduce_colors import reduce_colors
from .color_transfer import color_transfer_mkl, color_transfer_idt, color_hist_match, reinhard_color_transfer, linear_color_transfer, seamless_clone
from .color_transfer import color_hist_match
from .color_transfer import reinhard_color_transfer
from .color_transfer import linear_color_transfer
from .RankSRGAN import RankSRGAN
from .DCSCN import DCSCN
from .common import normalize_channels, overlay_alpha_image
from .common import normalize_channels
from .IEPolys import IEPolys

148
imagelib/blur.py
View file

@ -1,9 +1,143 @@
import cv2
import math
import numpy as np
from PIL import Image
from scipy.signal import convolve2d
from skimage.draw import line
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)
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])

125
imagelib/color_transfer.py
View file

@ -1,129 +1,6 @@
import cv2
import numpy as np
import cv2
import scipy as sp
import scipy.sparse
from scipy.sparse.linalg import spsolve
def color_transfer_mkl(x0, x1):
eps = np.finfo(float).eps
h,w,c = x0.shape
h1,w1,c1 = x1.shape
x0 = x0.reshape ( (h*w,c) )
x1 = x1.reshape ( (h1*w1,c1) )
a = np.cov(x0.T)
b = np.cov(x1.T)
Da2, Ua = np.linalg.eig(a)
Da = np.diag(np.sqrt(Da2.clip(eps, None)))
C = np.dot(np.dot(np.dot(np.dot(Da, Ua.T), b), Ua), Da)
Dc2, Uc = np.linalg.eig(C)
Dc = np.diag(np.sqrt(Dc2.clip(eps, None)))
Da_inv = np.diag(1./(np.diag(Da)))
t = np.dot(np.dot(np.dot(np.dot(np.dot(np.dot(Ua, Da_inv), Uc), Dc), Uc.T), Da_inv), Ua.T)
mx0 = np.mean(x0, axis=0)
mx1 = np.mean(x1, axis=0)
result = np.dot(x0-mx0, t) + mx1
return np.clip ( result.reshape ( (h,w,c) ), 0, 1)
def color_transfer_idt(i0, i1, bins=256, n_rot=20):
relaxation = 1 / n_rot
h,w,c = i0.shape
h1,w1,c1 = i1.shape
i0 = i0.reshape ( (h*w,c) )
i1 = i1.reshape ( (h1*w1,c1) )
n_dims = c
d0 = i0.T
d1 = i1.T
for i in range(n_rot):
r = sp.stats.special_ortho_group.rvs(n_dims).astype(np.float32)
d0r = np.dot(r, d0)
d1r = np.dot(r, d1)
d_r = np.empty_like(d0)
for j in range(n_dims):
lo = min(d0r[j].min(), d1r[j].min())
hi = max(d0r[j].max(), d1r[j].max())
p0r, edges = np.histogram(d0r[j], bins=bins, range=[lo, hi])
p1r, _ = np.histogram(d1r[j], bins=bins, range=[lo, hi])
cp0r = p0r.cumsum().astype(np.float32)
cp0r /= cp0r[-1]
cp1r = p1r.cumsum().astype(np.float32)
cp1r /= cp1r[-1]
f = np.interp(cp0r, cp1r, edges[1:])
d_r[j] = np.interp(d0r[j], edges[1:], f, left=0, right=bins)
d0 = relaxation * np.linalg.solve(r, (d_r - d0r)) + d0
return np.clip ( d0.T.reshape ( (h,w,c) ), 0, 1)
def laplacian_matrix(n, m):
mat_D = scipy.sparse.lil_matrix((m, m))
mat_D.setdiag(-1, -1)
mat_D.setdiag(4)
mat_D.setdiag(-1, 1)
mat_A = scipy.sparse.block_diag([mat_D] * n).tolil()
mat_A.setdiag(-1, 1*m)
mat_A.setdiag(-1, -1*m)
return mat_A
def seamless_clone(source, target, mask):
h, w,c = target.shape
result = []
mat_A = laplacian_matrix(h, w)
laplacian = mat_A.tocsc()
mask[0,:] = 1
mask[-1,:] = 1
mask[:,0] = 1
mask[:,-1] = 1
q = np.argwhere(mask==0)
k = q[:,1]+q[:,0]*w
mat_A[k, k] = 1
mat_A[k, k + 1] = 0
mat_A[k, k - 1] = 0
mat_A[k, k + w] = 0
mat_A[k, k - w] = 0
mat_A = mat_A.tocsc()
mask_flat = mask.flatten()
for channel in range(c):
source_flat = source[:, :, channel].flatten()
target_flat = target[:, :, channel].flatten()
mat_b = laplacian.dot(source_flat)*0.75
mat_b[mask_flat==0] = target_flat[mask_flat==0]
x = spsolve(mat_A, mat_b).reshape((h, w))
result.append (x)
return np.clip( np.dstack(result), 0, 1 )
def reinhard_color_transfer(target, source, clip=False, preserve_paper=False, source_mask=None, target_mask=None):
"""
Transfers the color distribution from the source to the target

31
imagelib/common.py
View file

@ -9,32 +9,13 @@ 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]
c = 1
img = img[...,np.newaxis]
if c == 1 and target_channels > 1:
img = np.repeat (img, target_channels, -1)
c = target_channels
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
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])
return img

6
imagelib/warp.py
View file

@ -4,8 +4,8 @@ from utils import random_utils
def gen_warp_params (source, flip, rotation_range=[-10,10], scale_range=[-0.5, 0.5], tx_range=[-0.05, 0.05], ty_range=[-0.05, 0.05] ):
h,w,c = source.shape
if (h != w):
raise ValueError ('gen_warp_params accepts only square images.')
if (h != w) or (w != 64 and w != 128 and w != 256 and w != 512 and w != 1024):
raise ValueError ('TrainingDataGenerator accepts only square power of 2 images.')
rotation = np.random.uniform( rotation_range[0], rotation_range[1] )
scale = np.random.uniform(1 +scale_range[0], 1 +scale_range[1])
@ -47,5 +47,5 @@ def warp_by_params (params, img, warp, transform, flip, is_border_replicate):
if transform:
img = cv2.warpAffine( img, params['rmat'], (params['w'], params['w']), borderMode=(cv2.BORDER_REPLICATE if is_border_replicate else cv2.BORDER_CONSTANT), flags=cv2.INTER_CUBIC )
if flip and params['flip']:
img = img[:,::-1,...]
img = img[:,::-1,:]
return img

26
interact/interact.py
View file

@ -33,7 +33,6 @@ class InteractBase(object):
self.key_events = {}
self.pg_bar = None
self.focus_wnd_name = None
self.error_log_line_prefix = '/!\\ '
def is_support_windows(self):
return False
@ -66,14 +65,10 @@ class InteractBase(object):
raise NotImplemented
def log_info(self, msg, end='\n'):
if self.pg_bar is not None:
print ("\n")
print (msg, end=end)
def log_err(self, msg, end='\n'):
if self.pg_bar is not None:
print ("\n")
print (f'{self.error_log_line_prefix}{msg}', end=end)
print (msg, end=end)
def named_window(self, wnd_name):
if wnd_name not in self.named_windows:
@ -137,9 +132,9 @@ class InteractBase(object):
else: print("capture_keys: already set for window ", wnd_name)
else: print("capture_keys: named_window ", wnd_name, " not found.")
def progress_bar(self, desc, total, leave=True, initial=0):
def progress_bar(self, desc, total, leave=True):
if self.pg_bar is None:
self.pg_bar = tqdm( total=total, desc=desc, leave=leave, ascii=True, initial=initial )
self.pg_bar = tqdm( total=total, desc=desc, leave=leave, ascii=True )
else: print("progress_bar: already set.")
def progress_bar_inc(self, c):
@ -154,12 +149,9 @@ class InteractBase(object):
self.pg_bar = None
else: print("progress_bar not set.")
def progress_bar_generator(self, data, desc, leave=True, initial=0):
self.pg_bar = tqdm( data, desc=desc, leave=leave, ascii=True, initial=initial )
for x in self.pg_bar:
def progress_bar_generator(self, data, desc, leave=True):
for x in tqdm( data, desc=desc, leave=leave, ascii=True ):
yield x
self.pg_bar.close()
self.pg_bar = None
def process_messages(self, sleep_time=0):
self.on_process_messages(sleep_time)
@ -252,15 +244,9 @@ class InteractBase(object):
print (help_message)
continue
if valid_list is not None:
if inp.lower() in valid_list:
return inp.lower()
if inp in valid_list:
return inp
if (valid_list is not None) and (inp.lower() not in valid_list):
return default_value
return inp
except:
print (default_value)
return default_value

9
joblib/SubprocessFunctionCaller.py
View file

@ -30,15 +30,8 @@ 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,15 +87,13 @@ class Subprocessor(object):
c2s.put ( {'op': 'error', 'data' : data} )
#overridable
def __init__(self, name, SubprocessorCli_class, no_response_time_sec = 0, io_loop_sleep_time=0.005, initialize_subprocesses_in_serial=True):
def __init__(self, name, SubprocessorCli_class, no_response_time_sec = 0):
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):
@ -135,8 +133,7 @@ class Subprocessor(object):
#overridable
def on_tick(self):
#tick in main loop
#return True if system can be finalized when no data in get_data, orelse False
return True
pass
#overridable
def on_check_run(self):
@ -160,24 +157,23 @@ class Subprocessor(object):
self.clis.append (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:
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
io.process_messages(0.005)
if cli.state == 0:
break
io.process_messages(0.005)
except:
raise Exception ("Unable to start subprocess %s" % (name))
@ -236,15 +232,6 @@ 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
@ -256,14 +243,19 @@ class Subprocessor(object):
cli.sent_data = data
cli.state = 1
if self.io_loop_sleep_time != 0:
io.process_messages(self.io_loop_sleep_time)
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.on_tick() and all ([cli.state == 0 for cli in self.clis]):
if 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[:]:

28
main.py
View file

@ -40,7 +40,7 @@ if __name__ == "__main__":
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the extracted files will be stored.")
p.add_argument('--debug-dir', action=fixPathAction, dest="debug_dir", help="Writes debug images to this directory.")
p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'head', 'full_face_no_align', 'mark_only'], default='full_face', help="Default 'full_face'. Don't change this option, currently all models uses 'full_face'")
p.add_argument('--face-type', dest="face_type", choices=['half_face', 'full_face', 'head', 'avatar', 'mark_only'], default='full_face', help="Default 'full_face'. Don't change this option, currently all models uses 'full_face'")
p.add_argument('--detector', dest="detector", choices=['dlib','mt','s3fd','manual'], default='dlib', help="Type of detector. Default 'dlib'. 'mt' (MTCNNv1) - faster, better, almost no jitter, perfect for gathering thousands faces for src-set. It is also good for dst-set, but can generate false faces in frames where main face not recognized! In this case for dst-set use either 'dlib' with '--manual-fix' or '--detector manual'. Manual detector suitable only for dst-set.")
p.add_argument('--multi-gpu', action="store_true", dest="multi_gpu", default=False, help="Enables multi GPU.")
p.add_argument('--manual-fix', action="store_true", dest="manual_fix", default=False, help="Enables manual extract only frames where faces were not recognized.")
@ -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,21 +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)
if arguments.remove_ie_polys:
Util.remove_ie_polys_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.add_argument('--remove-ie-polys', action="store_true", dest="remove_ie_polys", default=False, help="Remove ie_polys from aligned faces.")
p.set_defaults (func=process_util)
def process_train(arguments):
@ -155,8 +151,10 @@ if __name__ == "__main__":
args = {'input_dir' : arguments.input_dir,
'output_dir' : arguments.output_dir,
'aligned_dir' : arguments.aligned_dir,
'avaperator_aligned_dir' : arguments.avaperator_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,
@ -168,8 +166,10 @@ if __name__ == "__main__":
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory. A directory containing the files you wish to process.")
p.add_argument('--output-dir', required=True, action=fixPathAction, dest="output_dir", help="Output directory. This is where the converted files will be stored.")
p.add_argument('--aligned-dir', action=fixPathAction, dest="aligned_dir", help="Aligned directory. This is where the extracted of dst faces stored.")
p.add_argument('--avaperator-aligned-dir', action=fixPathAction, dest="avaperator_aligned_dir", help="Only for AVATAR model. Directory of aligned avatar operator faces.")
p.add_argument('--model-dir', required=True, action=fixPathAction, dest="model_dir", help="Model dir.")
p.add_argument('--model', 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)
@ -236,17 +236,15 @@ if __name__ == "__main__":
def process_labelingtool_edit_mask(arguments):
from mainscripts import MaskEditorTool
MaskEditorTool.mask_editor_main (arguments.input_dir, arguments.confirmed_dir, arguments.skipped_dir, no_default_mask=arguments.no_default_mask)
MaskEditorTool.mask_editor_main (arguments.input_dir, arguments.confirmed_dir, arguments.skipped_dir)
labeling_parser = subparsers.add_parser( "labelingtool", help="Labeling tool.").add_subparsers()
p = labeling_parser.add_parser ( "edit_mask", help="")
p.add_argument('--input-dir', required=True, action=fixPathAction, dest="input_dir", help="Input directory of aligned faces.")
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.add_argument('--no-default-mask', action="store_true", dest="no_default_mask", default=False, help="Don't use default mask.")
p.set_defaults(func=process_labelingtool_edit_mask)
def bad_args(arguments):
parser.print_help()
exit(0)

791
mainscripts/Converter.py
View file

@ -1,80 +1,37 @@
import math
import sys
import multiprocessing
import operator
import os
import pickle
import shutil
import sys
import time
import traceback
from pathlib import Path
import cv2
import numpy as np
import numpy.linalg as npla
import imagelib
from converters import (ConverterConfig, ConvertFaceAvatar, ConvertMasked,
FrameInfo)
from facelib import FaceType, FANSegmentator, LandmarksProcessor
from converters import Converter
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 .ConverterScreen import Screen, ScreenManager
CONVERTER_DEBUG = False
from imagelib import normalize_channels
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):
self.log_info ('Running on %s.' % (client_dict['device_name']) )
io.log_info ('Running on %s.' % (client_dict['device_name']) )
self.device_idx = client_dict['device_idx']
self.device_name = client_dict['device_name']
self.predictor_func = client_dict['predictor_func']
self.predictor_input_shape = client_dict['predictor_input_shape']
self.superres_func = client_dict['superres_func']
self.converter = client_dict['converter']
self.output_path = Path(client_dict['output_dir']) if 'output_dir' in client_dict.keys() else None
self.alignments = client_dict['alignments']
self.avatar_image_paths = client_dict['avatar_image_paths']
self.debug = client_dict['debug']
#transfer and set stdin in order to work code.interact in debug subprocess
stdin_fd = client_dict['stdin_fd']
@ -87,498 +44,191 @@ class ConvertSubprocessor(Subprocessor):
#therefore forcing active_DeviceConfig to CPU only
nnlib.active_DeviceConfig = nnlib.DeviceConfig (cpu_only=True)
def sharpen_func (img, sharpen_mode=0, kernel_size=3, amount=150):
if kernel_size % 2 == 0:
kernel_size += 1
self.converter.on_cli_initialize()
if sharpen_mode == 1: #box
kernel = np.zeros( (kernel_size, kernel_size), dtype=np.float32)
kernel[ kernel_size//2, kernel_size//2] = 1.0
box_filter = np.ones( (kernel_size, kernel_size), dtype=np.float32) / (kernel_size**2)
kernel = kernel + (kernel - box_filter) * amount
return cv2.filter2D(img, -1, kernel)
elif sharpen_mode == 2: #gaussian
blur = cv2.GaussianBlur(img, (kernel_size, kernel_size) , 0)
img = cv2.addWeighted(img, 1.0 + (0.5 * amount), blur, -(0.5 * amount), 0)
return img
return img
self.sharpen_func = sharpen_func
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
import ebsynth
def ebs_ct(*args, **kwargs):
return ebsynth.color_transfer(*args, **kwargs)
self.ebs_ct_func = ebs_ct
return None
#override
def process_data(self, pf): #pf=ProcessingFrame
cfg = pf.cfg.copy()
cfg.sharpen_func = self.sharpen_func
cfg.superres_func = self.superres_func
cfg.ebs_ct_func = self.ebs_ct_func
frame_info = pf.frame_info
filename = frame_info.filename
landmarks_list = frame_info.landmarks_list
def process_data(self, data):
idx, filename = data
filename_path = Path(filename)
output_filename = pf.output_filename
need_return_image = pf.need_return_image
files_processed = 1
faces_processed = 0
if len(landmarks_list) == 0:
self.log_info ( 'no faces found for %s, copying without faces' % (filename_path.name) )
output_filename_path = self.output_path / (filename_path.stem + '.png')
if filename_path.suffix == '.png':
shutil.copy (filename, output_filename )
else:
img_bgr = cv2_imread(filename)
cv2_imwrite (output_filename, img_bgr)
if (self.converter.type == Converter.TYPE_FACE or self.converter.type == Converter.TYPE_FACE_AVATAR ) \
and 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 need_return_image:
img_bgr = cv2_imread(filename)
pf.image = img_bgr
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 )
else:
if cfg.type == ConverterConfig.TYPE_MASKED:
cfg.fanseg_input_size = self.fanseg_input_size
cfg.fanseg_extract_func = self.fanseg_extract_func
image = (cv2_imread(str(filename_path)) / 255.0).astype(np.float32)
image = normalize_channels (image, 3)
if self.converter.type == Converter.TYPE_IMAGE:
image = self.converter.cli_convert_image(image, None, self.debug)
try:
final_img = ConvertMasked (self.predictor_func, self.predictor_input_shape, 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:
return (1, image)
faces_processed = 1
elif self.converter.type == Converter.TYPE_IMAGE_WITH_LANDMARKS:
#currently unused
if filename_path.suffix == '.png':
dflimg = DFLPNG.load( str(filename_path) )
elif filename_path.suffix == '.jpg':
dflimg = DFLJPG.load ( str(filename_path) )
else:
dflimg = None
elif cfg.type == ConverterConfig.TYPE_FACE_AVATAR:
final_img = ConvertFaceAvatar (self.predictor_func, self.predictor_input_shape,
cfg, pf.prev_temporal_frame_infos,
pf.frame_info,
pf.next_temporal_frame_infos )
if dflimg is not None:
image_landmarks = dflimg.get_landmarks()
if output_filename is not None and final_img is not None:
cv2_imwrite (output_filename, final_img )
image = self.converter.convert_image(image, image_landmarks, self.debug)
if need_return_image:
pf.image = final_img
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) )
return pf
elif self.converter.type == Converter.TYPE_FACE or self.converter.type == Converter.TYPE_FACE_AVATAR:
ava_face = None
if self.converter.type == Converter.TYPE_FACE_AVATAR:
ava_filename_path = self.avatar_image_paths[idx]
ava_face = (cv2_imread(str(ava_filename_path)) / 255.0).astype(np.float32)
ava_face = normalize_channels (ava_face, 3)
faces = self.alignments[filename_path.stem]
if self.debug:
debug_images = []
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) )
if self.debug:
debug_images += self.converter.cli_convert_face(image, image_landmarks, self.debug, avaperator_face_bgr=ava_face)
else:
image = self.converter.cli_convert_face(image, image_landmarks, self.debug, avaperator_face_bgr=ava_face)
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 self.debug:
return (1, debug_images)
faces_processed = len(faces)
if not self.debug:
cv2_imwrite (str(output_filename_path), (image*255).astype(np.uint8) )
return (0, files_processed, faces_processed)
#overridable
def get_data_name (self, pf):
def get_data_name (self, data):
#return string identificator of your data
return pf.frame_info.filename
idx, filename = data
return filename
#override
def __init__(self, is_interactive, converter_session_filepath, predictor_func, predictor_input_shape, converter_config, frames, output_path, model_iter):
if len (frames) == 0:
raise ValueError ("len (frames) == 0")
super().__init__('Converter', ConvertSubprocessor.Cli, 86400 if CONVERTER_DEBUG else 60, io_loop_sleep_time=0.001, initialize_subprocesses_in_serial=False)
self.is_interactive = is_interactive
self.converter_session_filepath = Path(converter_session_filepath)
self.converter_config = converter_config
#dummy predict and sleep, tensorflow caching kernels. If remove it, sometime conversion speed can be x2 slower
predictor_func (dummy_predict=True)
time.sleep(2)
self.predictor_func_host, self.predictor_func = SubprocessFunctionCaller.make_pair(predictor_func)
self.predictor_input_shape = predictor_input_shape
self.dcscn = None
self.ranksrgan = None
def superres_func(mode, *args, **kwargs):
if mode == 1:
if self.ranksrgan is None:
self.ranksrgan = imagelib.RankSRGAN()
return self.ranksrgan.upscale(*args, **kwargs)
self.dcscn_host, self.superres_func = SubprocessFunctionCaller.make_pair(superres_func)
def __init__(self, converter, input_path_image_paths, output_path, alignments, avatar_image_paths=None, debug = False):
super().__init__('Converter', ConvertSubprocessor.Cli, 86400 if debug == True else 60)
self.converter = converter
self.input_data = self.input_path_image_paths = input_path_image_paths
self.input_data_idxs = [ *range(len(self.input_data)) ]
self.output_path = output_path
self.model_iter = model_iter
self.prefetch_frame_count = self.process_count = min(6,multiprocessing.cpu_count())
session_data = None
if self.is_interactive and self.converter_session_filepath.exists():
try:
with open( str(self.converter_session_filepath), "rb") as f:
session_data = pickle.loads(f.read())
except Exception as e:
pass
self.frames = frames
self.frames_idxs = [ *range(len(self.frames)) ]
self.frames_done_idxs = []
if self.is_interactive and session_data is not None:
s_frames = session_data.get('frames', None)
s_frames_idxs = session_data.get('frames_idxs', None)
s_frames_done_idxs = session_data.get('frames_done_idxs', None)
s_model_iter = session_data.get('model_iter', None)
frames_equal = (s_frames is not None) and \
(s_frames_idxs is not None) and \
(s_frames_done_idxs is not None) and \
(s_model_iter is not None) and \
(len(frames) == len(s_frames))
if frames_equal:
for i in range(len(frames)):
frame = frames[i]
s_frame = s_frames[i]
if frame.frame_info.filename != s_frame.frame_info.filename:
frames_equal = False
if not frames_equal:
break
if frames_equal:
io.log_info ('Using saved session from ' + '/'.join (self.converter_session_filepath.parts[-2:]) )
self.frames = s_frames
self.frames_idxs = s_frames_idxs
self.frames_done_idxs = s_frames_done_idxs
if self.model_iter != s_model_iter:
#model is more trained, recompute all frames
for frame in self.frames:
frame.is_done = False
if self.model_iter != s_model_iter or \
len(self.frames_idxs) == 0:
#rewind to begin if model is more trained or all frames are done
while len(self.frames_done_idxs) > 0:
prev_frame = self.frames[self.frames_done_idxs.pop()]
self.frames_idxs.insert(0, prev_frame.idx)
if len(self.frames_idxs) != 0:
cur_frame = self.frames[self.frames_idxs[0]]
cur_frame.is_shown = False
if not frames_equal:
session_data = None
if session_data is None:
for filename in Path_utils.get_image_paths(self.output_path): #remove all images in output_path
Path(filename).unlink()
frames[0].cfg = self.converter_config.copy()
for i in range( len(self.frames) ):
frame = self.frames[i]
frame.idx = i
frame.output_filename = self.output_path / ( Path(frame.frame_info.filename).stem + '.png' )
self.alignments = alignments
self.avatar_image_paths = avatar_image_paths
self.debug = debug
self.files_processed = 0
self.faces_processed = 0
#override
def process_info_generator(self):
r = [0] if CONVERTER_DEBUG else range(self.process_count)
r = [0] if self.debug else range( min(6,multiprocessing.cpu_count()) )
for i in r:
yield 'CPU%d' % (i), {}, {'device_idx': i,
'device_name': 'CPU%d' % (i),
'predictor_func': self.predictor_func,
'predictor_input_shape' : self.predictor_input_shape,
'superres_func': self.superres_func,
'stdin_fd': sys.stdin.fileno() if CONVERTER_DEBUG else None
'converter' : self.converter,
'output_dir' : str(self.output_path),
'alignments' : self.alignments,
'avatar_image_paths' : self.avatar_image_paths,
'debug': self.debug,
'stdin_fd': sys.stdin.fileno() if self.debug else None
}
#overridable optional
def on_clients_initialized(self):
io.progress_bar ("Converting", len (self.frames_idxs), initial=len(self.frames_done_idxs) )
if self.debug:
io.named_window ("Debug convert")
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()
io.progress_bar ("Converting", len (self.input_data_idxs) )
#overridable optional
def on_clients_finalized(self):
io.progress_bar_close()
if self.is_interactive:
self.screen_manager.finalize()
for frame in self.frames:
frame.output_filename = None
frame.image = None
session_data = {
'frames': self.frames,
'frames_idxs': self.frames_idxs,
'frames_done_idxs': self.frames_done_idxs,
'model_iter' : self.model_iter,
}
self.converter_session_filepath.write_bytes( pickle.dumps(session_data) )
io.log_info ("Session is saved to " + '/'.join (self.converter_session_filepath.parts[-2:]) )
cfg_change_keys = ['`','1', '2', '3', '4', '5', '6', '7', '8', '9',
'q', 'a', 'w', 's', 'e', 'd', 'r', 'f', 't', 'g','y','h','u','j',
'z', 'x', 'c', 'v', 'b','n' ]
#override
def on_tick(self):
self.predictor_func_host.process_messages()
self.dcscn_host.process_messages()
go_prev_frame = False
go_prev_frame_overriding_cfg = 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.to_string( cur_frame.frame_info.filename_short) )
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_sharpen_amount(1 if not shift_pressed else 5)
elif chr_key == 'h':
cfg.add_sharpen_amount(-1 if not shift_pressed else -5)
elif chr_key == 'u':
cfg.add_output_face_scale(1 if not shift_pressed else 5)
elif chr_key == 'j':
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_mode()
elif chr_key == 'b':
cfg.toggle_export_mask_alpha()
elif chr_key == 'n':
cfg.toggle_sharpen_mode()
else:
if chr_key == 'y':
cfg.add_sharpen_amount(1 if not shift_pressed else 5)
elif chr_key == 'h':
cfg.add_sharpen_amount(-1 if not shift_pressed else -5)
elif chr_key == 's':
cfg.toggle_add_source_image()
elif chr_key == 'v':
cfg.toggle_super_resolution_mode()
elif chr_key == 'n':
cfg.toggle_sharpen_mode()
if prev_cfg != cfg:
io.log_info ( cfg.to_string(cur_frame.frame_info.filename_short) )
cur_frame.is_done = False
cur_frame.is_shown = False
else:
if chr_key == ',' or chr_key == 'm':
self.process_remain_frames = False
go_prev_frame = True
go_prev_frame_overriding_cfg = chr_key == 'm'
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 None or cur_frame.is_done:
if cur_frame is not None:
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)
if cur_frame is not None and go_prev_frame_overriding_cfg:
if prev_frame.cfg != cur_frame.cfg:
prev_frame.cfg = cur_frame.cfg.copy()
prev_frame.is_done = False
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
if self.debug:
io.destroy_all_windows()
#override
def get_data(self, host_dict):
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 )
if len (self.input_data_idxs) > 0:
idx = self.input_data_idxs.pop(0)
return (idx, self.input_data[idx])
return None
#override
def on_data_return (self, host_dict, data):
idx, filename = 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 0
return self.files_processed, self.faces_processed
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'])
@ -588,26 +238,24 @@ def main (args, device_args):
io.log_err('Input directory not found. Please ensure it exists.')
return
if not output_path.exists():
if output_path.exists():
for filename in Path_utils.get_image_paths(output_path):
Path(filename).unlink()
else:
output_path.mkdir(parents=True, exist_ok=True)
if not model_path.exists():
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_session_filepath = model.get_strpath_storage_for_file('converter_session.dat')
predictor_func, predictor_input_shape, cfg = model.get_ConverterConfig()
if not is_interactive:
cfg.ask_settings()
converter = model.get_converter()
input_path_image_paths = Path_utils.get_image_paths(input_path)
if cfg.type == ConverterConfig.TYPE_MASKED:
alignments = None
avatar_image_paths = None
if converter.type == Converter.TYPE_FACE or converter.type == Converter.TYPE_FACE_AVATAR:
if aligned_dir is None:
io.log_err('Aligned directory not found. Please ensure it exists.')
return
@ -618,7 +266,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)
@ -638,55 +286,22 @@ def main (args, device_args):
if source_filename_stem not in alignments.keys():
alignments[ source_filename_stem ] = []
alignments_ar = alignments[ source_filename_stem ]
alignments_ar.append (dflimg.get_source_landmarks())
if len(alignments_ar) > 1:
multiple_faces_detected = True
alignments[ source_filename_stem ].append (dflimg.get_source_landmarks())
if converter.type == Converter.TYPE_FACE_AVATAR:
if avaperator_aligned_dir is None:
io.log_err('Avatar operator aligned directory not found. Please ensure it exists.')
return
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)
avaperator_aligned_path = Path(avaperator_aligned_dir)
if not avaperator_aligned_path.exists():
io.log_err('Avatar operator aligned directory not found. Please ensure it exists.')
return
avatar_image_paths = []
for filename in io.progress_bar_generator( Path_utils.get_image_paths(avaperator_aligned_path) , "Sorting avaperator faces"):
filepath = Path(filename)
if filepath.suffix == '.png':
dflimg = DFLPNG.load( str(filepath) )
elif filepath.suffix == '.jpg':
@ -695,42 +310,23 @@ def main (args, device_args):
dflimg = None
if dflimg is None:
io.log_err ("%s is not a dfl image file" % (filepath.name) )
continue
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:
ConvertSubprocessor (
is_interactive = is_interactive,
converter_session_filepath = converter_session_filepath,
predictor_func = predictor_func,
predictor_input_shape = predictor_input_shape,
converter_config = cfg,
frames = frames,
output_path = output_path,
model_iter = model.get_iter()
io.log_err ("Fatal error: %s is not a dfl image file" % (filepath.name) )
return
avatar_image_paths += [ (filename, dflimg.get_source_filename() ) ]
avatar_image_paths = [ p[0] for p in sorted(avatar_image_paths, key=operator.itemgetter(1)) ]
if len(input_path_image_paths) < len(avatar_image_paths):
io.log_err("Input faces count must be >= avatar operator faces count.")
return
files_processed, faces_processed = ConvertSubprocessor (
converter = converter,
input_path_image_paths = input_path_image_paths,
output_path = output_path,
alignments = alignments,
avatar_image_paths = avatar_image_paths,
debug = args.get('debug',False)
).run()
model.finalize()
@ -739,6 +335,29 @@ 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

138
mainscripts/ConverterScreen/ConverterScreen.py
View file

@ -1,138 +0,0 @@
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
View file

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

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

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 3.2 KiB

46
mainscripts/Extractor.py
View file

@ -6,7 +6,6 @@ import multiprocessing
import shutil
from pathlib import Path
import numpy as np
import math
import mathlib
import imagelib
import cv2
@ -22,8 +21,6 @@ from nnlib import nnlib
from joblib import Subprocessor
from interact import interact as io
DEBUG = False
class ExtractSubprocessor(Subprocessor):
class Data(object):
def __init__(self, filename=None, rects=None, landmarks = None, landmarks_accurate=True, pitch_yaw_roll=None, final_output_files = None):
@ -47,11 +44,6 @@ class ExtractSubprocessor(Subprocessor):
self.cpu_only = client_dict['device_type'] == 'CPU'
self.final_output_path = Path(client_dict['final_output_dir']) if 'final_output_dir' in client_dict.keys() else None
self.debug_dir = client_dict['debug_dir']
#transfer and set stdin in order to work code.interact in debug subprocess
stdin_fd = client_dict['stdin_fd']
if stdin_fd is not None and DEBUG:
sys.stdin = os.fdopen(stdin_fd)
self.cached_image = (None, None)
@ -118,8 +110,18 @@ class ExtractSubprocessor(Subprocessor):
self.log_err ( 'Failed to extract %s, reason: cv2_imread() fail.' % ( str(filename_path) ) )
return data
image = imagelib.normalize_channels(image, 3)
h, w, ch = image.shape
image_shape = image.shape
if len(image_shape) == 2:
h, w = image.shape
image = image[:,:,np.newaxis]
ch = 1
else:
h, w, ch = image.shape
if ch == 1:
image = np.repeat (image, 3, -1)
elif ch == 4:
image = image[:,:,0:3]
wm, hm = w % 2, h % 2
if wm + hm != 0: #fix odd image
@ -212,22 +214,16 @@ class ExtractSubprocessor(Subprocessor):
else:
face_idx = 0
for rect, image_landmarks in zip( rects, landmarks ):
if src_dflimg is not None and face_idx > 1:
#cannot extract more than 1 face from dflimg
break
if image_landmarks is None:
continue
rect = np.array(rect)
if self.face_type == FaceType.MARK_ONLY:
image_to_face_mat = None
face_image = image
face_image_landmarks = image_landmarks
else:
image_to_face_mat = LandmarksProcessor.get_transform_mat (image_landmarks, self.image_size, self.face_type)
face_image = cv2.warpAffine(image, image_to_face_mat, (self.image_size, self.image_size), cv2.INTER_LANCZOS4)
face_image_landmarks = LandmarksProcessor.transform_points (image_landmarks, image_to_face_mat)
@ -239,11 +235,11 @@ class ExtractSubprocessor(Subprocessor):
if landmarks_area > 4*rect_area: #get rid of faces which umeyama-landmark-area > 4*detector-rect-area
continue
if self.debug_dir is not None:
LandmarksProcessor.draw_rect_landmarks (debug_image, rect, image_landmarks, self.image_size, self.face_type, transparent_mask=True)
if self.debug_dir is not None:
LandmarksProcessor.draw_rect_landmarks (debug_image, rect, image_landmarks, self.image_size, self.face_type, transparent_mask=True)
if src_dflimg is not None and filename_path.suffix == '.jpg':
#if extracting from dflimg and jpg copy it in order not to lose quality
if src_dflimg is not None:
#if extracting from dflimg copy it in order not to lose quality
output_file = str(self.final_output_path / filename_path.name)
if str(filename_path) != str(output_file):
shutil.copy ( str(filename_path), str(output_file) )
@ -296,7 +292,7 @@ class ExtractSubprocessor(Subprocessor):
self.devices = ExtractSubprocessor.get_devices_for_config(self.manual, self.type, multi_gpu, cpu_only)
no_response_time_sec = 60 if not self.manual and not DEBUG else 999999
no_response_time_sec = 60 if not self.manual else 999999
super().__init__('Extractor', ExtractSubprocessor.Cli, no_response_time_sec)
#override
@ -342,8 +338,7 @@ class ExtractSubprocessor(Subprocessor):
'image_size': self.image_size,
'face_type': self.face_type,
'debug_dir': self.debug_dir,
'final_output_dir': str(self.final_output_path),
'stdin_fd': sys.stdin.fileno() }
'final_output_dir': str(self.final_output_path)}
for (device_idx, device_type, device_name, device_total_vram_gb) in self.devices:
@ -383,8 +378,7 @@ class ExtractSubprocessor(Subprocessor):
if self.cache_original_image[0] == filename:
self.original_image = self.cache_original_image[1]
else:
self.original_image = imagelib.normalize_channels( cv2_imread( filename ), 3 )
self.original_image = cv2_imread( filename )
self.cache_original_image = (filename, self.original_image )
(h,w,c) = self.original_image.shape
@ -622,7 +616,7 @@ class ExtractSubprocessor(Subprocessor):
return [ (i, 'CPU', 'CPU%d' % (i), 0 ) for i in range( min(8, multiprocessing.cpu_count() // 2) ) ]
elif type == 'final':
return [ (i, 'CPU', 'CPU%d' % (i), 0 ) for i in (range(min(8, multiprocessing.cpu_count())) if not DEBUG else [0]) ]
return [ (i, 'CPU', 'CPU%d' % (i), 0 ) for i in range(min(8, multiprocessing.cpu_count())) ]
class DeletedFilesSearcherSubprocessor(Subprocessor):
class Cli(Subprocessor.Cli):

16
mainscripts/MaskEditorTool.py
View file

@ -320,7 +320,7 @@ class MaskEditor:
def get_ie_polys(self):
return self.ie_polys
def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None, no_default_mask=False):
def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None):
input_path = Path(input_dir)
confirmed_path = Path(confirmed_dir)
@ -334,11 +334,6 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None, no_default
if not skipped_path.exists():
skipped_path.mkdir(parents=True)
if not no_default_mask:
eyebrows_expand_mod = np.clip ( io.input_int ("Default eyebrows expand modifier? (0..400, skip:100) : ", 100), 0, 400 ) / 100.0
else:
eyebrows_expand_mod = None
wnd_name = "MaskEditor tool"
io.named_window (wnd_name)
@ -412,10 +407,7 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None, no_default
if fanseg_mask is not None:
mask = fanseg_mask
else:
if no_default_mask:
mask = np.zeros ( (target_wh,target_wh,3) )
else:
mask = LandmarksProcessor.get_image_hull_mask( img.shape, lmrks, eyebrows_expand_mod=eyebrows_expand_mod)
mask = LandmarksProcessor.get_image_hull_mask( img.shape, lmrks)
else:
img = np.zeros ( (target_wh,target_wh,3) )
mask = np.ones ( (target_wh,target_wh,3) )
@ -514,7 +506,7 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None, no_default
do_save_move_count -= 1
ed.mask_finish()
dflimg.embed_and_set (str(filepath), ie_polys=ed.get_ie_polys(), eyebrows_expand_mod=eyebrows_expand_mod )
dflimg.embed_and_set (str(filepath), ie_polys=ed.get_ie_polys() )
done_paths += [ confirmed_path / filepath.name ]
done_images_types[filepath.name] = 2
@ -525,7 +517,7 @@ def mask_editor_main(input_dir, confirmed_dir=None, skipped_dir=None, no_default
do_save_count -= 1
ed.mask_finish()
dflimg.embed_and_set (str(filepath), ie_polys=ed.get_ie_polys(), eyebrows_expand_mod=eyebrows_expand_mod )
dflimg.embed_and_set (str(filepath), ie_polys=ed.get_ie_polys() )
done_paths += [ filepath ]
done_images_types[filepath.name] = 2

11
mainscripts/Sorter.py
View file

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

12
mainscripts/Trainer.py
View file

@ -12,7 +12,7 @@ import cv2
import models
from interact import interact as io
def trainerThread (s2c, c2s, e, args, device_args):
def trainerThread (s2c, c2s, args, device_args):
while True:
try:
start_time = time.time()
@ -66,7 +66,6 @@ def trainerThread (s2c, c2s, e, args, device_args):
else:
previews = [( 'debug, press update for new', model.debug_one_iter())]
c2s.put ( {'op':'show', 'previews': previews} )
e.set() #Set the GUI Thread as Ready
if model.is_first_run():
@ -190,12 +189,9 @@ def main(args, device_args):
s2c = queue.Queue()
c2s = queue.Queue()
e = threading.Event()
thread = threading.Thread(target=trainerThread, args=(s2c, c2s, e, args, device_args) )
thread = threading.Thread(target=trainerThread, args=(s2c, c2s, args, device_args) )
thread.start()
e.wait() #Wait for inital load to occur.
if no_preview:
while True:
if not c2s.empty():
@ -295,7 +291,7 @@ def main(args, device_args):
key_events = io.get_key_events(wnd_name)
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 == ord('\n') or key == ord('\r'):
s2c.put ( {'op': 'close'} )
elif key == ord('s'):
@ -325,4 +321,4 @@ def main(args, device_args):
except KeyboardInterrupt:
s2c.put ( {'op': 'close'} )
io.destroy_all_windows()
io.destroy_all_windows()

27
mainscripts/Util.py
View file

@ -7,33 +7,6 @@ from utils.cv2_utils import *
from facelib import LandmarksProcessor
from interact import interact as io
def remove_ie_polys_file (filepath):
filepath = Path(filepath)
if filepath.suffix == '.png':
dflimg = DFLPNG.load( str(filepath) )
elif filepath.suffix == '.jpg':
dflimg = DFLJPG.load ( str(filepath) )
else:
return
if dflimg is None:
io.log_err ("%s is not a dfl image file" % (filepath.name) )
return
dflimg.remove_ie_polys()
dflimg.embed_and_set( str(filepath) )
def remove_ie_polys_folder(input_path):
input_path = Path(input_path)
io.log_info ("Removing ie_polys...\r\n")
for filepath in io.progress_bar_generator( Path_utils.get_image_paths(input_path), "Removing"):
filepath = Path(filepath)
remove_ie_polys_file(filepath)
def remove_fanseg_file (filepath):
filepath = Path(filepath)

17
mainscripts/VideoEd.py
View file

@ -167,10 +167,8 @@ def video_from_sequence( input_dir, output_file, reference_file=None, ext=None,
if not lossless and bitrate is None:
bitrate = max (1, io.input_int ("Bitrate of output file in MB/s ? (default:16) : ", 16) )
input_image_paths = Path_utils.get_image_paths(input_path)
i_in = ffmpeg.input(str (input_path / ('%5d.'+ext)), r=fps)
i_in = ffmpeg.input('pipe:', format='image2pipe', r=fps)
output_args = [i_in]
if ref_in_a is not None:
@ -195,16 +193,7 @@ def video_from_sequence( input_dir, output_file, reference_file=None, ext=None,
})
job = ( ffmpeg.output(*output_args, **output_kwargs).overwrite_output() )
try:
job_run = job.run_async(pipe_stdin=True)
for image_path in input_image_paths:
with open (image_path, "rb") as f:
image_bytes = f.read()
job_run.stdin.write (image_bytes)
job_run.stdin.close()
job_run.wait()
try:
job = job.run()
except:
io.log_err ("ffmpeg fail, job commandline:" + str(job.compile()) )

BIN
mainscripts/gfx/help_converter_face_avatar.jpg
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 159 KiB

BIN
mainscripts/gfx/help_converter_face_avatar_source.psd
View file

Binary file not shown.

BIN
mainscripts/gfx/help_converter_masked.jpg
View file

Binary file not shown.
Side by side

Before

Width:  |  Height:  |  Size: 288 KiB

BIN
mainscripts/gfx/help_converter_masked_source.psd
View file

Binary file not shown.

328
models/ModelBase.py
View file

@ -1,22 +1,19 @@
import colorsys
import inspect
import json
import os
import pickle
import shutil
import json
import time
from pathlib import Path
import cv2
import numpy as np
import inspect
import pickle
import colorsys
import imagelib
from interact import interact as io
from nnlib import nnlib
from samplelib import SampleGeneratorBase
from utils import Path_utils, std_utils
from pathlib import Path
from utils import Path_utils
from utils import std_utils
from utils.cv2_utils import *
import numpy as np
import cv2
from samplelib import SampleGeneratorBase
from nnlib import nnlib
from interact import interact as io
'''
You can implement your own model. Check examples.
'''
@ -24,12 +21,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)
@ -45,7 +41,7 @@ class ModelBase(object):
device_args['force_gpu_idx'] = io.input_int("Which GPU idx to choose? ( skip: best GPU ) : ", -1, [ x[0] for x in idxs_names_list] )
self.device_args = device_args
self.device_config = nnlib.DeviceConfig(allow_growth=True, **self.device_args)
self.device_config = nnlib.DeviceConfig(allow_growth=False, **self.device_args)
io.log_info ("Loading model...")
@ -55,7 +51,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
@ -78,8 +74,8 @@ class ModelBase(object):
self.options.pop('epoch')
if self.iter != 0:
self.options = model_data['options']
self.loss_history = model_data.get('loss_history', [])
self.sample_for_preview = model_data.get('sample_for_preview', None)
self.loss_history = model_data['loss_history'] if 'loss_history' in model_data.keys() else []
self.sample_for_preview = model_data['sample_for_preview'] if 'sample_for_preview' in model_data.keys() else None
ask_override = self.is_training_mode and self.iter != 0 and io.input_in_time ("Press enter in 2 seconds to override model settings.", 5 if io.is_colab() else 2 )
@ -88,12 +84,6 @@ class ModelBase(object):
if self.iter == 0:
io.log_info ("\nModel first run. Enter model options as default for each run.")
if ask_enable_autobackup and (self.iter == 0 or ask_override):
default_autobackup = False if self.iter == 0 else self.options.get('autobackup',False)
self.options['autobackup'] = io.input_bool("Enable autobackup? (y/n ?:help skip:%s) : " % (yn_str[default_autobackup]) , default_autobackup, help_message="Autobackup model files with preview every hour for last 15 hours. Latest backup located in model/<>_autobackups/01")
else:
self.options['autobackup'] = self.options.get('autobackup', False)
if ask_write_preview_history and (self.iter == 0 or ask_override):
default_write_preview_history = False if self.iter == 0 else self.options.get('write_preview_history',False)
self.options['write_preview_history'] = io.input_bool("Write preview history? (y/n ?:help skip:%s) : " % (yn_str[default_write_preview_history]) , default_write_preview_history, help_message="Preview history will be writed to <ModelName>_history folder.")
@ -102,11 +92,9 @@ class ModelBase(object):
if (self.iter == 0 or ask_override) and self.options['write_preview_history'] and io.is_support_windows():
choose_preview_history = io.input_bool("Choose image for the preview history? (y/n skip:%s) : " % (yn_str[False]) , False)
elif (self.iter == 0 or ask_override) and self.options['write_preview_history'] and io.is_colab():
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))
@ -117,14 +105,13 @@ class ModelBase(object):
if ask_batch_size and (self.iter == 0 or ask_override):
default_batch_size = 0 if self.iter == 0 else self.options.get('batch_size',0)
self.batch_size = max(0, io.input_int("Batch_size (?:help skip:%d) : " % (default_batch_size), default_batch_size, help_message="Larger batch size is better for NN's generalization, but it can cause Out of Memory error. Tune this value for your videocard manually."))
self.options['batch_size'] = max(0, io.input_int("Batch_size (?:help skip:%d) : " % (default_batch_size), default_batch_size, help_message="Larger batch size is better for NN's generalization, but it can cause Out of Memory error. Tune this value for your videocard manually."))
else:
self.batch_size = self.options.get('batch_size', 0)
self.options['batch_size'] = self.options.get('batch_size', 0)
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." )
if (self.iter == 0):
self.options['sort_by_yaw'] = io.input_bool("Feed faces to network sorted by yaw? (y/n ?:help skip:n) : ", False, 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." )
else:
self.options['sort_by_yaw'] = self.options.get('sort_by_yaw', False)
@ -140,10 +127,6 @@ class ModelBase(object):
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')
self.write_preview_history = self.options.get('write_preview_history', False)
if not self.write_preview_history and 'write_preview_history' in self.options:
self.options.pop('write_preview_history')
@ -152,7 +135,7 @@ class ModelBase(object):
if self.target_iter == 0 and 'target_iter' in self.options:
self.options.pop('target_iter')
#self.batch_size = self.options.get('batch_size',0)
self.batch_size = self.options.get('batch_size',0)
self.sort_by_yaw = self.options.get('sort_by_yaw',False)
self.random_flip = self.options.get('random_flip',True)
@ -176,16 +159,8 @@ class ModelBase(object):
if self.is_training_mode:
if self.device_args['force_gpu_idx'] == -1:
self.preview_history_path = self.model_path / ( '%s_history' % (self.get_model_name()) )
self.autobackups_path = self.model_path / ( '%s_autobackups' % (self.get_model_name()) )
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)
if self.write_preview_history or io.is_colab():
if not self.preview_history_path.exists():
@ -202,84 +177,67 @@ 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 choose_preview_history and io.is_support_windows():
io.log_info ("Choose image for the preview history. [p] - next. [enter] - confirm.")
wnd_name = "[p] - next. [enter] - confirm."
io.named_window(wnd_name)
io.capture_keys(wnd_name)
choosed = False
while not choosed:
if (self.sample_for_preview is None) or (self.iter == 0):
if self.iter == 0:
if choose_preview_history and io.is_support_windows():
wnd_name = "[p] - next. [enter] - confirm."
io.named_window(wnd_name)
io.capture_keys(wnd_name)
choosed = False
while not choosed:
self.sample_for_preview = self.generate_next_sample()
preview = self.get_static_preview()
io.show_image( wnd_name, (preview*255).astype(np.uint8) )
while True:
key_events = io.get_key_events(wnd_name)
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 == ord('\n') or key == ord('\r'):
choosed = True
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()
preview = self.get_static_preview()
io.show_image( wnd_name, (preview*255).astype(np.uint8) )
while True:
key_events = io.get_key_events(wnd_name)
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 == ord('\n') or key == ord('\r'):
choosed = True
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()
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_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}==']
model_summary_text += [f'=={"Model name": >{width_name}}: {self.get_model_name(): <{width_value}}=='] # Name
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}==']
model_summary_text += ["===== Model summary ====="]
model_summary_text += ["== Model name: " + self.get_model_name()]
model_summary_text += ["=="]
model_summary_text += ["== Current iteration: " + str(self.iter)]
model_summary_text += ["=="]
model_summary_text += ["== Model options:"]
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 += ["== |== %s : %s" % (key, self.options[key])]
model_summary_text += [f'=={" Running On ":-^{width_total}}=='] # Training hardware info
model_summary_text += [f'=={" "*width_total}==']
if self.device_config.multi_gpu:
model_summary_text += [f'=={"Using multi_gpu": >{width_name}}: {"True": <{width_value}}=='] # multi_gpu
model_summary_text += [f'=={" "*width_total}==']
model_summary_text += ["== |== multi_gpu : True "]
model_summary_text += ["== Running on:"]
if self.device_config.cpu_only:
model_summary_text += [f'=={"Using device": >{width_name}}: {"CPU": <{width_value}}=='] # cpu_only
model_summary_text += ["== |== [CPU]"]
else:
for idx in self.device_config.gpu_idxs:
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'=={" "*width_total}==']
model_summary_text += [f'=={"="*width_total}==']
model_summary_text += ["== |== [%d : %s]" % (idx, nnlib.device.getDeviceName(idx))]
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:"]
model_summary_text += ["/!\\ You are using a GPU with 2GB or less VRAM. This may significantly reduce the quality of your result!"]
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 += ["/!\\"]
if not self.device_config.cpu_only and self.device_config.gpu_vram_gb[0] == 2:
model_summary_text += ["=="]
model_summary_text += ["== WARNING: You are using 2GB GPU. Result quality may be significantly decreased."]
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 get extra free VRAM."]
model_summary_text += ["=="]
model_summary_text = "\n".join (model_summary_text)
model_summary_text += ["========================="]
model_summary_text = "\r\n".join (model_summary_text)
self.model_summary_text = model_summary_text
io.log_info(model_summary_text)
@ -320,14 +278,10 @@ class ModelBase(object):
def get_model_name(self):
return Path(inspect.getmodule(self).__file__).parent.name.rsplit("_", 1)[1]
#overridable , return [ [model, filename],... ] list
def get_model_filename_list(self):
return []
#overridable
def get_ConverterConfig(self):
#return predictor_func, predictor_input_shape, ConverterConfig() for the model
raise NotImplementedError
def get_converter(self):
raise NotImplementeError
#return existing or your own converter which derived from base
def get_target_iter(self):
return self.target_iter
@ -361,8 +315,7 @@ class ModelBase(object):
return self.onGetPreview (self.sample_for_preview)[0][1] #first preview, and bgr
def save(self):
summary_path = self.get_strpath_storage_for_file('summary.txt')
Path( summary_path ).write_text(self.model_summary_text)
Path( self.get_strpath_storage_for_file('summary.txt') ).write_text(self.model_summary_text)
self.onSave()
model_data = {
@ -373,56 +326,12 @@ 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) ]
if self.autobackup:
current_hour = time.localtime().tm_hour
if self.autobackup_current_hour != current_hour:
self.autobackup_current_hour = current_hour
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:
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:
shutil.copy ( str(filename), str(idx_backup_path / Path(filename).name) )
previews = self.get_previews()
plist = []
for i in range(len(previews)):
name, bgr = previews[i]
plist += [ (bgr, idx_backup_path / ( ('preview_%s.jpg') % (name)) ) ]
for preview, filepath in plist:
preview_lh = ModelBase.get_loss_history_preview(self.loss_history, self.iter, preview.shape[1], preview.shape[2])
img = (np.concatenate ( [preview_lh, preview], axis=0 ) * 255).astype(np.uint8)
cv2_imwrite (filepath, img )
def load_weights_safe(self, model_filename_list, optimizer_filename_list=[]):
loaded = []
not_loaded = []
for mf in model_filename_list:
model, filename = mf
for model, filename in model_filename_list:
filename = self.get_strpath_storage_for_file(filename)
if Path(filename).exists():
loaded += [ mf ]
model.load_weights(filename)
else:
not_loaded += [ mf ]
if len(optimizer_filename_list) != 0:
opt_filename = self.get_strpath_storage_for_file('opt.h5')
if Path(opt_filename).exists():
@ -439,19 +348,14 @@ class ModelBase(object):
print("set ok")
except Exception as e:
print ("Unable to load ", opt_filename)
return loaded, not_loaded
def save_weights_safe(self, model_filename_list):
def save_weights_safe(self, model_filename_list, optimizer_filename_list=[]):
for model, filename in model_filename_list:
filename = self.get_strpath_storage_for_file(filename)
model.save_weights( filename + '.tmp' )
rename_list = model_filename_list
"""
#unused
, optimizer_filename_list=[]
if len(optimizer_filename_list) != 0:
opt_filename = self.get_strpath_storage_for_file('opt.h5')
@ -471,7 +375,6 @@ class ModelBase(object):
rename_list += [('', 'opt.h5')]
except Exception as e:
print ("Unable to save ", opt_filename)
"""
for _, filename in rename_list:
filename = self.get_strpath_storage_for_file(filename)
@ -482,6 +385,7 @@ class ModelBase(object):
target_filename.unlink()
source_filename.rename ( str(target_filename) )
def debug_one_iter(self):
images = []
for generator in self.generator_list:
@ -492,7 +396,7 @@ class ModelBase(object):
return imagelib.equalize_and_stack_square (images)
def generate_next_sample(self):
return [ generator.generate_next() for generator in self.generator_list]
return [next(generator) for generator in self.generator_list]
def train_one_iter(self):
sample = self.generate_next_sample()
@ -587,47 +491,45 @@ class ModelBase(object):
lh_height = 100
lh_img = np.ones ( (lh_height,w,c) ) * 0.1
loss_count = len(loss_history[0])
lh_len = len(loss_history)
if len(loss_history) != 0:
loss_count = len(loss_history[0])
lh_len = len(loss_history)
l_per_col = lh_len / w
plist_max = [ [ max (0.0, loss_history[int(col*l_per_col)][p],
*[ loss_history[i_ab][p]
for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) )
]
)
for p in range(loss_count)
]
for col in range(w)
l_per_col = lh_len / w
plist_max = [ [ max (0.0, loss_history[int(col*l_per_col)][p],
*[ loss_history[i_ab][p]
for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) )
]
)
for p in range(loss_count)
]
for col in range(w)
]
plist_min = [ [ min (plist_max[col][p], loss_history[int(col*l_per_col)][p],
*[ loss_history[i_ab][p]
for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) )
]
)
for p in range(loss_count)
]
for col in range(w)
plist_min = [ [ min (plist_max[col][p], loss_history[int(col*l_per_col)][p],
*[ loss_history[i_ab][p]
for i_ab in range( int(col*l_per_col), int((col+1)*l_per_col) )
]
)
for p in range(loss_count)
]
for col in range(w)
]
plist_abs_max = np.mean(loss_history[ len(loss_history) // 5 : ]) * 2
plist_abs_max = np.mean(loss_history[ len(loss_history) // 5 : ]) * 2
for col in range(0, w):
for p in range(0,loss_count):
point_color = [1.0]*c
point_color[0:3] = colorsys.hsv_to_rgb ( p * (1.0/loss_count), 1.0, 1.0 )
for col in range(0, w):
for p in range(0,loss_count):
point_color = [1.0]*c
point_color[0:3] = colorsys.hsv_to_rgb ( p * (1.0/loss_count), 1.0, 1.0 )
ph_max = int ( (plist_max[col][p] / plist_abs_max) * (lh_height-1) )
ph_max = np.clip( ph_max, 0, lh_height-1 )
ph_max = int ( (plist_max[col][p] / plist_abs_max) * (lh_height-1) )
ph_max = np.clip( ph_max, 0, lh_height-1 )
ph_min = int ( (plist_min[col][p] / plist_abs_max) * (lh_height-1) )
ph_min = np.clip( ph_min, 0, lh_height-1 )
ph_min = int ( (plist_min[col][p] / plist_abs_max) * (lh_height-1) )
ph_min = np.clip( ph_min, 0, lh_height-1 )
for ph in range(ph_min, ph_max+1):
lh_img[ (lh_height-ph-1), col ] = point_color
for ph in range(ph_min, ph_max+1):
lh_img[ (lh_height-ph-1), col ] = point_color
lh_lines = 5
lh_line_height = (lh_height-1)/lh_lines

734
models/Model_AVATAR/Model.py
View file

@ -1,734 +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 *
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:head, 2:full_face ?:help skip:1) : ", 1, [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:'head',
2:'full_face'}[avatar_type]
self.options['avatar_type'] = avatar_type
else:
self.options['avatar_type'] = self.options.get('avatar_type', 'head')
if is_first_run or ask_override:
def_stage = self.options.get('stage', 1)
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', 1)
#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.get_model_filename_list():
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 self.predictor_func, (self.df_res, self.df_res, 3), converters.ConverterConfigFaceAvatar(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

5
models/Model_DEV_FANSEG/Model.py
View file

@ -11,7 +11,6 @@ class Model(ModelBase):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs,
ask_enable_autobackup=False,
ask_write_preview_history=False,
ask_target_iter=False,
ask_sort_by_yaw=False,
@ -48,7 +47,7 @@ class Model(ModelBase):
self.set_training_data_generators ([
SampleGeneratorFace(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=True),
output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR_SHUFFLE), 'resolution' : self.resolution, 'motion_blur':(25, 5), 'border_replicate':False },
output_sample_types=[ { 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_BGR_SHUFFLE), 'resolution' : self.resolution, 'motion_blur':(25, 1) },
{ 'types': (t.IMG_WARPED_TRANSFORMED, face_type, t.MODE_M), 'resolution': self.resolution },
]),
@ -66,7 +65,7 @@ class Model(ModelBase):
def onTrainOneIter(self, generators_samples, generators_list):
target_src, target_src_mask = generators_samples[0]
loss = self.fan_seg.train( target_src, target_src_mask )
loss = self.fan_seg.train_on_batch( [target_src], [target_src_mask] )
return ( ('loss', loss), )

1
models/Model_DEV_POSEEST/Model.py
View file

@ -12,7 +12,6 @@ class Model(ModelBase):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs,
ask_enable_autobackup=False,
ask_write_preview_history=False,
ask_target_iter=False,
ask_sort_by_yaw=False,

30
models/Model_DF/Model.py
View file

@ -59,16 +59,11 @@ 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() )
self.save_weights_safe( [[self.encoder, 'encoder.h5'],
[self.decoder_src, 'decoder_src.h5'],
[self.decoder_dst, 'decoder_dst.h5']] )
#override
def onTrainOneIter(self, sample, generators_list):
@ -111,17 +106,18 @@ class Model(ModelBase):
return [ ('DF', np.concatenate ( st, axis=0 ) ) ]
def predictor_func (self, face=None, dummy_predict=False):
if dummy_predict:
self.convert ([ np.zeros ( (1, 128, 128, 3), dtype=np.float32 ) ])
else:
x, mx = self.convert ( [ face[np.newaxis,...] ] )
return x[0], mx[0][...,0]
def predictor_func (self, face):
x, mx = self.convert ( [ face[np.newaxis,...] ] )
return x[0], mx[0][...,0]
#override
def get_ConverterConfig(self):
import converters
return self.predictor_func, (128,128,3), converters.ConverterConfigMasked(face_type=FaceType.FULL, default_mode=4)
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 Build(self, input_layer):
exec(nnlib.code_import_all, locals(), globals())

29
models/Model_H128/Model.py
View file

@ -70,15 +70,11 @@ class Model(ModelBase):
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() )
self.save_weights_safe( [[self.encoder, 'encoder.h5'],
[self.decoder_src, 'decoder_src.h5'],
[self.decoder_dst, 'decoder_dst.h5']] )
#override
def onTrainOneIter(self, sample, generators_list):
@ -119,17 +115,18 @@ class Model(ModelBase):
return [ ('H128', np.concatenate ( st, axis=0 ) ) ]
def predictor_func (self, face=None, dummy_predict=False):
if dummy_predict:
self.src_view ([ np.zeros ( (1, 128, 128, 3), dtype=np.float32 ) ])
else:
x, mx = self.src_view ( [ face[np.newaxis,...] ] )
return x[0], mx[0][...,0]
def predictor_func (self, face):
x, mx = self.src_view ( [ face[np.newaxis,...] ] )
return x[0], mx[0][...,0]
#override
def get_ConverterConfig(self):
import converters
return self.predictor_func, (128,128,3), converters.ConverterConfigMasked(face_type=FaceType.HALF, default_mode=4)
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 Build(self, lighter_ae):
exec(nnlib.code_import_all, locals(), globals())

29
models/Model_H64/Model.py
View file

@ -71,15 +71,11 @@ class Model(ModelBase):
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() )
self.save_weights_safe( [[self.encoder, 'encoder.h5'],
[self.decoder_src, 'decoder_src.h5'],
[self.decoder_dst, 'decoder_dst.h5']] )
#override
def onTrainOneIter(self, sample, generators_list):
@ -120,17 +116,18 @@ class Model(ModelBase):
return [ ('H64', np.concatenate ( st, axis=0 ) ) ]
def predictor_func (self, face=None, dummy_predict=False):
if dummy_predict:
self.src_view ([ np.zeros ( (1, 64, 64, 3), dtype=np.float32 ) ])
else:
x, mx = self.src_view ( [ face[np.newaxis,...] ] )
return x[0], mx[0][...,0]
def predictor_func (self, face):
x, mx = self.src_view ( [ face[np.newaxis,...] ] )
return x[0], mx[0][...,0]
#override
def get_ConverterConfig(self):
import converters
return self.predictor_func, (64,64,3), converters.ConverterConfigMasked(face_type=FaceType.HALF, default_mode=4)
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 Build(self, lighter_ae):
exec(nnlib.code_import_all, locals(), globals())

31
models/Model_LIAEF128/Model.py
View file

@ -65,16 +65,12 @@ class Model(ModelBase):
output_sample_types=output_sample_types)
])
#override
def get_model_filename_list(self):
return [[self.encoder, 'encoder.h5'],
[self.decoder, 'decoder.h5'],
[self.inter_B, 'inter_B.h5'],
[self.inter_AB, 'inter_AB.h5']]
#override
def onSave(self):
self.save_weights_safe( self.get_model_filename_list() )
self.save_weights_safe( [[self.encoder, 'encoder.h5'],
[self.decoder, 'decoder.h5'],
[self.inter_B, 'inter_B.h5'],
[self.inter_AB, 'inter_AB.h5']] )
#override
def onTrainOneIter(self, sample, generators_list):
@ -117,17 +113,18 @@ class Model(ModelBase):
return [ ('LIAEF128', np.concatenate ( st, axis=0 ) ) ]
def predictor_func (self, face=None, dummy_predict=False):
if dummy_predict:
self.convert ([ np.zeros ( (1, 128, 128, 3), dtype=np.float32 ) ])
else:
x, mx = self.convert ( [ face[np.newaxis,...] ] )
return x[0], mx[0][...,0]
def predictor_func (self, face):
x, mx = self.convert ( [ face[np.newaxis,...] ] )
return x[0], mx[0][...,0]
#override
def get_ConverterConfig(self):
import converters
return self.predictor_func, (128,128,3), converters.ConverterConfigMasked(face_type=FaceType.FULL, default_mode=4)
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 Build(self, input_layer):
exec(nnlib.code_import_all, locals(), globals())

479
models/Model_RecycleGAN/Model.py Normal file
View file

@ -0,0 +1,479 @@
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} ]
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, normalize_tanh = True),
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, normalize_tanh = True),
output_sample_types=output_sample_types ),
])
else:
self.G_convert = K.function([real_B0],[fake_A0])
#override
def onSave(self):
self.save_weights_safe( [[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 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
nnlib.UNet = UNet
@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

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

790
models/Model_SAE/Model.py
View file

@ -1,18 +1,26 @@
from functools import partial
import numpy as np
import mathlib
from facelib import FaceType
from interact import interact as io
from models import ModelBase
from nnlib import nnlib
from models import ModelBase
from facelib import FaceType
from samplelib import *
from interact import interact as io
#SAE - Styled AutoEncoder
class SAEModel(ModelBase):
encoderH5 = 'encoder.h5'
inter_BH5 = 'inter_B.h5'
inter_ABH5 = 'inter_AB.h5'
decoderH5 = 'decoder.h5'
decodermH5 = 'decoderm.h5'
decoder_srcH5 = 'decoder_src.h5'
decoder_srcmH5 = 'decoder_srcm.h5'
decoder_dstH5 = 'decoder_dst.h5'
decoder_dstmH5 = 'decoder_dstm.h5'
#override
def onInitializeOptions(self, is_first_run, ask_override):
yn_str = {True:'y',False:'n'}
@ -20,7 +28,6 @@ class SAEModel(ModelBase):
default_resolution = 128
default_archi = 'df'
default_face_type = 'f'
default_learn_mask = True
if is_first_run:
resolution = io.input_int("Resolution ( 64-256 ?:help skip:128) : ", default_resolution, help_message="More resolution requires more VRAM and time to train. Value will be adjusted to multiple of 16.")
@ -30,11 +37,12 @@ class SAEModel(ModelBase):
self.options['resolution'] = resolution
self.options['face_type'] = io.input_str ("Half or Full face? (h/f, ?:help skip:f) : ", default_face_type, ['h','f'], help_message="Half face has better resolution, but covers less area of cheeks.").lower()
self.options['learn_mask'] = io.input_bool ( f"Learn mask? (y/n, ?:help skip:{yn_str[default_learn_mask]} ) : " , default_learn_mask, help_message="Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case converter forced to use 'not predicted mask' that is not smooth as predicted. Model with style values can be learned without mask and produce same quality result.")
self.options['learn_mask'] = io.input_bool ("Learn mask? (y/n, ?:help skip:y) : ", True, help_message="Learning mask can help model to recognize face directions. Learn without mask can reduce model size, in this case converter forced to use 'not predicted mask' that is not smooth as predicted. Model with style values can be learned without mask and produce same quality result.")
else:
self.options['resolution'] = self.options.get('resolution', default_resolution)
self.options['face_type'] = self.options.get('face_type', default_face_type)
self.options['learn_mask'] = self.options.get('learn_mask', default_learn_mask)
self.options['learn_mask'] = self.options.get('learn_mask', True)
if (is_first_run or ask_override) and 'tensorflow' in self.device_config.backend:
def_optimizer_mode = self.options.get('optimizer_mode', 1)
@ -51,24 +59,26 @@ class SAEModel(ModelBase):
default_e_ch_dims = 42
default_d_ch_dims = default_e_ch_dims // 2
def_ca_weights = False
if is_first_run:
self.options['ae_dims'] = np.clip ( io.input_int("AutoEncoder dims (32-1024 ?:help skip:%d) : " % (default_ae_dims) , default_ae_dims, help_message="All face information will packed to AE dims. If amount of AE dims are not enough, then for example closed eyes will not be recognized. More dims are better, but require more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 )
self.options['e_ch_dims'] = np.clip ( io.input_int("Encoder dims per channel (21-85 ?:help skip:%d) : " % (default_e_ch_dims) , default_e_ch_dims, help_message="More encoder dims help to recognize more facial features, but require more VRAM. You can fine-tune model size to fit your GPU." ), 21, 85 )
default_d_ch_dims = self.options['e_ch_dims'] // 2
self.options['d_ch_dims'] = np.clip ( io.input_int("Decoder dims per channel (10-85 ?:help skip:%d) : " % (default_d_ch_dims) , default_d_ch_dims, help_message="More decoder dims help to get better details, but require more VRAM. You can fine-tune model size to fit your GPU." ), 10, 85 )
self.options['ca_weights'] = io.input_bool (f"Use CA weights? (y/n, ?:help skip:{yn_str[def_ca_weights]} ) : ", def_ca_weights, help_message="Initialize network with 'Convolution Aware' weights. This may help to achieve a higher accuracy model, but consumes a time at first run.")
self.options['multiscale_decoder'] = io.input_bool ("Use multiscale decoder? (y/n, ?:help skip:n) : ", False, help_message="Multiscale decoder helps to get better details.")
self.options['ca_weights'] = io.input_bool ("Use CA weights? (y/n, ?:help skip: %s ) : " % (yn_str[def_ca_weights]), def_ca_weights, help_message="Initialize network with 'Convolution Aware' weights. This may help to achieve a higher accuracy model, but consumes a time at first run.")
else:
self.options['ae_dims'] = self.options.get('ae_dims', default_ae_dims)
self.options['e_ch_dims'] = self.options.get('e_ch_dims', default_e_ch_dims)
self.options['d_ch_dims'] = self.options.get('d_ch_dims', default_d_ch_dims)
self.options['multiscale_decoder'] = self.options.get('multiscale_decoder', False)
self.options['ca_weights'] = self.options.get('ca_weights', def_ca_weights)
default_face_style_power = 0.0
default_bg_style_power = 0.0
if is_first_run or ask_override:
def_pixel_loss = self.options.get('pixel_loss', False)
self.options['pixel_loss'] = io.input_bool (f"Use pixel loss? (y/n, ?:help skip:{yn_str[def_pixel_loss]} ) : ", def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time. Enabling this option too early increases the chance of model collapse.")
self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: %s ) : " % (yn_str[def_pixel_loss]), def_pixel_loss, help_message="Pixel loss may help to enhance fine details and stabilize face color. Use it only if quality does not improve over time. Enabling this option too early increases the chance of model collapse.")
default_face_style_power = default_face_style_power if is_first_run else self.options.get('face_style_power', default_face_style_power)
self.options['face_style_power'] = np.clip ( io.input_number("Face style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_face_style_power), default_face_style_power,
@ -79,20 +89,13 @@ class SAEModel(ModelBase):
help_message="Learn to transfer image around face. This can make face more like dst. Enabling this option increases the chance of model collapse."), 0.0, 100.0 )
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 (f"Apply random color transfer to src faceset? (y/n, ?:help skip:{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 (f"Enable gradient clipping? (y/n, ?:help skip:{yn_str[default_clipgrad]}) : ", default_clipgrad, help_message="Gradient clipping reduces chance of model collapse, sacrificing speed of training.")
else:
self.options['clipgrad'] = 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 RCT color transfer from random dst samples.")
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.")
@ -102,11 +105,10 @@ class SAEModel(ModelBase):
#override
def onInitialize(self):
exec(nnlib.import_all(), locals(), globals())
SAEModel.initialize_nn_functions()
self.set_vram_batch_requirements({1.5:4})
resolution = self.options['resolution']
learn_mask = self.options['learn_mask']
ae_dims = self.options['ae_dims']
e_ch_dims = self.options['e_ch_dims']
d_ch_dims = self.options['d_ch_dims']
@ -118,333 +120,232 @@ class SAEModel(ModelBase):
bgr_shape = (resolution, resolution, 3)
mask_shape = (resolution, resolution, 1)
self.ms_count = ms_count = 3 if (self.options['multiscale_decoder']) else 1
apply_random_ct = self.options.get('apply_random_ct', False)
masked_training = True
class SAEDFModel(object):
def __init__(self, resolution, ae_dims, e_ch_dims, d_ch_dims, learn_mask):
super().__init__()
self.learn_mask = learn_mask
warped_src = Input(bgr_shape)
target_src = Input(bgr_shape)
target_srcm = Input(mask_shape)
output_nc = 3
bgr_shape = (resolution, resolution, output_nc)
mask_shape = (resolution, resolution, 1)
lowest_dense_res = resolution // 16
e_dims = output_nc*e_ch_dims
warped_dst = Input(bgr_shape)
target_dst = Input(bgr_shape)
target_dstm = Input(mask_shape)
def upscale (dim):
def func(x):
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same')(x)))
return func
target_src_ar = [ Input ( ( bgr_shape[0] // (2**i) ,)*2 + (bgr_shape[-1],) ) for i in range(ms_count-1, -1, -1)]
target_srcm_ar = [ Input ( ( mask_shape[0] // (2**i) ,)*2 + (mask_shape[-1],) ) for i in range(ms_count-1, -1, -1)]
target_dst_ar = [ Input ( ( bgr_shape[0] // (2**i) ,)*2 + (bgr_shape[-1],) ) for i in range(ms_count-1, -1, -1)]
target_dstm_ar = [ Input ( ( mask_shape[0] // (2**i) ,)*2 + (mask_shape[-1],) ) for i in range(ms_count-1, -1, -1)]
def enc_flow(e_dims, ae_dims, lowest_dense_res):
def func(x):
x = LeakyReLU(0.1)(Conv2D(e_dims, kernel_size=5, strides=2, padding='same')(x))
x = LeakyReLU(0.1)(Conv2D(e_dims*2, kernel_size=5, strides=2, padding='same')(x))
x = LeakyReLU(0.1)(Conv2D(e_dims*4, kernel_size=5, strides=2, padding='same')(x))
x = LeakyReLU(0.1)(Conv2D(e_dims*8, kernel_size=5, strides=2, padding='same')(x))
common_flow_kwargs = { 'padding': 'zero',
'norm': '',
'act':'' }
models_list = []
weights_to_load = []
if 'liae' in self.options['archi']:
self.encoder = modelify(SAEModel.LIAEEncFlow(resolution, ch_dims=e_ch_dims, **common_flow_kwargs) ) (Input(bgr_shape))
x = Dense(ae_dims)(Flatten()(x))
x = Dense(lowest_dense_res * lowest_dense_res * ae_dims)(x)
x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims))(x)
x = upscale(ae_dims)(x)
return x
return func
enc_output_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ]
def dec_flow(output_nc, d_ch_dims):
def ResidualBlock(dim):
def func(inp):
x = Conv2D(dim, kernel_size=3, padding='same')(inp)
x = LeakyReLU(0.2)(x)
x = Conv2D(dim, kernel_size=3, padding='same')(x)
x = Add()([x, inp])
x = LeakyReLU(0.2)(x)
return x
return func
self.inter_B = modelify(SAEModel.LIAEInterFlow(resolution, ae_dims=ae_dims, **common_flow_kwargs)) (enc_output_Inputs)
self.inter_AB = modelify(SAEModel.LIAEInterFlow(resolution, ae_dims=ae_dims, **common_flow_kwargs)) (enc_output_Inputs)
def func(x):
dims = output_nc * d_ch_dims
x = upscale(dims*8)(x)
x = ResidualBlock(dims*8)(x)
x = ResidualBlock(dims*8)(x)
inter_output_Inputs = [ Input( np.array(K.int_shape(x)[1:])*(1,1,2) ) for x in self.inter_B.outputs ]
x = upscale(dims*4)(x)
x = ResidualBlock(dims*4)(x)
x = ResidualBlock(dims*4)(x)
self.decoder = modelify(SAEModel.LIAEDecFlow (bgr_shape[2],ch_dims=d_ch_dims, multiscale_count=self.ms_count, add_residual_blocks=d_residual_blocks, **common_flow_kwargs)) (inter_output_Inputs)
models_list += [self.encoder, self.inter_B, self.inter_AB, self.decoder]
x = upscale(dims*2)(x)
x = ResidualBlock(dims*2)(x)
x = ResidualBlock(dims*2)(x)
if self.options['learn_mask']:
self.decoderm = modelify(SAEModel.LIAEDecFlow (mask_shape[2],ch_dims=d_ch_dims, **common_flow_kwargs)) (inter_output_Inputs)
models_list += [self.decoderm]
return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid')(x)
return func
if not self.is_first_run():
weights_to_load += [ [self.encoder , 'encoder.h5'],
[self.inter_B , 'inter_B.h5'],
[self.inter_AB, 'inter_AB.h5'],
[self.decoder , 'decoder.h5'],
]
if self.options['learn_mask']:
weights_to_load += [ [self.decoderm, 'decoderm.h5'] ]
self.encoder = modelify(enc_flow(e_dims, ae_dims, lowest_dense_res)) ( Input(bgr_shape) )
warped_src_code = self.encoder (warped_src)
warped_src_inter_AB_code = self.inter_AB (warped_src_code)
warped_src_inter_code = Concatenate()([warped_src_inter_AB_code,warped_src_inter_AB_code])
sh = K.int_shape( self.encoder.outputs[0] )[1:]
self.decoder_src = modelify(dec_flow(output_nc, d_ch_dims)) ( Input(sh) )
self.decoder_dst = modelify(dec_flow(output_nc, d_ch_dims)) ( Input(sh) )
warped_dst_code = self.encoder (warped_dst)
warped_dst_inter_B_code = self.inter_B (warped_dst_code)
warped_dst_inter_AB_code = self.inter_AB (warped_dst_code)
warped_dst_inter_code = Concatenate()([warped_dst_inter_B_code,warped_dst_inter_AB_code])
if learn_mask:
self.decoder_srcm = modelify(dec_flow(1, d_ch_dims)) ( Input(sh) )
self.decoder_dstm = modelify(dec_flow(1, d_ch_dims)) ( Input(sh) )
warped_src_dst_inter_code = Concatenate()([warped_dst_inter_AB_code,warped_dst_inter_AB_code])
self.src_dst_trainable_weights = self.encoder.trainable_weights + self.decoder_src.trainable_weights + self.decoder_dst.trainable_weights
pred_src_src = self.decoder(warped_src_inter_code)
pred_dst_dst = self.decoder(warped_dst_inter_code)
pred_src_dst = self.decoder(warped_src_dst_inter_code)
if learn_mask:
self.src_dst_mask_trainable_weights = self.encoder.trainable_weights + self.decoder_srcm.trainable_weights + self.decoder_dstm.trainable_weights
if self.options['learn_mask']:
pred_src_srcm = self.decoderm(warped_src_inter_code)
pred_dst_dstm = self.decoderm(warped_dst_inter_code)
pred_src_dstm = self.decoderm(warped_src_dst_inter_code)
self.warped_src, self.warped_dst = Input(bgr_shape), Input(bgr_shape)
src_code, dst_code = self.encoder(self.warped_src), self.encoder(self.warped_dst)
elif 'df' in self.options['archi']:
self.encoder = modelify(SAEModel.DFEncFlow(resolution, ae_dims=ae_dims, ch_dims=e_ch_dims, **common_flow_kwargs) ) (Input(bgr_shape))
self.pred_src_src = self.decoder_src(src_code)
self.pred_dst_dst = self.decoder_dst(dst_code)
self.pred_src_dst = self.decoder_src(dst_code)
dec_Inputs = [ Input(K.int_shape(x)[1:]) for x in self.encoder.outputs ]
if learn_mask:
self.pred_src_srcm = self.decoder_srcm(src_code)
self.pred_dst_dstm = self.decoder_dstm(dst_code)
self.pred_src_dstm = self.decoder_srcm(dst_code)
self.decoder_src = modelify(SAEModel.DFDecFlow (bgr_shape[2],ch_dims=d_ch_dims, multiscale_count=self.ms_count, add_residual_blocks=d_residual_blocks, **common_flow_kwargs )) (dec_Inputs)
self.decoder_dst = modelify(SAEModel.DFDecFlow (bgr_shape[2],ch_dims=d_ch_dims, multiscale_count=self.ms_count, add_residual_blocks=d_residual_blocks, **common_flow_kwargs )) (dec_Inputs)
models_list += [self.encoder, self.decoder_src, self.decoder_dst]
def get_model_filename_list(self, exclude_for_pretrain=False):
ar = []
if not exclude_for_pretrain:
ar += [ [self.encoder, 'encoder.h5'] ]
ar += [ [self.decoder_src, 'decoder_src.h5'],
[self.decoder_dst, 'decoder_dst.h5'] ]
if self.learn_mask:
ar += [ [self.decoder_srcm, 'decoder_srcm.h5'],
[self.decoder_dstm, 'decoder_dstm.h5'] ]
return ar
if self.options['learn_mask']:
self.decoder_srcm = modelify(SAEModel.DFDecFlow (mask_shape[2],ch_dims=d_ch_dims, **common_flow_kwargs )) (dec_Inputs)
self.decoder_dstm = modelify(SAEModel.DFDecFlow (mask_shape[2],ch_dims=d_ch_dims, **common_flow_kwargs )) (dec_Inputs)
models_list += [self.decoder_srcm, self.decoder_dstm]
class SAELIAEModel(object):
def __init__(self, resolution, ae_dims, e_ch_dims, d_ch_dims, learn_mask):
super().__init__()
self.learn_mask = learn_mask
if not self.is_first_run():
weights_to_load += [ [self.encoder , 'encoder.h5'],
[self.decoder_src, 'decoder_src.h5'],
[self.decoder_dst, 'decoder_dst.h5']
]
if self.options['learn_mask']:
weights_to_load += [ [self.decoder_srcm, 'decoder_srcm.h5'],
[self.decoder_dstm, 'decoder_dstm.h5'],
]
output_nc = 3
bgr_shape = (resolution, resolution, output_nc)
mask_shape = (resolution, resolution, 1)
warped_src_code = self.encoder (warped_src)
warped_dst_code = self.encoder (warped_dst)
pred_src_src = self.decoder_src(warped_src_code)
pred_dst_dst = self.decoder_dst(warped_dst_code)
pred_src_dst = self.decoder_src(warped_dst_code)
e_dims = output_nc*e_ch_dims
d_dims = output_nc*d_ch_dims
lowest_dense_res = resolution // 16
if self.options['learn_mask']:
pred_src_srcm = self.decoder_srcm(warped_src_code)
pred_dst_dstm = self.decoder_dstm(warped_dst_code)
pred_src_dstm = self.decoder_srcm(warped_dst_code)
def upscale (dim):
def func(x):
return SubpixelUpscaler()(LeakyReLU(0.1)(Conv2D(dim * 4, kernel_size=3, strides=1, padding='same')(x)))
return func
if self.is_first_run():
if self.options.get('ca_weights',False):
conv_weights_list = []
for model in models_list:
for layer in model.layers:
if type(layer) == keras.layers.Conv2D:
conv_weights_list += [layer.weights[0]] #Conv2D kernel_weights
CAInitializerMP ( conv_weights_list )
else:
self.load_weights_safe(weights_to_load)
def enc_flow(e_dims):
def func(x):
x = LeakyReLU(0.1)(Conv2D(e_dims, kernel_size=5, strides=2, padding='same')(x))
x = LeakyReLU(0.1)(Conv2D(e_dims*2, kernel_size=5, strides=2, padding='same')(x))
x = LeakyReLU(0.1)(Conv2D(e_dims*4, kernel_size=5, strides=2, padding='same')(x))
x = LeakyReLU(0.1)(Conv2D(e_dims*8, kernel_size=5, strides=2, padding='same')(x))
x = Flatten()(x)
return x
return func
pred_src_src, pred_dst_dst, pred_src_dst, = [ [x] if type(x) != list else x for x in [pred_src_src, pred_dst_dst, pred_src_dst, ] ]
def inter_flow(lowest_dense_res, ae_dims):
def func(x):
x = Dense(ae_dims)(x)
x = Dense(lowest_dense_res * lowest_dense_res * ae_dims*2)(x)
x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims*2))(x)
x = upscale(ae_dims*2)(x)
return x
return func
if self.options['learn_mask']:
pred_src_srcm, pred_dst_dstm, pred_src_dstm = [ [x] if type(x) != list else x for x in [pred_src_srcm, pred_dst_dstm, pred_src_dstm] ]
def dec_flow(output_nc, d_dims):
def ResidualBlock(dim):
def func(inp):
x = Conv2D(dim, kernel_size=3, padding='same')(inp)
x = LeakyReLU(0.2)(x)
x = Conv2D(dim, kernel_size=3, padding='same')(x)
x = Add()([x, inp])
x = LeakyReLU(0.2)(x)
return x
return func
target_srcm_blurred_ar = [ gaussian_blur( max(1, K.int_shape(x)[1] // 32) )(x) for x in target_srcm_ar]
target_srcm_sigm_ar = target_srcm_blurred_ar #[ x / 2.0 + 0.5 for x in target_srcm_blurred_ar]
target_srcm_anti_sigm_ar = [ 1.0 - x for x in target_srcm_sigm_ar]
def func(x):
x = upscale(d_dims*8)(x)
x = ResidualBlock(d_dims*8)(x)
x = ResidualBlock(d_dims*8)(x)
target_dstm_blurred_ar = [ gaussian_blur( max(1, K.int_shape(x)[1] // 32) )(x) for x in target_dstm_ar]
target_dstm_sigm_ar = target_dstm_blurred_ar#[ x / 2.0 + 0.5 for x in target_dstm_blurred_ar]
target_dstm_anti_sigm_ar = [ 1.0 - x for x in target_dstm_sigm_ar]
x = upscale(d_dims*4)(x)
x = ResidualBlock(d_dims*4)(x)
x = ResidualBlock(d_dims*4)(x)
target_src_sigm_ar = target_src_ar#[ x + 1 for x in target_src_ar]
target_dst_sigm_ar = target_dst_ar#[ x + 1 for x in target_dst_ar]
x = upscale(d_dims*2)(x)
x = ResidualBlock(d_dims*2)(x)
x = ResidualBlock(d_dims*2)(x)
pred_src_src_sigm_ar = pred_src_src#[ x + 1 for x in pred_src_src]
pred_dst_dst_sigm_ar = pred_dst_dst#[ x + 1 for x in pred_dst_dst]
pred_src_dst_sigm_ar = pred_src_dst#[ x + 1 for x in pred_src_dst]
return Conv2D(output_nc, kernel_size=5, padding='same', activation='sigmoid')(x)
return func
target_src_masked_ar = [ target_src_sigm_ar[i]*target_srcm_sigm_ar[i] for i in range(len(target_src_sigm_ar))]
target_dst_masked_ar = [ target_dst_sigm_ar[i]*target_dstm_sigm_ar[i] for i in range(len(target_dst_sigm_ar))]
target_dst_anti_masked_ar = [ target_dst_sigm_ar[i]*target_dstm_anti_sigm_ar[i] for i in range(len(target_dst_sigm_ar))]
self.encoder = modelify(enc_flow(e_dims)) ( Input(bgr_shape) )
pred_src_src_masked_ar = [ pred_src_src_sigm_ar[i] * target_srcm_sigm_ar[i] for i in range(len(pred_src_src_sigm_ar))]
pred_dst_dst_masked_ar = [ pred_dst_dst_sigm_ar[i] * target_dstm_sigm_ar[i] for i in range(len(pred_dst_dst_sigm_ar))]
sh = K.int_shape( self.encoder.outputs[0] )[1:]
self.inter_B = modelify(inter_flow(lowest_dense_res, ae_dims)) ( Input(sh) )
self.inter_AB = modelify(inter_flow(lowest_dense_res, ae_dims)) ( Input(sh) )
target_src_masked_ar_opt = target_src_masked_ar if masked_training else target_src_sigm_ar
target_dst_masked_ar_opt = target_dst_masked_ar if masked_training else target_dst_sigm_ar
sh = np.array(K.int_shape( self.inter_B.outputs[0] )[1:])*(1,1,2)
self.decoder = modelify(dec_flow(output_nc, d_dims)) ( Input(sh) )
pred_src_src_masked_ar_opt = pred_src_src_masked_ar if masked_training else pred_src_src_sigm_ar
pred_dst_dst_masked_ar_opt = pred_dst_dst_masked_ar if masked_training else pred_dst_dst_sigm_ar
if learn_mask:
self.decoderm = modelify(dec_flow(1, d_dims)) ( Input(sh) )
self.src_dst_trainable_weights = self.encoder.trainable_weights + self.inter_B.trainable_weights + self.inter_AB.trainable_weights + self.decoder.trainable_weights
if learn_mask:
self.src_dst_mask_trainable_weights = self.encoder.trainable_weights + self.inter_B.trainable_weights + self.inter_AB.trainable_weights + self.decoderm.trainable_weights
self.warped_src, self.warped_dst = Input(bgr_shape), Input(bgr_shape)
warped_src_code = self.encoder (self.warped_src)
warped_src_inter_AB_code = self.inter_AB (warped_src_code)
warped_src_inter_code = Concatenate()([warped_src_inter_AB_code,warped_src_inter_AB_code])
warped_dst_code = self.encoder (self.warped_dst)
warped_dst_inter_B_code = self.inter_B (warped_dst_code)
warped_dst_inter_AB_code = self.inter_AB (warped_dst_code)
warped_dst_inter_code = Concatenate()([warped_dst_inter_B_code,warped_dst_inter_AB_code])
warped_src_dst_inter_code = Concatenate()([warped_dst_inter_AB_code,warped_dst_inter_AB_code])
self.pred_src_src = self.decoder(warped_src_inter_code)
self.pred_dst_dst = self.decoder(warped_dst_inter_code)
self.pred_src_dst = self.decoder(warped_src_dst_inter_code)
if learn_mask:
self.pred_src_srcm = self.decoderm(warped_src_inter_code)
self.pred_dst_dstm = self.decoderm(warped_dst_inter_code)
self.pred_src_dstm = self.decoderm(warped_src_dst_inter_code)
def get_model_filename_list(self, exclude_for_pretrain=False):
ar = [ [self.encoder, 'encoder.h5'],
[self.inter_B, 'inter_B.h5'] ]
if not exclude_for_pretrain:
ar += [ [self.inter_AB, 'inter_AB.h5'] ]
ar += [ [self.decoder, 'decoder.h5'] ]
if self.learn_mask:
ar += [ [self.decoderm, 'decoderm.h5'] ]
return ar
if 'df' in self.options['archi']:
self.model = SAEDFModel (resolution, ae_dims, e_ch_dims, d_ch_dims, learn_mask)
elif 'liae' in self.options['archi']:
self.model = SAELIAEModel (resolution, ae_dims, e_ch_dims, d_ch_dims, learn_mask)
loaded, not_loaded = [], self.model.get_model_filename_list()
if not self.is_first_run():
loaded, not_loaded = self.load_weights_safe(not_loaded)
CA_models = []
if self.options.get('ca_weights', False):
CA_models += [ model for model, _ in not_loaded ]
CA_conv_weights_list = []
for model in CA_models:
for layer in model.layers:
if type(layer) == keras.layers.Conv2D:
CA_conv_weights_list += [layer.weights[0]] #- is Conv2D kernel_weights
if len(CA_conv_weights_list) != 0:
CAInitializerMP ( CA_conv_weights_list )
warped_src = self.model.warped_src
target_src = Input ( (resolution, resolution, 3) )
target_srcm = Input ( (resolution, resolution, 1) )
warped_dst = self.model.warped_dst
target_dst = Input ( (resolution, resolution, 3) )
target_dstm = Input ( (resolution, resolution, 1) )
target_src_sigm = target_src
target_dst_sigm = target_dst
target_srcm_sigm = gaussian_blur( max(1, K.int_shape(target_srcm)[1] // 32) )(target_srcm)
target_dstm_sigm = gaussian_blur( max(1, K.int_shape(target_dstm)[1] // 32) )(target_dstm)
target_dstm_anti_sigm = 1.0 - target_dstm_sigm
target_src_masked = target_src_sigm*target_srcm_sigm
target_dst_masked = target_dst_sigm*target_dstm_sigm
target_dst_anti_masked = target_dst_sigm*target_dstm_anti_sigm
target_src_masked_opt = target_src_masked if masked_training else target_src_sigm
target_dst_masked_opt = target_dst_masked if masked_training else target_dst_sigm
pred_src_src = self.model.pred_src_src
pred_dst_dst = self.model.pred_dst_dst
pred_src_dst = self.model.pred_src_dst
if learn_mask:
pred_src_srcm = self.model.pred_src_srcm
pred_dst_dstm = self.model.pred_dst_dstm
pred_src_dstm = self.model.pred_src_dstm
pred_src_src_sigm = self.model.pred_src_src
pred_dst_dst_sigm = self.model.pred_dst_dst
pred_src_dst_sigm = self.model.pred_src_dst
pred_src_src_masked = pred_src_src_sigm*target_srcm_sigm
pred_dst_dst_masked = pred_dst_dst_sigm*target_dstm_sigm
pred_src_src_masked_opt = pred_src_src_masked if masked_training else pred_src_src_sigm
pred_dst_dst_masked_opt = pred_dst_dst_masked if masked_training else pred_dst_dst_sigm
psd_target_dst_masked = pred_src_dst_sigm*target_dstm_sigm
psd_target_dst_anti_masked = pred_src_dst_sigm*target_dstm_anti_sigm
psd_target_dst_masked_ar = [ pred_src_dst_sigm_ar[i]*target_dstm_sigm_ar[i] for i in range(len(pred_src_dst_sigm_ar))]
psd_target_dst_anti_masked_ar = [ pred_src_dst_sigm_ar[i]*target_dstm_anti_sigm_ar[i] for i in range(len(pred_src_dst_sigm_ar))]
if self.is_training_mode:
self.src_dst_opt = Adam(lr=5e-5, beta_1=0.5, beta_2=0.999, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1)
self.src_dst_mask_opt = Adam(lr=5e-5, beta_1=0.5, beta_2=0.999, clipnorm=1.0 if self.options['clipgrad'] else 0.0, tf_cpu_mode=self.options['optimizer_mode']-1)
self.sr_opt = Adam(lr=5e-5, beta_1=0.9, beta_2=0.999, tf_cpu_mode=self.options['optimizer_mode']-1)
self.src_dst_opt = Adam(lr=5e-5, beta_1=0.5, beta_2=0.999, tf_cpu_mode=self.options['optimizer_mode']-1)
self.src_dst_mask_opt = Adam(lr=5e-5, beta_1=0.5, beta_2=0.999, tf_cpu_mode=self.options['optimizer_mode']-1)
if 'liae' in self.options['archi']:
src_dst_loss_train_weights = self.encoder.trainable_weights + self.inter_B.trainable_weights + self.inter_AB.trainable_weights + self.decoder.trainable_weights
if self.options['learn_mask']:
src_dst_mask_loss_train_weights = self.encoder.trainable_weights + self.inter_B.trainable_weights + self.inter_AB.trainable_weights + self.decoderm.trainable_weights
else:
src_dst_loss_train_weights = self.encoder.trainable_weights + self.decoder_src.trainable_weights + self.decoder_dst.trainable_weights
if self.options['learn_mask']:
src_dst_mask_loss_train_weights = self.encoder.trainable_weights + self.decoder_srcm.trainable_weights + self.decoder_dstm.trainable_weights
if not self.options['pixel_loss']:
src_loss = K.mean ( 10*dssim(kernel_size=int(resolution/11.6),max_value=1.0)( target_src_masked_opt, pred_src_src_masked_opt) )
src_loss_batch = sum([ 10*dssim(kernel_size=int(resolution/11.6),max_value=1.0)( target_src_masked_ar_opt[i], pred_src_src_masked_ar_opt[i]) for i in range(len(target_src_masked_ar_opt)) ])
else:
src_loss = K.mean ( 50*K.square( target_src_masked_opt - pred_src_src_masked_opt ) )
src_loss_batch = sum([ K.mean ( 50*K.square( target_src_masked_ar_opt[i] - pred_src_src_masked_ar_opt[i] ), axis=[1,2,3]) for i in range(len(target_src_masked_ar_opt)) ])
src_loss = K.mean(src_loss_batch)
face_style_power = self.options['face_style_power'] / 100.0
face_style_power = self.options['face_style_power'] / 100.0
if face_style_power != 0:
src_loss += style_loss(gaussian_blur_radius=resolution//16, loss_weight=face_style_power, wnd_size=0)( psd_target_dst_masked, target_dst_masked )
src_loss += style_loss(gaussian_blur_radius=resolution//16, loss_weight=face_style_power, wnd_size=0)( psd_target_dst_masked_ar[-1], target_dst_masked_ar[-1] )
bg_style_power = self.options['bg_style_power'] / 100.0
if bg_style_power != 0:
if not self.options['pixel_loss']:
src_loss += K.mean( (10*bg_style_power)*dssim(kernel_size=int(resolution/11.6),max_value=1.0)( psd_target_dst_anti_masked, target_dst_anti_masked ))
bg_loss = K.mean( (10*bg_style_power)*dssim(kernel_size=int(resolution/11.6),max_value=1.0)( psd_target_dst_anti_masked_ar[-1], target_dst_anti_masked_ar[-1] ))
else:
src_loss += K.mean( (50*bg_style_power)*K.square( psd_target_dst_anti_masked - target_dst_anti_masked ))
bg_loss = K.mean( (50*bg_style_power)*K.square( psd_target_dst_anti_masked_ar[-1] - target_dst_anti_masked_ar[-1] ))
src_loss += bg_loss
if not self.options['pixel_loss']:
dst_loss = K.mean( 10*dssim(kernel_size=int(resolution/11.6),max_value=1.0)(target_dst_masked_opt, pred_dst_dst_masked_opt) )
dst_loss_batch = sum([ 10*dssim(kernel_size=int(resolution/11.6),max_value=1.0)(target_dst_masked_ar_opt[i], pred_dst_dst_masked_ar_opt[i]) for i in range(len(target_dst_masked_ar_opt)) ])
else:
dst_loss = K.mean( 50*K.square( target_dst_masked_opt - pred_dst_dst_masked_opt ) )
dst_loss_batch = sum([ K.mean ( 50*K.square( target_dst_masked_ar_opt[i] - pred_dst_dst_masked_ar_opt[i] ), axis=[1,2,3]) for i in range(len(target_dst_masked_ar_opt)) ])
self.src_dst_train = K.function ([warped_src, warped_dst, target_src, target_srcm, target_dst, target_dstm],[src_loss,dst_loss], self.src_dst_opt.get_updates(src_loss+dst_loss, self.model.src_dst_trainable_weights) )
dst_loss = K.mean(dst_loss_batch)
feed = [warped_src, warped_dst]
feed += target_src_ar[::-1]
feed += target_srcm_ar[::-1]
feed += target_dst_ar[::-1]
feed += target_dstm_ar[::-1]
self.src_dst_train = K.function (feed,[src_loss,dst_loss], self.src_dst_opt.get_updates(src_loss+dst_loss, src_dst_loss_train_weights) )
if self.options['learn_mask']:
src_mask_loss = K.mean(K.square(target_srcm-pred_src_srcm))
dst_mask_loss = K.mean(K.square(target_dstm-pred_dst_dstm))
self.src_dst_mask_train = K.function ([warped_src, warped_dst, target_srcm, target_dstm],[src_mask_loss, dst_mask_loss], self.src_dst_mask_opt.get_updates(src_mask_loss+dst_mask_loss, self.model.src_dst_mask_trainable_weights ) )
src_mask_loss = sum([ K.mean(K.square(target_srcm_ar[-1]-pred_src_srcm[-1])) for i in range(len(target_srcm_ar)) ])
dst_mask_loss = sum([ K.mean(K.square(target_dstm_ar[-1]-pred_dst_dstm[-1])) for i in range(len(target_dstm_ar)) ])
feed = [ warped_src, warped_dst]
feed += target_srcm_ar[::-1]
feed += target_dstm_ar[::-1]
self.src_dst_mask_train = K.function (feed,[src_mask_loss, dst_mask_loss], self.src_dst_mask_opt.get_updates(src_mask_loss+dst_mask_loss, src_dst_mask_loss_train_weights) )
if self.options['learn_mask']:
self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src, pred_dst_dst, pred_dst_dstm, pred_src_dst, pred_src_dstm])
self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src[-1], pred_dst_dst[-1], pred_dst_dstm[-1], pred_src_dst[-1], pred_src_dstm[-1]])
else:
self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src, pred_dst_dst, pred_src_dst ])
self.AE_view = K.function ([warped_src, warped_dst], [pred_src_src[-1], pred_dst_dst[-1], pred_src_dst[-1] ] )
else:
if self.options['learn_mask']:
self.AE_convert = K.function ([warped_dst],[ pred_src_dst, pred_dst_dstm, pred_src_dstm ])
self.AE_convert = K.function ([warped_dst],[ pred_src_dst[-1], pred_dst_dstm[-1], pred_src_dstm[-1] ])
else:
self.AE_convert = K.function ([warped_dst],[ pred_src_dst ])
self.AE_convert = K.function ([warped_dst],[ pred_src_dst[-1] ])
if self.is_training_mode:
self.src_sample_losses = []
self.dst_sample_losses = []
t = SampleProcessor.Types
face_type = t.FACE_TYPE_FULL if self.options['face_type'] == 'f' else t.FACE_TYPE_HALF
@ -464,48 +365,79 @@ class SAEModel(ModelBase):
random_ct_samples_path=training_data_dst_path if apply_random_ct else None,
debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, scale_range=np.array([-0.05, 0.05])+self.src_scale_mod / 100.0 ),
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t_mode_bgr), 'resolution':resolution, 'apply_ct': apply_random_ct},
{'types' : (t.IMG_TRANSFORMED, face_type, t_mode_bgr), 'resolution': resolution, 'apply_ct': apply_random_ct },
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution } ]
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t_mode_bgr), 'resolution':resolution, 'apply_ct': apply_random_ct} ] + \
[ {'types' : (t.IMG_TRANSFORMED, face_type, t_mode_bgr), 'resolution': resolution // (2**i), 'apply_ct': apply_random_ct } 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)]
),
SampleGeneratorFace(training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
sample_process_options=SampleProcessor.Options(random_flip=self.random_flip, ),
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t_mode_bgr), 'resolution':resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t_mode_bgr), 'resolution': resolution},
{'types' : (t.IMG_TRANSFORMED, face_type, t.MODE_M), 'resolution': resolution} ])
])
#override
def get_model_filename_list(self):
ar = self.model.get_model_filename_list ( exclude_for_pretrain=(self.pretrain and self.iter != 0) )
return ar
output_sample_types = [ {'types' : (t.IMG_WARPED_TRANSFORMED, face_type, t_mode_bgr), 'resolution':resolution} ] + \
[ {'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 onSave(self):
self.save_weights_safe( self.get_model_filename_list() )
opt_ar = [ [self.src_dst_opt, 'src_dst_opt'],
[self.src_dst_mask_opt, 'src_dst_mask_opt']
]
ar = []
if 'liae' in self.options['archi']:
ar += [[self.encoder, 'encoder.h5'],
[self.inter_B, 'inter_B.h5'],
[self.decoder, 'decoder.h5']
]
if not self.pretrain or self.iter == 0:
ar += [ [self.inter_AB, 'inter_AB.h5'],
]
if self.options['learn_mask']:
ar += [ [self.decoderm, 'decoderm.h5'] ]
elif 'df' in self.options['archi']:
if not self.pretrain or self.iter == 0:
ar += [ [self.encoder, 'encoder.h5'],
]
ar += [ [self.decoder_src, 'decoder_src.h5'],
[self.decoder_dst, 'decoder_dst.h5']
]
if self.options['learn_mask']:
ar += [ [self.decoder_srcm, 'decoder_srcm.h5'],
[self.decoder_dstm, 'decoder_dstm.h5'] ]
self.save_weights_safe(ar)
#override
def onTrainOneIter(self, generators_samples, generators_list):
warped_src, target_src, target_srcm = generators_samples[0]
warped_dst, target_dst, target_dstm = generators_samples[1]
src_samples = generators_samples[0]
dst_samples = generators_samples[1]
feed = [warped_src, warped_dst, target_src, target_srcm, target_dst, target_dstm]
feed = [src_samples[0], dst_samples[0] ] + \
src_samples[1:1+self.ms_count*2] + \
dst_samples[1:1+self.ms_count*2]
src_loss, dst_loss, = self.src_dst_train (feed)
if self.options['learn_mask']:
feed = [ warped_src, warped_dst, target_srcm, target_dstm ]
feed = [ src_samples[0], dst_samples[0] ] + \
src_samples[1+self.ms_count:1+self.ms_count*2] + \
dst_samples[1+self.ms_count:1+self.ms_count*2]
src_mask_loss, dst_mask_loss, = self.src_dst_mask_train (feed)
return ( ('src_loss', src_loss), ('dst_loss', dst_loss), )
return ( ('src_loss', src_loss), ('dst_loss', dst_loss) )
#override
def onGetPreview(self, sample):
test_S = sample[0][1][0:4] #first 4 samples
test_S_m = sample[0][2][0:4] #first 4 samples
test_S_m = sample[0][1+self.ms_count][0:4] #first 4 samples
test_D = sample[1][1][0:4]
test_D_m = sample[1][2][0:4]
test_D_m = sample[1][1+self.ms_count][0:4]
if self.options['learn_mask']:
S, D, SS, DD, DDM, SD, SDM = [ np.clip(x, 0.0, 1.0) for x in ([test_S,test_D] + self.AE_view ([test_S, test_D]) ) ]
@ -515,16 +447,15 @@ class SAEModel(ModelBase):
result = []
st = []
for i in range(len(test_S)):
for i in range(0, len(test_S)):
ar = S[i], SS[i], D[i], DD[i], SD[i]
st.append ( np.concatenate ( ar, axis=1) )
result += [ ('SAE', np.concatenate (st, axis=0 )), ]
if self.options['learn_mask']:
st_m = []
for i in range(len(test_S)):
for i in range(0, len(test_S)):
ar = S[i]*test_S_m[i], SS[i], D[i]*test_D_m[i], DD[i]*DDM[i], SD[i]*(DDM[i]*SDM[i])
st_m.append ( np.concatenate ( ar, axis=1) )
@ -532,26 +463,233 @@ class SAEModel(ModelBase):
return result
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 ) ])
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]
else:
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]
bgr, = self.AE_convert ([face[np.newaxis,...]])
return bgr[0]
#override
def get_ConverterConfig(self):
def get_converter(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
default_erode_mask_modifier = 0
default_blur_mask_modifier = 100 if (self.options['face_style_power'] or self.options['bg_style_power']) and \
self.options['face_type'] == 'f' else 0
face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
import converters
return self.predictor_func, (self.options['resolution'], self.options['resolution'], 3), converters.ConverterConfigMasked(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,
clip_hborder_mask_per=0.0625 if (self.options['face_type'] == 'f') else 0,
)
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['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)
Model = SAEModel
@staticmethod
def initialize_nn_functions():
exec (nnlib.import_all(), locals(), globals())
def NormPass(x):
return x
def Norm(norm=''):
if norm == 'bn':
return BatchNormalization(axis=-1)
else:
return NormPass
def Act(act='', lrelu_alpha=0.1):
if act == 'prelu':
return PReLU()
else:
return LeakyReLU(alpha=lrelu_alpha)
class ResidualBlock(object):
def __init__(self, filters, kernel_size=3, padding='zero', norm='', act='', **kwargs):
self.filters = filters
self.kernel_size = kernel_size
self.padding = padding
self.norm = norm
self.act = act
def __call__(self, inp):
x = inp
x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding)(x)
x = Act(self.act, lrelu_alpha=0.2)(x)
x = Norm(self.norm)(x)
x = Conv2D(self.filters, kernel_size=self.kernel_size, padding=self.padding)(x)
x = Add()([x, inp])
x = Act(self.act, lrelu_alpha=0.2)(x)
x = Norm(self.norm)(x)
return x
SAEModel.ResidualBlock = ResidualBlock
def downscale (dim, padding='zero', norm='', act='', **kwargs):
def func(x):
return Norm(norm)( Act(act) (Conv2D(dim, kernel_size=5, strides=2, padding=padding)(x)) )
return func
SAEModel.downscale = downscale
def upscale (dim, padding='zero', norm='', act='', **kwargs):
def func(x):
return SubpixelUpscaler()(Norm(norm)(Act(act)(Conv2D(dim * 4, kernel_size=3, strides=1, padding=padding)(x))))
return func
SAEModel.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
SAEModel.to_bgr = to_bgr
@staticmethod
def LIAEEncFlow(resolution, ch_dims, **kwargs):
exec (nnlib.import_all(), locals(), globals())
upscale = partial(SAEModel.upscale, **kwargs)
downscale = partial(SAEModel.downscale, **kwargs)
def func(input):
dims = K.int_shape(input)[-1]*ch_dims
x = input
x = downscale(dims)(x)
x = downscale(dims*2)(x)
x = downscale(dims*4)(x)
x = downscale(dims*8)(x)
x = Flatten()(x)
return x
return func
@staticmethod
def LIAEInterFlow(resolution, ae_dims=256, **kwargs):
exec (nnlib.import_all(), locals(), globals())
upscale = partial(SAEModel.upscale, **kwargs)
lowest_dense_res=resolution // 16
def func(input):
x = input[0]
x = Dense(ae_dims)(x)
x = Dense(lowest_dense_res * lowest_dense_res * ae_dims*2)(x)
x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims*2))(x)
x = upscale(ae_dims*2)(x)
return x
return func
@staticmethod
def LIAEDecFlow(output_nc,ch_dims, multiscale_count=1, add_residual_blocks=False, **kwargs):
exec (nnlib.import_all(), locals(), globals())
upscale = partial(SAEModel.upscale, **kwargs)
to_bgr = partial(SAEModel.to_bgr, **kwargs)
dims = output_nc * ch_dims
ResidualBlock = partial(SAEModel.ResidualBlock, **kwargs)
def func(input):
x = input[0]
outputs = []
x1 = upscale(dims*8)( x )
if add_residual_blocks:
x1 = ResidualBlock(dims*8)(x1)
x1 = ResidualBlock(dims*8)(x1)
if multiscale_count >= 3:
outputs += [ to_bgr(output_nc) ( x1 ) ]
x2 = upscale(dims*4)( x1 )
if add_residual_blocks:
x2 = ResidualBlock(dims*4)(x2)
x2 = ResidualBlock(dims*4)(x2)
if multiscale_count >= 2:
outputs += [ to_bgr(output_nc) ( x2 ) ]
x3 = upscale(dims*2)( x2 )
if add_residual_blocks:
x3 = ResidualBlock( dims*2)(x3)
x3 = ResidualBlock( dims*2)(x3)
outputs += [ to_bgr(output_nc) ( x3 ) ]
return outputs
return func
@staticmethod
def DFEncFlow(resolution, ae_dims, ch_dims, **kwargs):
exec (nnlib.import_all(), locals(), globals())
upscale = partial(SAEModel.upscale, **kwargs)
downscale = partial(SAEModel.downscale, **kwargs)#, kernel_regularizer=keras.regularizers.l2(0.0),
lowest_dense_res = resolution // 16
def func(input):
x = input
dims = K.int_shape(input)[-1]*ch_dims
x = downscale(dims)(x)
x = downscale(dims*2)(x)
x = downscale(dims*4)(x)
x = downscale(dims*8)(x)
x = Dense(ae_dims)(Flatten()(x))
x = Dense(lowest_dense_res * lowest_dense_res * ae_dims)(x)
x = Reshape((lowest_dense_res, lowest_dense_res, ae_dims))(x)
x = upscale(ae_dims)(x)
return x
return func
@staticmethod
def DFDecFlow(output_nc, ch_dims, multiscale_count=1, add_residual_blocks=False, **kwargs):
exec (nnlib.import_all(), locals(), globals())
upscale = partial(SAEModel.upscale, **kwargs)
to_bgr = partial(SAEModel.to_bgr, **kwargs)
dims = output_nc * ch_dims
ResidualBlock = partial(SAEModel.ResidualBlock, **kwargs)
def func(input):
x = input[0]
outputs = []
x1 = upscale(dims*8)( x )
if add_residual_blocks:
x1 = ResidualBlock( dims*8 )(x1)
x1 = ResidualBlock( dims*8 )(x1)
if multiscale_count >= 3:
outputs += [ to_bgr(output_nc) ( x1 ) ]
x2 = upscale(dims*4)( x1 )
if add_residual_blocks:
x2 = ResidualBlock( dims*4)(x2)
x2 = ResidualBlock( dims*4)(x2)
if multiscale_count >= 2:
outputs += [ to_bgr(output_nc) ( x2 ) ]
x3 = upscale(dims*2)( x2 )
if add_residual_blocks:
x3 = ResidualBlock( dims*2)(x3)
x3 = ResidualBlock( dims*2)(x3)
outputs += [ to_bgr(output_nc) ( x3 ) ]
return outputs
return func
Model = SAEModel

2
nnlib/device.py
View file

@ -299,7 +299,7 @@ def get_plaidML_devices():
except:
pass
return plaidML_devices
if not has_nvidia_device:
get_plaidML_devices()

187
nnlib/nnlib.py
View file

@ -53,17 +53,12 @@ Input = KL.Input
Dense = KL.Dense
Conv2D = nnlib.Conv2D
Conv2DTranspose = nnlib.Conv2DTranspose
EqualConv2D = nnlib.EqualConv2D
SeparableConv2D = KL.SeparableConv2D
MaxPooling2D = KL.MaxPooling2D
AveragePooling2D = KL.AveragePooling2D
GlobalAveragePooling2D = KL.GlobalAveragePooling2D
UpSampling2D = KL.UpSampling2D
BatchNormalization = KL.BatchNormalization
PixelNormalization = nnlib.PixelNormalization
LeakyReLU = KL.LeakyReLU
ELU = KL.ELU
ReLU = KL.ReLU
PReLU = KL.PReLU
tanh = KL.Activation('tanh')
@ -73,7 +68,6 @@ Softmax = KL.Softmax
Lambda = KL.Lambda
Add = KL.Add
Multiply = KL.Multiply
Concatenate = KL.Concatenate
@ -95,8 +89,6 @@ dssim = nnlib.dssim
PixelShuffler = nnlib.PixelShuffler
SubpixelUpscaler = nnlib.SubpixelUpscaler
Scale = nnlib.Scale
BlurPool = nnlib.BlurPool
SelfAttention = nnlib.SelfAttention
CAInitializerMP = nnlib.CAInitializerMP
@ -139,10 +131,6 @@ 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
@ -173,6 +161,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
return nnlib.code_import_keras
nnlib.backend = device_config.backend
if "tensorflow" in nnlib.backend:
nnlib._import_tf(device_config)
elif nnlib.backend == "plaidML":
@ -185,9 +174,6 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
import keras as keras_
nnlib.keras = keras_
if 'KERAS_BACKEND' in os.environ:
os.environ.pop('KERAS_BACKEND')
if nnlib.backend == "plaidML":
import plaidml
import plaidml.tile
@ -468,51 +454,6 @@ 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
"""
def __init__(self, filt_size=3, stride=2, **kwargs):
self.strides = (stride,stride)
self.filt_size = filt_size
self.padding = ( (int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)) ), (int(1.*(filt_size-1)/2), int(np.ceil(1.*(filt_size-1)/2)) ) )
if(self.filt_size==1):
self.a = np.array([1.,])
elif(self.filt_size==2):
self.a = np.array([1., 1.])
elif(self.filt_size==3):
self.a = np.array([1., 2., 1.])
elif(self.filt_size==4):
self.a = np.array([1., 3., 3., 1.])
elif(self.filt_size==5):
self.a = np.array([1., 4., 6., 4., 1.])
elif(self.filt_size==6):
self.a = np.array([1., 5., 10., 10., 5., 1.])
elif(self.filt_size==7):
self.a = np.array([1., 6., 15., 20., 15., 6., 1.])
super(BlurPool, self).__init__(**kwargs)
def compute_output_shape(self, input_shape):
height = input_shape[1] // self.strides[0]
width = input_shape[2] // self.strides[1]
channels = input_shape[3]
return (input_shape[0], height, width, channels)
def call(self, x):
k = self.a
k = k[:,None]*k[None,:]
k = k / np.sum(k)
k = np.tile (k[:,:,None,None], (1,1,K.int_shape(x)[-1],1) )
k = K.constant (k, dtype=K.floatx() )
x = K.spatial_2d_padding(x, padding=self.padding)
x = K.depthwise_conv2d(x, k, strides=self.strides, padding='valid')
return x
nnlib.BlurPool = BlurPool
class Scale(KL.Layer):
"""
GAN Custom Scal Layer
@ -544,43 +485,6 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
return dict(list(base_config.items()) + list(config.items()))
nnlib.Scale = Scale
class SelfAttention(KL.Layer):
def __init__(self, nc, squeeze_factor=8, **kwargs):
assert nc//squeeze_factor > 0, f"Input channels must be >= {squeeze_factor}, recieved nc={nc}"
self.nc = nc
self.squeeze_factor = squeeze_factor
super(SelfAttention, self).__init__(**kwargs)
def compute_output_shape(self, input_shape):
return (input_shape[0], input_shape[1], input_shape[2], self.nc)
def call(self, inp):
x = inp
shape_x = x.get_shape().as_list()
f = Conv2D(self.nc//self.squeeze_factor, 1, kernel_regularizer=keras.regularizers.l2(1e-4))(x)
g = Conv2D(self.nc//self.squeeze_factor, 1, kernel_regularizer=keras.regularizers.l2(1e-4))(x)
h = Conv2D(self.nc, 1, kernel_regularizer=keras.regularizers.l2(1e-4))(x)
shape_f = f.get_shape().as_list()
shape_g = g.get_shape().as_list()
shape_h = h.get_shape().as_list()
flat_f = Reshape( (-1, shape_f[-1]) )(f)
flat_g = Reshape( (-1, shape_g[-1]) )(g)
flat_h = Reshape( (-1, shape_h[-1]) )(h)
s = Lambda(lambda x: K.batch_dot(x[0], keras.layers.Permute((2,1))(x[1]) ))([flat_g, flat_f])
beta = keras.layers.Softmax(axis=-1)(s)
o = Lambda(lambda x: K.batch_dot(x[0], x[1]))([beta, flat_h])
o = Reshape(shape_x[1:])(o)
o = Scale()(o)
out = Add()([o, inp])
return out
nnlib.SelfAttention = SelfAttention
class Adam(keras.optimizers.Optimizer):
"""Adam optimizer.
@ -688,9 +592,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
def CAInitializerMP( conv_weights_list ):
#Convolution Aware Initialization https://arxiv.org/abs/1702.06295
data = [ (i, K.int_shape(conv_weights)) for i, conv_weights in enumerate(conv_weights_list) ]
data = sorted(data, key=lambda data: np.prod(data[1]) )
result = CAInitializerMPSubprocessor (data, K.floatx(), K.image_data_format() ).run()
result = CAInitializerMPSubprocessor ( [ (i, K.int_shape(conv_weights)) for i, conv_weights in enumerate(conv_weights_list) ], K.floatx(), K.image_data_format() ).run()
for idx, weights in result:
K.set_value ( conv_weights_list[idx], weights )
nnlib.CAInitializerMP = CAInitializerMP
@ -794,7 +696,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
x = ReflectionPadding2D( self.pad ) (x)
return self.func(x)
nnlib.Conv2D = Conv2D
class Conv2DTranspose():
def __init__ (self, *args, **kwargs):
self.reflect_pad = False
@ -814,89 +716,7 @@ NLayerDiscriminator = nnlib.NLayerDiscriminator
x = ReflectionPadding2D( self.pad ) (x)
return self.func(x)
nnlib.Conv2DTranspose = Conv2DTranspose
class EqualConv2D(KL.Conv2D):
def __init__(self, filters,
kernel_size,
strides=(1, 1),
padding='valid',
data_format=None,
dilation_rate=(1, 1),
activation=None,
use_bias=True,
gain=np.sqrt(2),
**kwargs):
super().__init__(
filters=filters,
kernel_size=kernel_size,
strides=strides,
padding=padding,
data_format=data_format,
dilation_rate=dilation_rate,
activation=activation,
use_bias=use_bias,
kernel_initializer=keras.initializers.RandomNormal(mean=0.0, stddev=1.0),
bias_initializer='zeros',
kernel_regularizer=None,
bias_regularizer=None,
activity_regularizer=None,
kernel_constraint=None,
bias_constraint=None,
**kwargs)
self.gain = gain
def build(self, input_shape):
super().build(input_shape)
self.wscale = self.gain / np.sqrt( np.prod( K.int_shape(self.kernel)[:-1]) )
self.wscale_t = K.constant (self.wscale, dtype=K.floatx() )
def call(self, inputs):
k = self.kernel * self.wscale_t
outputs = K.conv2d(
inputs,
k,
strides=self.strides,
padding=self.padding,
data_format=self.data_format,
dilation_rate=self.dilation_rate)
if self.use_bias:
outputs = K.bias_add(
outputs,
self.bias,
data_format=self.data_format)
if self.activation is not None:
return self.activation(outputs)
return outputs
nnlib.EqualConv2D = EqualConv2D
class PixelNormalization(KL.Layer):
# initialize the layer
def __init__(self, **kwargs):
super(PixelNormalization, self).__init__(**kwargs)
# perform the operation
def call(self, inputs):
# calculate square pixel values
values = inputs**2.0
# calculate the mean pixel values
mean_values = K.mean(values, axis=-1, keepdims=True)
# ensure the mean is not zero
mean_values += 1.0e-8
# calculate the sqrt of the mean squared value (L2 norm)
l2 = K.sqrt(mean_values)
# normalize values by the l2 norm
normalized = inputs / l2
return normalized
# define the output shape of the layer
def compute_output_shape(self, input_shape):
return input_shape
nnlib.PixelNormalization = PixelNormalization
@staticmethod
def import_keras_contrib(device_config):
if nnlib.keras_contrib is not None:
@ -1180,6 +1000,7 @@ class CAInitializerMPSubprocessor(Subprocessor):
#override
def __init__(self, idx_shapes_list, floatx, data_format ):
self.idx_shapes_list = idx_shapes_list
self.floatx = floatx
self.data_format = data_format

4
requirements-colab.txt
View file

@ -1,7 +1,7 @@
numpy==1.17.0
numpy==1.16.3
h5py==2.9.0
Keras==2.2.4
opencv-python==4.1.0.25
opencv-python==4.0.0.21
tensorflow-gpu==1.13.1
plaidml-keras==0.5.0
scikit-image

4
requirements-cpu.txt
View file

@ -1,7 +1,7 @@
numpy==1.17.0
numpy==1.16.3
h5py==2.9.0
Keras==2.2.4
opencv-python==4.1.0.25
opencv-python==4.0.0.21
tensorflow==1.12.0
scikit-image
tqdm

6
requirements-cuda.txt
View file

@ -1,9 +1,9 @@
numpy==1.17.0
numpy==1.16.3
h5py==2.9.0
Keras==2.2.4
opencv-python==4.1.0.25
opencv-python==4.0.0.21
tensorflow-gpu==1.12.0
plaidml==0.6.0
plaidml==0.6.0rc1
plaidml-keras==0.5.0
scikit-image
tqdm

6
requirements-opencl.txt
View file

@ -1,9 +1,9 @@
numpy==1.17.0
numpy==1.16.3
h5py==2.9.0
Keras==2.2.4
opencv-python==4.1.0.25
opencv-python==4.0.0.21
tensorflow==1.12.0
plaidml==0.6.0
plaidml==0.6.0rc1
plaidml-keras==0.5.0
scikit-image
tqdm

6
samplelib/Sample.py
View file

@ -22,7 +22,7 @@ class SampleType(IntEnum):
QTY = 5
class Sample(object):
def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch_yaw_roll=None, eyebrows_expand_mod=None, source_filename=None, mirror=None, close_target_list=None, fanseg_mask_exist=False):
def __init__(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch_yaw_roll=None, source_filename=None, mirror=None, close_target_list=None, fanseg_mask_exist=False):
self.sample_type = sample_type if sample_type is not None else SampleType.IMAGE
self.filename = filename
self.face_type = face_type
@ -30,13 +30,12 @@ class Sample(object):
self.landmarks = np.array(landmarks) if landmarks is not None else None
self.ie_polys = ie_polys
self.pitch_yaw_roll = pitch_yaw_roll
self.eyebrows_expand_mod = eyebrows_expand_mod
self.source_filename = source_filename
self.mirror = mirror
self.close_target_list = close_target_list
self.fanseg_mask_exist = fanseg_mask_exist
def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch_yaw_roll=None, eyebrows_expand_mod=None, source_filename=None, mirror=None, close_target_list=None, fanseg_mask=None, fanseg_mask_exist=None):
def copy_and_set(self, sample_type=None, filename=None, face_type=None, shape=None, landmarks=None, ie_polys=None, pitch_yaw_roll=None, source_filename=None, mirror=None, close_target_list=None, fanseg_mask=None, fanseg_mask_exist=None):
return Sample(
sample_type=sample_type if sample_type is not None else self.sample_type,
filename=filename if filename is not None else self.filename,
@ -45,7 +44,6 @@ class Sample(object):
landmarks=landmarks if landmarks is not None else self.landmarks.copy(),
ie_polys=ie_polys if ie_polys is not None else self.ie_polys,
pitch_yaw_roll=pitch_yaw_roll if pitch_yaw_roll is not None else self.pitch_yaw_roll,
eyebrows_expand_mod=eyebrows_expand_mod if eyebrows_expand_mod is not None else self.eyebrows_expand_mod,
source_filename=source_filename if source_filename is not None else self.source_filename,
mirror=mirror if mirror is not None else self.mirror,
close_target_list=close_target_list if close_target_list is not None else self.close_target_list,

13
samplelib/SampleGeneratorBase.py
View file

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

9
samplelib/SampleGeneratorFace.py
View file

@ -39,7 +39,6 @@ class SampleGeneratorFace(SampleGeneratorBase):
samples = SampleLoader.load (self.sample_type, self.samples_path, sort_by_yaw_target_samples_path)
ct_samples = SampleLoader.load (SampleType.FACE, random_ct_samples_path) if random_ct_samples_path is not None else None
self.random_ct_sample_chance = 100
if self.debug:
self.generators_count = 1
@ -114,12 +113,8 @@ class SampleGeneratorFace(SampleGeneratorBase):
if sample is not None:
try:
ct_sample=None
if ct_samples is not None:
if np.random.randint(100) < self.random_ct_sample_chance:
ct_sample=ct_samples[np.random.randint(ct_samples_len)]
x = SampleProcessor.process (sample, self.sample_process_options, self.output_sample_types, self.debug, ct_sample=ct_sample)
x = SampleProcessor.process (sample, self.sample_process_options, self.output_sample_types, self.debug,
ct_sample=ct_samples[np.random.randint(ct_samples_len)] if ct_samples is not None else None )
except:
raise Exception ("Exception occured in sample %s. Error: %s" % (sample.filename, traceback.format_exc() ) )

8
samplelib/SampleGeneratorFaceTemporal.py
View file

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

5
samplelib/SampleGeneratorImageTemporal.py
View file

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

6
samplelib/SampleLoader.py
View file

@ -71,10 +71,9 @@ class SampleLoader:
landmarks = dflimg.get_landmarks()
pitch_yaw_roll = dflimg.get_pitch_yaw_roll()
eyebrows_expand_mod = dflimg.get_eyebrows_expand_mod()
if pitch_yaw_roll is None:
pitch_yaw_roll = LandmarksProcessor.estimate_pitch_yaw_roll(landmarks)
pitch_yaw_roll = LandmarksProcessor.estimate_pitch_yaw_roll(landmarks)
sample_list.append( s.copy_and_set(sample_type=SampleType.FACE,
face_type=FaceType.fromString (dflimg.get_face_type()),
@ -82,7 +81,6 @@ class SampleLoader:
landmarks=landmarks,
ie_polys=dflimg.get_ie_polys(),
pitch_yaw_roll=pitch_yaw_roll,
eyebrows_expand_mod=eyebrows_expand_mod,
source_filename=dflimg.get_source_filename(),
fanseg_mask_exist=dflimg.get_fanseg_mask() is not None, ) )
except:

120
samplelib/SampleProcessor.py
View file

@ -36,9 +36,10 @@ opts:
'MODE_BGR_SHUFFLE' #BGR shuffle
'resolution' : N
'motion_blur' : (chance_int, range) - chance 0..100 to apply to face (not mask), and max_size of motion blur
'motion_blur' : (chance_int, range) - chance 0..100 to apply to face (not mask), and range [1..3] where 3 is highest power of motion blur
'apply_ct' : bool
'normalize_tanh' : bool
"""
@ -61,8 +62,6 @@ class SampleProcessor(object):
FACE_TYPE_FULL = 11
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
@ -74,20 +73,15 @@ 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] ):
def __init__(self, random_flip = True, normalize_tanh = False, 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.normalize_tanh = normalize_tanh
self.rotation_range = rotation_range
self.scale_range = scale_range
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
@ -108,7 +102,10 @@ 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_AVATAR : FaceType.AVATAR }
outputs = []
for opts in output_sample_types:
@ -116,13 +113,11 @@ class SampleProcessor(object):
resolution = opts.get('resolution', 0)
types = opts.get('types', [] )
border_replicate = opts.get('border_replicate', True)
random_sub_res = opts.get('random_sub_res', 0)
normalize_std_dev = opts.get('normalize_std_dev', False)
normalize_vgg = opts.get('normalize_vgg', False)
motion_blur = opts.get('motion_blur', None)
apply_ct = opts.get('apply_ct', False)
normalize_tanh = opts.get('normalize_tanh', False)
img_type = SPTF.NONE
target_face_type = SPTF.NONE
@ -163,63 +158,48 @@ class SampleProcessor(object):
if mode_type == SPTF.NONE:
raise ValueError ('expected MODE_ type')
def do_transform(img, mask):
img = cached_images.get(img_type, None)
if img is None:
img = sample_bgr
mask = None
cur_sample = sample
if is_face_sample:
if motion_blur is not None:
chance, mb_range = motion_blur
chance = np.clip(chance, 0, 100)
if np.random.randint(100) < chance:
mb_range = [3,5,7,9][ : np.clip(mb_range, 0, 3)+1 ]
dim = mb_range[ np.random.randint(len(mb_range) ) ]
img = imagelib.LinearMotionBlur (img, dim, np.random.randint(180) )
mask = cur_sample.load_fanseg_mask() #using fanseg_mask if exist
if mask is None:
mask = LandmarksProcessor.get_image_hull_mask (img.shape, cur_sample.landmarks)
if cur_sample.ie_polys is not None:
cur_sample.ie_polys.overlay_mask(mask)
warp = (img_type==SPTF.IMG_WARPED or img_type==SPTF.IMG_WARPED_TRANSFORMED)
transform = (img_type==SPTF.IMG_WARPED_TRANSFORMED or img_type==SPTF.IMG_TRANSFORMED)
flip = img_type != SPTF.IMG_WARPED
img = imagelib.warp_by_params (params, img, warp, transform, flip, border_replicate)
img = imagelib.warp_by_params (params, img, warp, transform, flip, True)
if mask is not None:
mask = imagelib.warp_by_params (params, mask, warp, transform, flip, False)
if len(mask.shape) == 2:
mask = mask[...,np.newaxis]
mask = imagelib.warp_by_params (params, mask, warp, transform, flip, False)[...,np.newaxis]
img = np.concatenate( (img, mask ), -1 )
return img
img = sample_bgr
### Prepare a mask
mask = None
if is_face_sample:
mask = sample.load_fanseg_mask() #using fanseg_mask if exist
if mask is None:
if sample.eyebrows_expand_mod is not None:
mask = LandmarksProcessor.get_image_hull_mask (img.shape, sample.landmarks, eyebrows_expand_mod=sample.eyebrows_expand_mod )
else:
mask = LandmarksProcessor.get_image_hull_mask (img.shape, sample.landmarks)
if sample.ie_polys is not None:
sample.ie_polys.overlay_mask(mask)
##################
if motion_blur is not None:
chance, mb_max_size = motion_blur
chance = np.clip(chance, 0, 100)
if np.random.randint(100) < chance:
img = imagelib.LinearMotionBlur (img, np.random.randint( mb_max_size )+1, np.random.randint(360) )
cached_images[img_type] = img
if is_face_sample and target_face_type != SPTF.NONE:
target_ft = SampleProcessor.SPTF_FACETYPE_TO_FACETYPE[target_face_type]
if target_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, target_ft) )
if sample.face_type == FaceType.MARK_ONLY:
#first warp to target facetype
img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, sample.shape[0], target_ft), (sample.shape[0],sample.shape[0]), flags=cv2.INTER_CUBIC )
mask = cv2.warpAffine( mask, LandmarksProcessor.get_transform_mat (sample.landmarks, sample.shape[0], target_ft), (sample.shape[0],sample.shape[0]), flags=cv2.INTER_CUBIC )
#then apply transforms
img = do_transform (img, mask)
img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC )
else:
img = do_transform (img, mask)
img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, resolution, target_ft), (resolution,resolution), borderMode=(cv2.BORDER_REPLICATE if border_replicate else cv2.BORDER_CONSTANT), flags=cv2.INTER_CUBIC )
ft = SPTF_FACETYPE_TO_FACETYPE[target_face_type]
if ft > sample.face_type:
raise Exception ('sample %s type %s does not match model requirement %s. Consider extract necessary type of faces.' % (sample.filename, sample.face_type, ft) )
img = cv2.warpAffine( img, LandmarksProcessor.get_transform_mat (sample.landmarks, resolution, ft), (resolution,resolution), flags=cv2.INTER_CUBIC )
else:
img = do_transform (img, mask)
img = cv2.resize( img, (resolution,resolution), cv2.INTER_CUBIC )
if random_sub_res != 0:
@ -233,14 +213,19 @@ class SampleProcessor(object):
img_bgr = img[...,0:3]
img_mask = img[...,3:4]
if apply_ct and ct_sample is not None:
if apply_ct:
if ct_sample_bgr is None:
ct_sample_bgr = ct_sample.load_bgr()
ct_sample_mask = LandmarksProcessor.get_image_hull_mask (ct_sample_bgr.shape, ct_sample.landmarks)
ct_sample_bgr_resized = cv2.resize( ct_sample_bgr, (resolution,resolution), cv2.INTER_LINEAR )
ct_sample_mask_resized = cv2.resize( ct_sample_mask, (resolution,resolution), cv2.INTER_LINEAR )[...,np.newaxis]
img_bgr = imagelib.linear_color_transfer (img_bgr, ct_sample_bgr_resized)
img_bgr = np.clip( img_bgr, 0.0, 1.0)
img_bgr = imagelib.reinhard_color_transfer ( np.clip( (img_bgr*255) .astype(np.uint8), 0, 255),
np.clip( (ct_sample_bgr_resized*255).astype(np.uint8), 0, 255),
source_mask=img_mask, target_mask=ct_sample_mask_resized)
img_bgr = np.clip( img_bgr.astype(np.float32) / 255.0, 0.0, 1.0)
if normalize_std_dev:
img_bgr = (img_bgr - img_bgr.mean( (0,1)) ) / img_bgr.std( (0,1) )
@ -254,7 +239,8 @@ class SampleProcessor(object):
img = img_bgr
elif mode_type == SPTF.MODE_BGR_SHUFFLE:
rnd_state = np.random.RandomState (sample_rnd_seed)
img = np.take (img_bgr, rnd_state.permutation(img_bgr.shape[-1]), axis=-1)
img_bgr = np.take (img_bgr, rnd_state.permutation(img_bgr.shape[-1]), axis=-1)
img = np.concatenate ( (img_bgr,img_mask) , -1 )
elif mode_type == SPTF.MODE_G:
img = np.concatenate ( (np.expand_dims(cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY),-1),img_mask) , -1 )
elif mode_type == SPTF.MODE_GGG:
@ -263,7 +249,7 @@ class SampleProcessor(object):
img = img_mask
if not debug:
if normalize_tanh:
if sample_process_options.normalize_tanh:
img = np.clip (img * 2.0 - 1.0, -1.0, 1.0)
else:
img = np.clip (img, 0.0, 1.0)

20
utils/DFLJPG.py
View file

@ -168,8 +168,7 @@ class DFLJPG(object):
source_landmarks=None,
image_to_face_mat=None,
fanseg_mask=None,
pitch_yaw_roll=None,
eyebrows_expand_mod=None,
pitch_yaw_roll=None,
**kwargs
):
@ -194,8 +193,7 @@ class DFLJPG(object):
'source_landmarks': source_landmarks,
'image_to_face_mat': image_to_face_mat,
'fanseg_mask' : fanseg_mask,
'pitch_yaw_roll' : pitch_yaw_roll,
'eyebrows_expand_mod' : eyebrows_expand_mod
'pitch_yaw_roll' : pitch_yaw_roll
})
try:
@ -213,7 +211,6 @@ class DFLJPG(object):
image_to_face_mat=None,
fanseg_mask=None,
pitch_yaw_roll=None,
eyebrows_expand_mod=None,
**kwargs
):
if face_type is None: face_type = self.get_face_type()
@ -225,8 +222,6 @@ class DFLJPG(object):
if image_to_face_mat is None: image_to_face_mat = self.get_image_to_face_mat()
if fanseg_mask is None: fanseg_mask = self.get_fanseg_mask()
if pitch_yaw_roll is None: pitch_yaw_roll = self.get_pitch_yaw_roll()
if eyebrows_expand_mod is None: eyebrows_expand_mod = self.get_eyebrows_expand_mod()
DFLJPG.embed_data (filename, face_type=face_type,
landmarks=landmarks,
ie_polys=ie_polys,
@ -235,12 +230,7 @@ class DFLJPG(object):
source_landmarks=source_landmarks,
image_to_face_mat=image_to_face_mat,
fanseg_mask=fanseg_mask,
pitch_yaw_roll=pitch_yaw_roll,
eyebrows_expand_mod=eyebrows_expand_mod)
def remove_ie_polys(self):
self.dfl_dict['ie_polys'] = None
pitch_yaw_roll=pitch_yaw_roll)
def remove_fanseg_mask(self):
self.dfl_dict['fanseg_mask'] = None
@ -310,6 +300,4 @@ class DFLJPG(object):
return None
def get_pitch_yaw_roll(self):
return self.dfl_dict.get ('pitch_yaw_roll', None)
def get_eyebrows_expand_mod(self):
return self.dfl_dict.get ('eyebrows_expand_mod', None)

16
utils/DFLPNG.py
View file

@ -285,7 +285,6 @@ class DFLPNG(object):
image_to_face_mat=None,
fanseg_mask=None,
pitch_yaw_roll=None,
eyebrows_expand_mod=None,
**kwargs
):
@ -310,8 +309,7 @@ class DFLPNG(object):
'source_landmarks': source_landmarks,
'image_to_face_mat':image_to_face_mat,
'fanseg_mask' : fanseg_mask,
'pitch_yaw_roll' : pitch_yaw_roll,
'eyebrows_expand_mod' : eyebrows_expand_mod,
'pitch_yaw_roll' : pitch_yaw_roll
})
try:
@ -329,7 +327,6 @@ class DFLPNG(object):
image_to_face_mat=None,
fanseg_mask=None,
pitch_yaw_roll=None,
eyebrows_expand_mod=None,
**kwargs
):
if face_type is None: face_type = self.get_face_type()
@ -341,8 +338,6 @@ class DFLPNG(object):
if image_to_face_mat is None: image_to_face_mat = self.get_image_to_face_mat()
if fanseg_mask is None: fanseg_mask = self.get_fanseg_mask()
if pitch_yaw_roll is None: pitch_yaw_roll = self.get_pitch_yaw_roll()
if eyebrows_expand_mod is None: eyebrows_expand_mod = self.get_eyebrows_expand_mod()
DFLPNG.embed_data (filename, face_type=face_type,
landmarks=landmarks,
ie_polys=ie_polys,
@ -351,11 +346,7 @@ class DFLPNG(object):
source_landmarks=source_landmarks,
image_to_face_mat=image_to_face_mat,
fanseg_mask=fanseg_mask,
pitch_yaw_roll=pitch_yaw_roll,
eyebrows_expand_mod=eyebrows_expand_mod)
def remove_ie_polys(self):
self.dfl_dict['ie_polys'] = None
pitch_yaw_roll=pitch_yaw_roll)
def remove_fanseg_mask(self):
self.dfl_dict['fanseg_mask'] = None
@ -415,8 +406,5 @@ class DFLPNG(object):
return None
def get_pitch_yaw_roll(self):
return self.dfl_dict.get ('pitch_yaw_roll', None)
def get_eyebrows_expand_mod(self):
return self.dfl_dict.get ('eyebrows_expand_mod', None)
def __str__(self):
return "<PNG length={length} chunks={}>".format(len(self.chunks), **self.__dict__)

36
utils/Path_utils.py
View file

@ -11,7 +11,7 @@ def get_image_paths(dir_path, image_extensions=image_extensions):
for x in list(scandir(str(dir_path))):
if any([x.name.lower().endswith(ext) for ext in image_extensions]):
result.append(x.path)
return sorted(result)
return result
def get_image_unique_filestem_paths(dir_path, verbose_print_func=None):
result = get_image_paths(dir_path)
@ -26,24 +26,8 @@ def get_image_unique_filestem_paths(dir_path, verbose_print_func=None):
continue
result_dup.add(f_stem)
return sorted(result)
def get_file_paths(dir_path):
dir_path = Path (dir_path)
return result
if dir_path.exists():
return sorted([ x.path for x in list(scandir(str(dir_path))) if x.is_file() ])
else:
return []
def get_all_dir_names (dir_path):
dir_path = Path (dir_path)
if dir_path.exists():
return sorted([ x.name for x in list(scandir(str(dir_path))) if x.is_dir() ])
else:
return []
def get_all_dir_names_startswith (dir_path, startswith):
dir_path = Path (dir_path)
startswith = startswith.lower()
@ -53,14 +37,14 @@ def get_all_dir_names_startswith (dir_path, startswith):
for x in list(scandir(str(dir_path))):
if x.name.lower().startswith(startswith):
result.append ( x.name[len(startswith):] )
return sorted(result)
return result
def get_first_file_by_stem (dir_path, stem, exts=None):
dir_path = Path (dir_path)
stem = stem.lower()
if dir_path.exists():
for x in sorted(list(scandir(str(dir_path))), key=lambda x: x.name):
for x in list(scandir(str(dir_path))):
if not x.is_file():
continue
xp = Path(x.path)
@ -68,15 +52,3 @@ def get_first_file_by_stem (dir_path, stem, exts=None):
return xp
return None
def move_all_files (src_dir_path, dst_dir_path):
paths = get_file_paths(src_dir_path)
for p in paths:
p = Path(p)
p.rename ( Path(dst_dir_path) / p.name )
def delete_all_files (dir_path):
paths = get_file_paths(dir_path)
for p in paths:
p = Path(p)
p.unlink()

5
utils/iter_utils.py
View file

@ -48,11 +48,6 @@ class SubprocessGenerator(object):
def __iter__(self):
return self
def __getstate__(self):
self_dict = self.__dict__.copy()
del self_dict['p']
return self_dict
def __next__(self):
if self.p == None:
user_param = self.user_param

11
utils/os_utils.py
View file

@ -23,14 +23,3 @@ 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)