mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-03-12 20:42:45 -07:00
d46fb5cfd3
added FacesetEnhancer 4.2.other) data_src util faceset enhance best GPU.bat 4.2.other) data_src util faceset enhance multi GPU.bat FacesetEnhancer greatly increases details in your source face set, same as Gigapixel enhancer, but in fully automatic mode. In OpenCL build it works on CPU only. Please consider a donation.
155 lines
6.1 KiB
Python
155 lines
6.1 KiB
Python
import operator
|
|
from pathlib import Path
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
|
|
|
|
class FaceEnhancer(object):
|
|
"""
|
|
x4 face enhancer
|
|
"""
|
|
def __init__(self):
|
|
from nnlib import nnlib
|
|
exec( nnlib.import_all(), locals(), globals() )
|
|
|
|
model_path = Path(__file__).parent / "FaceEnhancer.h5"
|
|
if not model_path.exists():
|
|
return
|
|
|
|
bgr_inp = Input ( (192,192,3) )
|
|
t_param_inp = Input ( (1,) )
|
|
t_param1_inp = Input ( (1,) )
|
|
x = Conv2D (64, 3, strides=1, padding='same' )(bgr_inp)
|
|
|
|
a = Dense (64, use_bias=False) ( t_param_inp )
|
|
a = Reshape( (1,1,64) )(a)
|
|
b = Dense (64, use_bias=False ) ( t_param1_inp )
|
|
b = Reshape( (1,1,64) )(b)
|
|
x = Add()([x,a,b])
|
|
|
|
x = LeakyReLU(0.1)(x)
|
|
|
|
x = LeakyReLU(0.1)(Conv2D (64, 3, strides=1, padding='same' )(x))
|
|
x = e0 = LeakyReLU(0.1)(Conv2D (64, 3, strides=1, padding='same')(x))
|
|
|
|
x = AveragePooling2D()(x)
|
|
x = LeakyReLU(0.1)(Conv2D (112, 3, strides=1, padding='same')(x))
|
|
x = e1 = LeakyReLU(0.1)(Conv2D (112, 3, strides=1, padding='same')(x))
|
|
|
|
x = AveragePooling2D()(x)
|
|
x = LeakyReLU(0.1)(Conv2D (192, 3, strides=1, padding='same')(x))
|
|
x = e2 = LeakyReLU(0.1)(Conv2D (192, 3, strides=1, padding='same')(x))
|
|
|
|
x = AveragePooling2D()(x)
|
|
x = LeakyReLU(0.1)(Conv2D (336, 3, strides=1, padding='same')(x))
|
|
x = e3 = LeakyReLU(0.1)(Conv2D (336, 3, strides=1, padding='same')(x))
|
|
|
|
x = AveragePooling2D()(x)
|
|
x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
|
|
x = e4 = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
|
|
|
|
x = AveragePooling2D()(x)
|
|
x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
|
|
x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
|
|
x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
|
|
x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
|
|
|
|
x = Concatenate()([ BilinearInterpolation()(x), e4 ])
|
|
|
|
x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
|
|
x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
|
|
|
|
x = Concatenate()([ BilinearInterpolation()(x), e3 ])
|
|
x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
|
|
x = LeakyReLU(0.1)(Conv2D (512, 3, strides=1, padding='same')(x))
|
|
|
|
x = Concatenate()([ BilinearInterpolation()(x), e2 ])
|
|
x = LeakyReLU(0.1)(Conv2D (288, 3, strides=1, padding='same')(x))
|
|
x = LeakyReLU(0.1)(Conv2D (288, 3, strides=1, padding='same')(x))
|
|
|
|
x = Concatenate()([ BilinearInterpolation()(x), e1 ])
|
|
x = LeakyReLU(0.1)(Conv2D (160, 3, strides=1, padding='same')(x))
|
|
x = LeakyReLU(0.1)(Conv2D (160, 3, strides=1, padding='same')(x))
|
|
|
|
x = Concatenate()([ BilinearInterpolation()(x), e0 ])
|
|
x = LeakyReLU(0.1)(Conv2D (96, 3, strides=1, padding='same')(x))
|
|
x = d0 = LeakyReLU(0.1)(Conv2D (96, 3, strides=1, padding='same')(x))
|
|
|
|
x = LeakyReLU(0.1)(Conv2D (48, 3, strides=1, padding='same')(x))
|
|
|
|
x = Conv2D (3, 3, strides=1, padding='same', activation='tanh')(x)
|
|
out1x = Add()([bgr_inp, x])
|
|
|
|
x = d0
|
|
x = LeakyReLU(0.1)(Conv2D (96, 3, strides=1, padding='same')(x))
|
|
x = LeakyReLU(0.1)(Conv2D (96, 3, strides=1, padding='same')(x))
|
|
x = d2x = BilinearInterpolation()(x)
|
|
|
|
x = LeakyReLU(0.1)(Conv2D (48, 3, strides=1, padding='same')(x))
|
|
x = Conv2D (3, 3, strides=1, padding='same', activation='tanh')(x)
|
|
|
|
out2x = Add()([BilinearInterpolation()(out1x), x])
|
|
|
|
x = d2x
|
|
x = LeakyReLU(0.1)(Conv2D (72, 3, strides=1, padding='same')(x))
|
|
x = LeakyReLU(0.1)(Conv2D (72, 3, strides=1, padding='same')(x))
|
|
x = d4x = BilinearInterpolation()(x)
|
|
|
|
x = LeakyReLU(0.1)(Conv2D (36, 3, strides=1, padding='same')(x))
|
|
x = Conv2D (3, 3, strides=1, padding='same', activation='tanh')(x)
|
|
out4x = Add()([BilinearInterpolation()(out2x), x ])
|
|
|
|
self.model = keras.models.Model ( [bgr_inp,t_param_inp,t_param1_inp], [out4x] )
|
|
self.model.load_weights (str(model_path))
|
|
|
|
|
|
def enhance (self, inp_img, is_tanh=False, preserve_size=True):
|
|
if not is_tanh:
|
|
inp_img = np.clip( inp_img * 2 -1, -1, 1 )
|
|
|
|
param = np.array([0.2])
|
|
param1 = np.array([1.0])
|
|
up_res = 4
|
|
patch_size = 192
|
|
patch_size_half = patch_size // 2
|
|
|
|
h,w,c = inp_img.shape
|
|
|
|
i_max = w-patch_size+1
|
|
j_max = h-patch_size+1
|
|
|
|
final_img = np.zeros ( (h*up_res,w*up_res,c), dtype=np.float32 )
|
|
final_img_div = np.zeros ( (h*up_res,w*up_res,1), dtype=np.float32 )
|
|
|
|
x = np.concatenate ( [ np.linspace (0,1,patch_size_half*up_res), np.linspace (1,0,patch_size_half*up_res) ] )
|
|
x,y = np.meshgrid(x,x)
|
|
patch_mask = (x*y)[...,None]
|
|
|
|
j=0
|
|
while j < j_max:
|
|
i = 0
|
|
while i < i_max:
|
|
patch_img = inp_img[j:j+patch_size, i:i+patch_size,:]
|
|
x = self.model.predict( [ patch_img[None,...], param, param1 ] )[0]
|
|
final_img [j*up_res:(j+patch_size)*up_res, i*up_res:(i+patch_size)*up_res,:] += x*patch_mask
|
|
final_img_div[j*up_res:(j+patch_size)*up_res, i*up_res:(i+patch_size)*up_res,:] += patch_mask
|
|
if i == i_max-1:
|
|
break
|
|
i = min( i+patch_size_half, i_max-1)
|
|
if j == j_max-1:
|
|
break
|
|
j = min( j+patch_size_half, j_max-1)
|
|
|
|
final_img_div[final_img_div==0] = 1.0
|
|
final_img /= final_img_div
|
|
|
|
if preserve_size:
|
|
final_img = cv2.resize (final_img, (w,h), cv2.INTER_LANCZOS4)
|
|
|
|
if not is_tanh:
|
|
final_img = np.clip( final_img/2+0.5, 0, 1 )
|
|
|
|
return final_img
|