mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-14 02:37:00 -07:00
manual extractor: increased FPS,
sort by final : now you can specify target number of images, converter: fix seamless mask and exception, huge refactoring
This commit is contained in:
iperov
parent
7db469a1da
commit
438213e97c
30 changed files with 1834 additions and 1718 deletions
|
|
@ -1,46 +0,0 @@
|
|||
import copy
|
||||
'''
|
||||
You can implement your own Converter, check example ConverterMasked.py
|
||||
'''
|
||||
|
||||
class ConverterBase(object):
|
||||
MODE_FACE = 0
|
||||
MODE_IMAGE = 1
|
||||
MODE_IMAGE_WITH_LANDMARKS = 2
|
||||
|
||||
#overridable
|
||||
def __init__(self, predictor):
|
||||
self.predictor = predictor
|
||||
|
||||
#overridable
|
||||
def get_mode(self):
|
||||
#MODE_FACE calls convert_face
|
||||
#MODE_IMAGE calls convert_image without landmarks
|
||||
#MODE_IMAGE_WITH_LANDMARKS calls convert_image with landmarks
|
||||
return ConverterBase.MODE_FACE
|
||||
|
||||
#overridable
|
||||
def convert_face (self, img_bgr, img_face_landmarks, debug):
|
||||
#return float32 image
|
||||
#if debug , return tuple ( images of any size and channels, ...)
|
||||
return image
|
||||
|
||||
#overridable
|
||||
def 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):
|
||||
result = self.copy()
|
||||
result.predictor = predictor
|
||||
return result
|
||||
|
|
@ -1,46 +0,0 @@
|
|||
from models import ConverterBase
|
||||
from facelib import LandmarksProcessor
|
||||
from facelib import FaceType
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from utils import image_utils
|
||||
|
||||
'''
|
||||
predictor:
|
||||
input: [predictor_input_size, predictor_input_size, BGR]
|
||||
output: [predictor_input_size, predictor_input_size, BGR]
|
||||
'''
|
||||
|
||||
class ConverterImage(ConverterBase):
|
||||
|
||||
#override
|
||||
def __init__(self, predictor,
|
||||
predictor_input_size=0,
|
||||
output_size=0,
|
||||
**in_options):
|
||||
|
||||
super().__init__(predictor)
|
||||
|
||||
self.predictor_input_size = predictor_input_size
|
||||
self.output_size = output_size
|
||||
|
||||
#override
|
||||
def get_mode(self):
|
||||
return ConverterBase.MODE_IMAGE
|
||||
|
||||
#override
|
||||
def dummy_predict(self):
|
||||
self.predictor ( np.zeros ( (self.predictor_input_size, self.predictor_input_size,3), dtype=np.float32) )
|
||||
|
||||
#override
|
||||
def 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 )
|
||||
predicted_bgr = self.predictor ( predictor_input_bgr )
|
||||
|
||||
output = cv2.resize ( predicted_bgr, (self.output_size, self.output_size), cv2.INTER_LANCZOS4 )
|
||||
if debug:
|
||||
return (predictor_input_bgr,output,)
|
||||
return output
|
||||
|
|
@ -1,313 +0,0 @@
|
|||
from models import ConverterBase
|
||||
from facelib import LandmarksProcessor
|
||||
from facelib import FaceType
|
||||
import cv2
|
||||
import numpy as np
|
||||
from utils import image_utils
|
||||
from utils.console_utils import *
|
||||
|
||||
'''
|
||||
default_mode = {1:'overlay',
|
||||
2:'hist-match',
|
||||
3:'hist-match-bw',
|
||||
4:'seamless',
|
||||
5:'seamless-hist-match',
|
||||
6:'raw'}
|
||||
'''
|
||||
class ConverterMasked(ConverterBase):
|
||||
|
||||
#override
|
||||
def __init__(self, predictor,
|
||||
predictor_input_size=0,
|
||||
output_size=0,
|
||||
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,
|
||||
**in_options):
|
||||
|
||||
super().__init__(predictor)
|
||||
self.predictor_input_size = predictor_input_size
|
||||
self.output_size = output_size
|
||||
self.face_type = face_type
|
||||
self.clip_hborder_mask_per = clip_hborder_mask_per
|
||||
|
||||
mode = input_int ("Choose mode: (1) overlay, (2) hist match, (3) hist match bw, (4) seamless, (5) seamless hist match, (6) raw. Default - %d : " % (default_mode) , default_mode)
|
||||
|
||||
mode_dict = {1:'overlay',
|
||||
2:'hist-match',
|
||||
3:'hist-match-bw',
|
||||
4:'seamless',
|
||||
5:'seamless-hist-match',
|
||||
6:'raw'}
|
||||
|
||||
self.mode = mode_dict.get (mode, mode_dict[default_mode] )
|
||||
|
||||
if self.mode == 'raw':
|
||||
mode = 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 == 'hist-match' or self.mode == 'hist-match-bw':
|
||||
self.masked_hist_match = 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 ( input_int("Hist match threshold [0..255] (skip:255) : ", 255), 0, 255)
|
||||
|
||||
self.use_predicted_mask = input_bool("Use predicted mask? (y/n skip:y) : ", True)
|
||||
|
||||
if self.mode != 'raw':
|
||||
self.erode_mask_modifier = base_erode_mask_modifier + np.clip ( 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 ( input_int ("Choose blur mask modifier [-200..200] (skip:%d) : " % (default_blur_mask_modifier), default_blur_mask_modifier), -200, 200)
|
||||
|
||||
self.seamless_erode_mask_modifier = 0
|
||||
if self.mode == 'seamless' or self.mode == 'seamless-hist-match':
|
||||
self.seamless_erode_mask_modifier = np.clip ( input_int ("Choose seamless erode mask modifier [-100..100] (skip:0) : ", 0), -100, 100)
|
||||
|
||||
self.output_face_scale = np.clip ( 1.0 + input_int ("Choose output face scale modifier [-50..50] (skip:0) : ", 0)*0.01, 0.5, 1.5)
|
||||
self.color_transfer_mode = input_str ("Apply color transfer to predicted face? Choose mode ( rct/lct skip:None ) : ", None, ['rct','lct'])
|
||||
|
||||
if self.mode != 'raw':
|
||||
self.final_image_color_degrade_power = np.clip ( input_int ("Degrade color power of final image [0..100] (skip:0) : ", 0), 0, 100)
|
||||
self.alpha = input_bool("Export png with alpha channel? (y/n skip:n) : ", False)
|
||||
|
||||
print ("")
|
||||
|
||||
#override
|
||||
def get_mode(self):
|
||||
return ConverterBase.MODE_FACE
|
||||
|
||||
#override
|
||||
def dummy_predict(self):
|
||||
self.predictor ( np.zeros ( (self.predictor_input_size,self.predictor_input_size,4), dtype=np.float32 ) )
|
||||
|
||||
#override
|
||||
def convert_face (self, img_bgr, img_face_landmarks, debug):
|
||||
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)
|
||||
|
||||
face_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, self.output_size, face_type=self.face_type)
|
||||
face_output_mat = LandmarksProcessor.get_transform_mat (img_face_landmarks, self.output_size, face_type=self.face_type, scale=self.output_face_scale)
|
||||
|
||||
dst_face_bgr = cv2.warpAffine( img_bgr , face_mat, (self.output_size, self.output_size), flags=cv2.INTER_LANCZOS4 )
|
||||
dst_face_mask_a_0 = cv2.warpAffine( img_face_mask_a, face_mat, (self.output_size, self.output_size), flags=cv2.INTER_LANCZOS4 )
|
||||
|
||||
predictor_input_bgr = cv2.resize (dst_face_bgr, (self.predictor_input_size,self.predictor_input_size))
|
||||
predictor_input_mask_a_0 = cv2.resize (dst_face_mask_a_0, (self.predictor_input_size,self.predictor_input_size))
|
||||
predictor_input_mask_a = np.expand_dims (predictor_input_mask_a_0, -1)
|
||||
|
||||
predicted_bgra = self.predictor ( np.concatenate( (predictor_input_bgr, predictor_input_mask_a), -1) )
|
||||
|
||||
prd_face_bgr = np.clip (predicted_bgra[:,:,0:3], 0, 1.0 )
|
||||
prd_face_mask_a_0 = np.clip (predicted_bgra[:,:,3], 0.0, 1.0)
|
||||
|
||||
if not self.use_predicted_mask:
|
||||
prd_face_mask_a_0 = predictor_input_mask_a_0
|
||||
|
||||
prd_face_mask_a_0[ prd_face_mask_a_0 < 0.001 ] = 0.0
|
||||
|
||||
prd_face_mask_a = np.expand_dims (prd_face_mask_a_0, axis=-1)
|
||||
prd_face_mask_aaa = np.repeat (prd_face_mask_a, (3,), axis=-1)
|
||||
|
||||
img_prd_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_prd_face_mask_aaa = np.clip (img_prd_face_mask_aaa, 0.0, 1.0)
|
||||
|
||||
img_face_mask_aaa = img_prd_face_mask_aaa
|
||||
|
||||
if debug:
|
||||
debugs += [img_face_mask_aaa.copy()]
|
||||
|
||||
img_face_mask_aaa [ img_face_mask_aaa <= 0.1 ] = 0.0
|
||||
|
||||
img_face_mask_flatten_aaa = img_face_mask_aaa.copy()
|
||||
img_face_mask_flatten_aaa[img_face_mask_flatten_aaa > 0.9] = 1.0
|
||||
|
||||
maxregion = np.argwhere(img_face_mask_flatten_aaa==1.0)
|
||||
|
||||
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:
|
||||
if maxregion.size != 0:
|
||||
miny,minx = maxregion.min(axis=0)[:2]
|
||||
maxy,maxx = maxregion.max(axis=0)[:2]
|
||||
|
||||
if debug:
|
||||
print ("maxregion.size: %d, minx:%d, maxx:%d miny:%d, maxy:%d" % (maxregion.size, minx, maxx, miny, maxy ) )
|
||||
|
||||
lenx = maxx - minx
|
||||
leny = maxy - miny
|
||||
if lenx >= 4 and leny >= 4:
|
||||
masky = int(minx+(lenx//2))
|
||||
maskx = int(miny+(leny//2))
|
||||
lowest_len = min (lenx, leny)
|
||||
|
||||
if debug:
|
||||
print ("lowest_len = %f" % (lowest_len) )
|
||||
|
||||
img_mask_blurry_aaa = img_face_mask_aaa
|
||||
|
||||
if self.erode_mask_modifier != 0:
|
||||
ero = int( lowest_len * ( 0.126 - lowest_len * 0.00004551365 ) * 0.01*self.erode_mask_modifier )
|
||||
if debug:
|
||||
print ("erode_size = %d" % (ero) )
|
||||
if ero > 0:
|
||||
img_mask_blurry_aaa = cv2.erode(img_mask_blurry_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero,ero)), iterations = 1 )
|
||||
elif ero < 0:
|
||||
img_mask_blurry_aaa = cv2.dilate(img_mask_blurry_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(-ero,-ero)), iterations = 1 )
|
||||
|
||||
if self.seamless_erode_mask_modifier != 0:
|
||||
ero = int( lowest_len * ( 0.126 - lowest_len * 0.00004551365 ) * 0.01*self.seamless_erode_mask_modifier )
|
||||
if debug:
|
||||
print ("seamless_erode_size = %d" % (ero) )
|
||||
if ero > 0:
|
||||
img_face_mask_flatten_aaa = cv2.erode(img_face_mask_flatten_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(ero,ero)), iterations = 1 )
|
||||
elif ero < 0:
|
||||
img_face_mask_flatten_aaa = cv2.dilate(img_face_mask_flatten_aaa, cv2.getStructuringElement(cv2.MORPH_ELLIPSE,(-ero,-ero)), iterations = 1 )
|
||||
|
||||
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:
|
||||
print ("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 )
|
||||
|
||||
if debug:
|
||||
debugs += [img_mask_blurry_aaa.copy()]
|
||||
|
||||
if 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 = image_utils.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 = image_utils.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
|
||||
|
||||
hist_match_1 = prd_face_bgr*hist_mask_a + (1.0-hist_mask_a)* np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32)
|
||||
hist_match_1[ hist_match_1 > 1.0 ] = 1.0
|
||||
|
||||
hist_match_2 = dst_face_bgr*hist_mask_a + (1.0-hist_mask_a)* np.ones ( prd_face_bgr.shape[:2] + (1,) , dtype=np.float32)
|
||||
hist_match_2[ hist_match_1 > 1.0 ] = 1.0
|
||||
|
||||
prd_face_bgr = image_utils.color_hist_match(hist_match_1, hist_match_2, self.hist_match_threshold )
|
||||
|
||||
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 self.mode == 'seamless' or self.mode == 'seamless-hist-match':
|
||||
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()]
|
||||
|
||||
out_img = cv2.seamlessClone( (out_img*255).astype(np.uint8), (img_bgr*255).astype(np.uint8), (img_face_mask_flatten_aaa*255).astype(np.uint8), (masky,maskx) , cv2.NORMAL_CLONE )
|
||||
out_img = out_img.astype(dtype=np.float32) / 255.0
|
||||
|
||||
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 self.mode == 'seamless-hist-match':
|
||||
out_face_bgr = cv2.warpAffine( out_img, face_mat, (self.output_size, self.output_size) )
|
||||
new_out_face_bgr = image_utils.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 = image_utils.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
|
||||
|
||||
|
||||
|
|
@ -7,31 +7,34 @@ from pathlib import Path
|
|||
from utils import Path_utils
|
||||
from utils import std_utils
|
||||
from utils import image_utils
|
||||
from utils.console_utils import *
|
||||
from utils.cv2_utils import *
|
||||
import numpy as np
|
||||
import cv2
|
||||
from samples import SampleGeneratorBase
|
||||
from nnlib import nnlib
|
||||
from interact import interact as io
|
||||
'''
|
||||
You can implement your own model. Check examples.
|
||||
'''
|
||||
class ModelBase(object):
|
||||
|
||||
#DONT OVERRIDE
|
||||
def __init__(self, model_path, training_data_src_path=None, training_data_dst_path=None, debug = False, force_gpu_idx=-1, **in_options):
|
||||
|
||||
if force_gpu_idx == -1:
|
||||
def __init__(self, model_path, training_data_src_path=None, training_data_dst_path=None, debug = False, device_args = None):
|
||||
|
||||
device_args['force_gpu_idx'] = device_args.get('force_gpu_idx',-1)
|
||||
|
||||
if device_args['force_gpu_idx'] == -1:
|
||||
idxs_names_list = nnlib.device.getValidDevicesIdxsWithNamesList()
|
||||
if len(idxs_names_list) > 1:
|
||||
print ("You have multi GPUs in a system: ")
|
||||
io.log_info ("You have multi GPUs in a system: ")
|
||||
for idx, name in idxs_names_list:
|
||||
print ("[%d] : %s" % (idx, name) )
|
||||
io.log_info ("[%d] : %s" % (idx, name) )
|
||||
|
||||
force_gpu_idx = input_int("Which GPU idx to choose? ( skip: best GPU ) : ", -1, [ x[0] for x in idxs_names_list] )
|
||||
self.force_gpu_idx = force_gpu_idx
|
||||
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
|
||||
|
||||
print ("Loading model...")
|
||||
io.log_info ("Loading model...")
|
||||
|
||||
self.model_path = model_path
|
||||
self.model_data_path = Path( self.get_strpath_storage_for_file('data.dat') )
|
||||
|
||||
|
|
@ -46,9 +49,7 @@ class ModelBase(object):
|
|||
self.dst_data_generator = None
|
||||
self.debug = debug
|
||||
self.is_training_mode = (training_data_src_path is not None and training_data_dst_path is not None)
|
||||
|
||||
self.supress_std_once = os.environ.get('TF_SUPPRESS_STD', '0') == '1'
|
||||
|
||||
|
||||
self.epoch = 0
|
||||
self.options = {}
|
||||
self.loss_history = []
|
||||
|
|
@ -61,40 +62,40 @@ class ModelBase(object):
|
|||
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.epoch != 0 and input_in_time ("Press enter in 2 seconds to override model settings.", 2)
|
||||
ask_override = self.is_training_mode and self.epoch != 0 and io.input_in_time ("Press enter in 2 seconds to override model settings.", 2)
|
||||
|
||||
if self.epoch == 0:
|
||||
print ("\nModel first run. Enter model options as default for each run.")
|
||||
io.log_info ("\nModel first run. Enter model options as default for each run.")
|
||||
|
||||
if self.epoch == 0 or ask_override:
|
||||
default_write_preview_history = False if self.epoch == 0 else self.options.get('write_preview_history',False)
|
||||
self.options['write_preview_history'] = input_bool("Write preview history? (y/n ?:help skip:n/default) : ", default_write_preview_history, help_message="Preview history will be writed to <ModelName>_history folder.")
|
||||
self.options['write_preview_history'] = io.input_bool("Write preview history? (y/n ?:help skip:n/default) : ", default_write_preview_history, help_message="Preview history will be writed to <ModelName>_history folder.")
|
||||
else:
|
||||
self.options['write_preview_history'] = self.options.get('write_preview_history', False)
|
||||
|
||||
if self.epoch == 0 or ask_override:
|
||||
self.options['target_epoch'] = max(0, input_int("Target epoch (skip:unlimited/default) : ", 0))
|
||||
self.options['target_epoch'] = max(0, io.input_int("Target epoch (skip:unlimited/default) : ", 0))
|
||||
else:
|
||||
self.options['target_epoch'] = self.options.get('target_epoch', 0)
|
||||
|
||||
if self.epoch == 0 or ask_override:
|
||||
default_batch_size = 0 if self.epoch == 0 else self.options.get('batch_size',0)
|
||||
self.options['batch_size'] = max(0, input_int("Batch_size (?:help skip:0/default) : ", default_batch_size, help_message="Larger batch size is always 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:0/default) : ", default_batch_size, help_message="Larger batch size is always better for NN's generalization, but it can cause Out of Memory error. Tune this value for your videocard manually."))
|
||||
else:
|
||||
self.options['batch_size'] = self.options.get('batch_size', 0)
|
||||
|
||||
if self.epoch == 0:
|
||||
self.options['sort_by_yaw'] = 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." )
|
||||
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." )
|
||||
else:
|
||||
self.options['sort_by_yaw'] = self.options.get('sort_by_yaw', False)
|
||||
|
||||
if self.epoch == 0:
|
||||
self.options['random_flip'] = input_bool("Flip faces randomly? (y/n ?:help skip:y) : ", True, help_message="Predicted face will look more naturally without this option, but src faceset should cover all face directions as dst faceset.")
|
||||
self.options['random_flip'] = io.input_bool("Flip faces randomly? (y/n ?:help skip:y) : ", True, help_message="Predicted face will look more naturally without this option, but src faceset should cover all face directions as dst faceset.")
|
||||
else:
|
||||
self.options['random_flip'] = self.options.get('random_flip', True)
|
||||
|
||||
if self.epoch == 0:
|
||||
self.options['src_scale_mod'] = np.clip( input_int("Src face scale modifier % ( -30...30, ?:help skip:0) : ", 0, help_message="If src face shape is wider than dst, try to decrease this value to get a better result."), -30, 30)
|
||||
self.options['src_scale_mod'] = np.clip( io.input_int("Src face scale modifier % ( -30...30, ?:help skip:0) : ", 0, help_message="If src face shape is wider than dst, try to decrease this value to get a better result."), -30, 30)
|
||||
else:
|
||||
self.options['src_scale_mod'] = self.options.get('src_scale_mod', 0)
|
||||
|
||||
|
|
@ -116,10 +117,10 @@ class ModelBase(object):
|
|||
|
||||
self.onInitializeOptions(self.epoch == 0, ask_override)
|
||||
|
||||
nnlib.import_all ( nnlib.DeviceConfig(allow_growth=False, force_gpu_idx=self.force_gpu_idx, **in_options) )
|
||||
nnlib.import_all ( nnlib.DeviceConfig(allow_growth=False, **self.device_args) )
|
||||
self.device_config = nnlib.active_DeviceConfig
|
||||
|
||||
self.onInitialize(**in_options)
|
||||
self.onInitialize()
|
||||
|
||||
self.options['batch_size'] = self.batch_size
|
||||
|
||||
|
|
@ -128,10 +129,10 @@ class ModelBase(object):
|
|||
|
||||
if self.is_training_mode:
|
||||
if self.write_preview_history:
|
||||
if self.force_gpu_idx == -1:
|
||||
if self.device_args['force_gpu_idx'] == -1:
|
||||
self.preview_history_path = self.model_path / ( '%s_history' % (self.get_model_name()) )
|
||||
else:
|
||||
self.preview_history_path = self.model_path / ( '%d_%s_history' % (self.force_gpu_idx, self.get_model_name()) )
|
||||
self.preview_history_path = self.model_path / ( '%d_%s_history' % (self.device_args['force_gpu_idx'], self.get_model_name()) )
|
||||
|
||||
if not self.preview_history_path.exists():
|
||||
self.preview_history_path.mkdir(exist_ok=True)
|
||||
|
|
@ -141,11 +142,11 @@ class ModelBase(object):
|
|||
Path(filename).unlink()
|
||||
|
||||
if self.generator_list is None:
|
||||
raise Exception( 'You didnt set_training_data_generators()')
|
||||
raise ValueError( 'You didnt set_training_data_generators()')
|
||||
else:
|
||||
for i, generator in enumerate(self.generator_list):
|
||||
if not isinstance(generator, SampleGeneratorBase):
|
||||
raise Exception('training data generator is not subclass of SampleGeneratorBase')
|
||||
raise ValueError('training data generator is not subclass of SampleGeneratorBase')
|
||||
|
||||
if (self.sample_for_preview is None) or (self.epoch == 0):
|
||||
self.sample_for_preview = self.generate_next_sample()
|
||||
|
|
@ -181,14 +182,14 @@ class ModelBase(object):
|
|||
model_summary_text += ["========================="]
|
||||
model_summary_text = "\r\n".join (model_summary_text)
|
||||
self.model_summary_text = model_summary_text
|
||||
print(model_summary_text)
|
||||
io.log_info(model_summary_text)
|
||||
|
||||
#overridable
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
pass
|
||||
|
||||
#overridable
|
||||
def onInitialize(self, **in_options):
|
||||
def onInitialize(self):
|
||||
'''
|
||||
initialize your keras models
|
||||
|
||||
|
|
@ -221,10 +222,9 @@ class ModelBase(object):
|
|||
return Path(inspect.getmodule(self).__file__).parent.name.rsplit("_", 1)[1]
|
||||
|
||||
#overridable
|
||||
def get_converter(self, **in_options):
|
||||
#return existing or your own converter which derived from base
|
||||
from .ConverterBase import ConverterBase
|
||||
return ConverterBase(self, **in_options)
|
||||
def get_converter(self):
|
||||
raise NotImplementeError
|
||||
#return existing or your own converter which derived from base
|
||||
|
||||
def get_target_epoch(self):
|
||||
return self.target_epoch
|
||||
|
|
@ -258,17 +258,10 @@ class ModelBase(object):
|
|||
return self.onGetPreview (self.sample_for_preview)[0][1] #first preview, and bgr
|
||||
|
||||
def save(self):
|
||||
print ("Saving...")
|
||||
io.log_info ("Saving...")
|
||||
|
||||
if self.supress_std_once:
|
||||
supressor = std_utils.suppress_stdout_stderr()
|
||||
supressor.__enter__()
|
||||
|
||||
Path( self.get_strpath_storage_for_file('summary.txt') ).write_text(self.model_summary_text)
|
||||
self.onSave()
|
||||
|
||||
if self.supress_std_once:
|
||||
supressor.__exit__()
|
||||
|
||||
model_data = {
|
||||
'epoch': self.epoch,
|
||||
|
|
@ -310,11 +303,7 @@ class ModelBase(object):
|
|||
def generate_next_sample(self):
|
||||
return [next(generator) for generator in self.generator_list]
|
||||
|
||||
def train_one_epoch(self):
|
||||
if self.supress_std_once:
|
||||
supressor = std_utils.suppress_stdout_stderr()
|
||||
supressor.__enter__()
|
||||
|
||||
def train_one_epoch(self):
|
||||
sample = self.generate_next_sample()
|
||||
epoch_time = time.time()
|
||||
losses = self.onTrainOneEpoch(sample, self.generator_list)
|
||||
|
|
@ -322,11 +311,7 @@ class ModelBase(object):
|
|||
self.last_sample = sample
|
||||
|
||||
self.loss_history.append ( [float(loss[1]) for loss in losses] )
|
||||
|
||||
if self.supress_std_once:
|
||||
supressor.__exit__()
|
||||
self.supress_std_once = False
|
||||
|
||||
|
||||
if self.write_preview_history:
|
||||
if self.epoch % 10 == 0:
|
||||
preview = self.get_static_preview()
|
||||
|
|
@ -377,10 +362,10 @@ class ModelBase(object):
|
|||
return self.generator_list
|
||||
|
||||
def get_strpath_storage_for_file(self, filename):
|
||||
if self.force_gpu_idx == -1:
|
||||
if self.device_args['force_gpu_idx'] == -1:
|
||||
return str( self.model_path / ( self.get_model_name() + '_' + filename) )
|
||||
else:
|
||||
return str( self.model_path / ( str(self.force_gpu_idx) + '_' + self.get_model_name() + '_' + filename) )
|
||||
return str( self.model_path / ( str(self.device_args['force_gpu_idx']) + '_' + self.get_model_name() + '_' + filename) )
|
||||
|
||||
def set_vram_batch_requirements (self, d):
|
||||
#example d = {2:2,3:4,4:8,5:16,6:32,7:32,8:32,9:48}
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from nnlib import nnlib
|
|||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
from utils.console_utils import *
|
||||
from interact import interact as io
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
|
|
@ -12,12 +12,12 @@ class Model(ModelBase):
|
|||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
if is_first_run or ask_override:
|
||||
def_pixel_loss = self.options.get('pixel_loss', False)
|
||||
self.options['pixel_loss'] = input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k epochs to enhance fine details and decrease face jitter.")
|
||||
self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k epochs to enhance fine details and decrease face jitter.")
|
||||
else:
|
||||
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
def onInitialize(self):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {4.5:4} )
|
||||
|
||||
|
|
@ -113,15 +113,14 @@ class Model(ModelBase):
|
|||
return np.concatenate ( (x,mx), -1 )
|
||||
|
||||
#override
|
||||
def get_converter(self, **in_options):
|
||||
from models import ConverterMasked
|
||||
def get_converter(self):
|
||||
from converters import ConverterMasked
|
||||
return ConverterMasked(self.predictor_func,
|
||||
predictor_input_size=128,
|
||||
output_size=128,
|
||||
face_type=FaceType.FULL,
|
||||
base_erode_mask_modifier=30,
|
||||
base_blur_mask_modifier=0,
|
||||
**in_options)
|
||||
base_blur_mask_modifier=0)
|
||||
|
||||
def Build(self, input_layer):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
|
|
|||
|
|
@ -4,14 +4,14 @@ from nnlib import nnlib
|
|||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
from utils.console_utils import *
|
||||
from interact import interact as io
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
if is_first_run:
|
||||
self.options['lighter_ae'] = input_bool ("Use lightweight autoencoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight autoencoder is faster, requires less VRAM, sacrificing overall quality. If your GPU VRAM <= 4, you should to choose this option.")
|
||||
self.options['lighter_ae'] = io.input_bool ("Use lightweight autoencoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight autoencoder is faster, requires less VRAM, sacrificing overall quality. If your GPU VRAM <= 4, you should to choose this option.")
|
||||
else:
|
||||
default_lighter_ae = self.options.get('created_vram_gb', 99) <= 4 #temporally support old models, deprecate in future
|
||||
if 'created_vram_gb' in self.options.keys():
|
||||
|
|
@ -20,12 +20,12 @@ class Model(ModelBase):
|
|||
|
||||
if is_first_run or ask_override:
|
||||
def_pixel_loss = self.options.get('pixel_loss', False)
|
||||
self.options['pixel_loss'] = input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k epochs to enhance fine details and decrease face jitter.")
|
||||
self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k epochs to enhance fine details and decrease face jitter.")
|
||||
else:
|
||||
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
def onInitialize(self):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {2.5:4} )
|
||||
|
||||
|
|
@ -125,15 +125,14 @@ class Model(ModelBase):
|
|||
return np.concatenate ( (x,mx), -1 )
|
||||
|
||||
#override
|
||||
def get_converter(self, **in_options):
|
||||
from models import ConverterMasked
|
||||
def get_converter(self):
|
||||
from converters import ConverterMasked
|
||||
return ConverterMasked(self.predictor_func,
|
||||
predictor_input_size=128,
|
||||
output_size=128,
|
||||
face_type=FaceType.HALF,
|
||||
base_erode_mask_modifier=100,
|
||||
base_blur_mask_modifier=100,
|
||||
**in_options)
|
||||
base_blur_mask_modifier=100)
|
||||
|
||||
def Build(self, lighter_ae):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
|
|
|||
|
|
@ -4,14 +4,14 @@ from nnlib import nnlib
|
|||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
from utils.console_utils import *
|
||||
from interact import interact as io
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
#override
|
||||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
if is_first_run:
|
||||
self.options['lighter_ae'] = input_bool ("Use lightweight autoencoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight autoencoder is faster, requires less VRAM, sacrificing overall quality. If your GPU VRAM <= 4, you should to choose this option.")
|
||||
self.options['lighter_ae'] = io.input_bool ("Use lightweight autoencoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight autoencoder is faster, requires less VRAM, sacrificing overall quality. If your GPU VRAM <= 4, you should to choose this option.")
|
||||
else:
|
||||
default_lighter_ae = self.options.get('created_vram_gb', 99) <= 4 #temporally support old models, deprecate in future
|
||||
if 'created_vram_gb' in self.options.keys():
|
||||
|
|
@ -20,12 +20,12 @@ class Model(ModelBase):
|
|||
|
||||
if is_first_run or ask_override:
|
||||
def_pixel_loss = self.options.get('pixel_loss', False)
|
||||
self.options['pixel_loss'] = input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k epochs to enhance fine details and decrease face jitter.")
|
||||
self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k epochs to enhance fine details and decrease face jitter.")
|
||||
else:
|
||||
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
def onInitialize(self):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {1.5:4} )
|
||||
|
||||
|
|
@ -127,15 +127,14 @@ class Model(ModelBase):
|
|||
return np.concatenate ( (x,mx), -1 )
|
||||
|
||||
#override
|
||||
def get_converter(self, **in_options):
|
||||
from models import ConverterMasked
|
||||
def get_converter(self):
|
||||
from converters import ConverterMasked
|
||||
return ConverterMasked(self.predictor_func,
|
||||
predictor_input_size=64,
|
||||
output_size=64,
|
||||
face_type=FaceType.HALF,
|
||||
base_erode_mask_modifier=100,
|
||||
base_blur_mask_modifier=100,
|
||||
**in_options)
|
||||
base_blur_mask_modifier=100)
|
||||
|
||||
def Build(self, lighter_ae):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from nnlib import nnlib
|
|||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
from utils.console_utils import *
|
||||
from interact import interact as io
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
|
|
@ -12,12 +12,12 @@ class Model(ModelBase):
|
|||
def onInitializeOptions(self, is_first_run, ask_override):
|
||||
if is_first_run or ask_override:
|
||||
def_pixel_loss = self.options.get('pixel_loss', False)
|
||||
self.options['pixel_loss'] = input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k epochs to enhance fine details and decrease face jitter.")
|
||||
self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 20k epochs to enhance fine details and decrease face jitter.")
|
||||
else:
|
||||
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
def onInitialize(self):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {4.5:4} )
|
||||
|
||||
|
|
@ -121,15 +121,14 @@ class Model(ModelBase):
|
|||
return np.concatenate ( (x,mx), -1 )
|
||||
|
||||
#override
|
||||
def get_converter(self, **in_options):
|
||||
from models import ConverterMasked
|
||||
def get_converter(self):
|
||||
from converters import ConverterMasked
|
||||
return ConverterMasked(self.predictor_func,
|
||||
predictor_input_size=128,
|
||||
output_size=128,
|
||||
face_type=FaceType.FULL,
|
||||
base_erode_mask_modifier=30,
|
||||
base_blur_mask_modifier=0,
|
||||
**in_options)
|
||||
base_blur_mask_modifier=0)
|
||||
|
||||
def Build(self, input_layer):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
|
|
|||
|
|
@ -4,7 +4,7 @@ from nnlib import nnlib
|
|||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
from utils.console_utils import *
|
||||
from interact import interact as io
|
||||
|
||||
#SAE - Styled AutoEncoder
|
||||
class SAEModel(ModelBase):
|
||||
|
|
@ -27,15 +27,15 @@ class SAEModel(ModelBase):
|
|||
default_face_type = 'f'
|
||||
|
||||
if is_first_run:
|
||||
resolution = 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.")
|
||||
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.")
|
||||
resolution = np.clip (resolution, 64, 256)
|
||||
while np.modf(resolution / 16)[0] != 0.0:
|
||||
resolution -= 1
|
||||
self.options['resolution'] = resolution
|
||||
|
||||
self.options['face_type'] = 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'] = 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.")
|
||||
self.options['archi'] = input_str ("AE architecture (df, liae, vg ?:help skip:%s) : " % (default_archi) , default_archi, ['df','liae','vg'], help_message="'df' keeps faces more natural. 'liae' can fix overly different face shapes. 'vg' - currently testing.").lower()
|
||||
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 ("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.")
|
||||
self.options['archi'] = io.input_str ("AE architecture (df, liae, vg ?:help skip:%s) : " % (default_archi) , default_archi, ['df','liae','vg'], help_message="'df' keeps faces more natural. 'liae' can fix overly different face shapes. 'vg' - currently testing.").lower()
|
||||
else:
|
||||
self.options['resolution'] = self.options.get('resolution', default_resolution)
|
||||
self.options['face_type'] = self.options.get('face_type', default_face_type)
|
||||
|
|
@ -45,17 +45,17 @@ class SAEModel(ModelBase):
|
|||
default_ae_dims = 256 if self.options['archi'] == 'liae' else 512
|
||||
default_ed_ch_dims = 42
|
||||
if is_first_run:
|
||||
self.options['ae_dims'] = np.clip ( input_int("AutoEncoder dims (32-1024 ?:help skip:%d) : " % (default_ae_dims) , default_ae_dims, help_message="More dims are better, but requires more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 )
|
||||
self.options['ed_ch_dims'] = np.clip ( input_int("Encoder/Decoder dims per channel (21-85 ?:help skip:%d) : " % (default_ed_ch_dims) , default_ed_ch_dims, help_message="More dims are better, but requires more VRAM. You can fine-tune model size to fit your GPU." ), 21, 85 )
|
||||
self.options['ae_dims'] = np.clip ( io.input_int("AutoEncoder dims (32-1024 ?:help skip:%d) : " % (default_ae_dims) , default_ae_dims, help_message="More dims are better, but requires more VRAM. You can fine-tune model size to fit your GPU." ), 32, 1024 )
|
||||
self.options['ed_ch_dims'] = np.clip ( io.input_int("Encoder/Decoder dims per channel (21-85 ?:help skip:%d) : " % (default_ed_ch_dims) , default_ed_ch_dims, help_message="More dims are better, but requires more VRAM. You can fine-tune model size to fit your GPU." ), 21, 85 )
|
||||
else:
|
||||
self.options['ae_dims'] = self.options.get('ae_dims', default_ae_dims)
|
||||
self.options['ed_ch_dims'] = self.options.get('ed_ch_dims', default_ed_ch_dims)
|
||||
|
||||
if is_first_run:
|
||||
self.options['lighter_encoder'] = input_bool ("Use lightweight encoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight encoder is 35% faster, requires less VRAM, but sacrificing overall quality.")
|
||||
self.options['lighter_encoder'] = io.input_bool ("Use lightweight encoder? (y/n, ?:help skip:n) : ", False, help_message="Lightweight encoder is 35% faster, requires less VRAM, but sacrificing overall quality.")
|
||||
|
||||
if self.options['archi'] != 'vg':
|
||||
self.options['multiscale_decoder'] = input_bool ("Use multiscale decoder? (y/n, ?:help skip:y) : ", True, help_message="Multiscale decoder helps to get better details.")
|
||||
self.options['multiscale_decoder'] = io.input_bool ("Use multiscale decoder? (y/n, ?:help skip:y) : ", True, help_message="Multiscale decoder helps to get better details.")
|
||||
else:
|
||||
self.options['lighter_encoder'] = self.options.get('lighter_encoder', False)
|
||||
|
||||
|
|
@ -66,14 +66,14 @@ class SAEModel(ModelBase):
|
|||
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'] = input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 15-25k epochs to enhance fine details and decrease face jitter.")
|
||||
self.options['pixel_loss'] = io.input_bool ("Use pixel loss? (y/n, ?:help skip: n/default ) : ", def_pixel_loss, help_message="Default DSSIM loss good for initial understanding structure of faces. Use pixel loss after 15-25k epochs to enhance fine details and decrease face jitter.")
|
||||
|
||||
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 ( input_number("Face style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_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,
|
||||
help_message="Learn to transfer face style details such as light and color conditions. Warning: Enable it only after 10k epochs, when predicted face is clear enough to start learn style. Start from 0.1 value and check history changes."), 0.0, 100.0 )
|
||||
|
||||
default_bg_style_power = default_bg_style_power if is_first_run else self.options.get('bg_style_power', default_bg_style_power)
|
||||
self.options['bg_style_power'] = np.clip ( input_number("Background style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_bg_style_power), default_bg_style_power,
|
||||
self.options['bg_style_power'] = np.clip ( io.input_number("Background style power ( 0.0 .. 100.0 ?:help skip:%.2f) : " % (default_bg_style_power), default_bg_style_power,
|
||||
help_message="Learn to transfer image around face. This can make face more like dst."), 0.0, 100.0 )
|
||||
else:
|
||||
self.options['pixel_loss'] = self.options.get('pixel_loss', False)
|
||||
|
|
@ -81,7 +81,7 @@ class SAEModel(ModelBase):
|
|||
self.options['bg_style_power'] = self.options.get('bg_style_power', default_bg_style_power)
|
||||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
def onInitialize(self):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
SAEModel.initialize_nn_functions()
|
||||
|
||||
|
|
@ -427,9 +427,7 @@ class SAEModel(ModelBase):
|
|||
return np.concatenate ( [prd[0], prd[1]], -1 )
|
||||
|
||||
#override
|
||||
def get_converter(self, **in_options):
|
||||
from models import ConverterMasked
|
||||
|
||||
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
|
||||
|
||||
|
|
@ -439,6 +437,7 @@ class SAEModel(ModelBase):
|
|||
|
||||
face_type = FaceType.FULL if self.options['face_type'] == 'f' else FaceType.HALF
|
||||
|
||||
from converters import ConverterMasked
|
||||
return ConverterMasked(self.predictor_func,
|
||||
predictor_input_size=self.options['resolution'],
|
||||
output_size=self.options['resolution'],
|
||||
|
|
@ -448,8 +447,7 @@ class SAEModel(ModelBase):
|
|||
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,
|
||||
**in_options)
|
||||
clip_hborder_mask_per=0.0625 if self.options['face_type'] == 'f' else 0)
|
||||
|
||||
@staticmethod
|
||||
def initialize_nn_functions():
|
||||
|
|
|
|||
|
|
@ -1,7 +1,4 @@
|
|||
from .ModelBase import ModelBase
|
||||
from .ConverterBase import ConverterBase
|
||||
from .ConverterMasked import ConverterMasked
|
||||
from .ConverterImage import ConverterImage
|
||||
|
||||
def import_model(name):
|
||||
module = __import__('Model_'+name, globals(), locals(), [], 1)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue