mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-20 21:43:21 -07:00
global refactoring and fixes,
removed support of extracted(aligned) PNG faces. Use old builds to convert from PNG to JPG. fanseg model file in facelib/ is renamed
This commit is contained in:
Colombo
parent
921b464d5b
commit
61472cdaf7
82 changed files with 3838 additions and 3812 deletions
95
facelib/DFLSegNet.py
Normal file
95
facelib/DFLSegNet.py
Normal file
|
|
@ -0,0 +1,95 @@
|
|||
import os
|
||||
import pickle
|
||||
from functools import partial
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
|
||||
from core.interact import interact as io
|
||||
from core.leras import nn
|
||||
|
||||
|
||||
class DFLSegNet(object):
|
||||
VERSION = 1
|
||||
|
||||
def __init__ (self, name,
|
||||
resolution,
|
||||
load_weights=True,
|
||||
weights_file_root=None,
|
||||
training=False,
|
||||
place_model_on_cpu=False,
|
||||
run_on_cpu=False,
|
||||
optimizer=None,
|
||||
data_format="NHWC"):
|
||||
|
||||
nn.initialize(data_format=data_format)
|
||||
tf = nn.tf
|
||||
|
||||
self.weights_file_root = Path(weights_file_root) if weights_file_root is not None else Path(__file__).parent
|
||||
|
||||
with tf.device ('/CPU:0'):
|
||||
#Place holders on CPU
|
||||
self.input_t = tf.placeholder (nn.floatx, nn.get4Dshape(resolution,resolution,3) )
|
||||
self.target_t = tf.placeholder (nn.floatx, nn.get4Dshape(resolution,resolution,1) )
|
||||
|
||||
# Initializing model classes
|
||||
archi = nn.DFLSegnetArchi()
|
||||
with tf.device ('/CPU:0' if place_model_on_cpu else '/GPU:0'):
|
||||
self.enc = archi.Encoder(3, 64, name='Encoder')
|
||||
self.dec = archi.Decoder(64, 1, name='Decoder')
|
||||
self.enc_dec_weights = self.enc.get_weights()+self.dec.get_weights()
|
||||
|
||||
model_name = f'{name}_{resolution}'
|
||||
|
||||
self.model_filename_list = [ [self.enc, f'{model_name}_enc.npy'],
|
||||
[self.dec, f'{model_name}_dec.npy'],
|
||||
]
|
||||
|
||||
if training:
|
||||
if optimizer is None:
|
||||
raise ValueError("Optimizer should be provided for training mode.")
|
||||
|
||||
self.opt = optimizer
|
||||
self.opt.initialize_variables (self.enc_dec_weights, vars_on_cpu=place_model_on_cpu)
|
||||
self.model_filename_list += [ [self.opt, f'{model_name}_opt.npy' ] ]
|
||||
else:
|
||||
with tf.device ('/CPU:0' if run_on_cpu else '/GPU:0'):
|
||||
_, pred = self.dec(self.enc(self.input_t))
|
||||
|
||||
def net_run(input_np):
|
||||
return nn.tf_sess.run ( [pred], feed_dict={self.input_t :input_np})[0]
|
||||
self.net_run = net_run
|
||||
|
||||
# Loading/initializing all models/optimizers weights
|
||||
for model, filename in self.model_filename_list:
|
||||
do_init = not load_weights
|
||||
|
||||
if not do_init:
|
||||
do_init = not model.load_weights( self.weights_file_root / filename )
|
||||
|
||||
if do_init:
|
||||
model.init_weights()
|
||||
|
||||
def flow(self, x):
|
||||
return self.dec(self.enc(x))
|
||||
|
||||
def get_weights(self):
|
||||
return self.enc_dec_weights
|
||||
|
||||
def save_weights(self):
|
||||
for model, filename in io.progress_bar_generator(self.model_filename_list, "Saving", leave=False):
|
||||
model.save_weights( self.weights_file_root / filename )
|
||||
|
||||
def extract (self, input_image):
|
||||
input_shape_len = len(input_image.shape)
|
||||
if input_shape_len == 3:
|
||||
input_image = input_image[None,...]
|
||||
|
||||
result = np.clip ( self.net_run(input_image), 0, 1.0 )
|
||||
result[result < 0.1] = 0 #get rid of noise
|
||||
|
||||
if input_shape_len == 3:
|
||||
result = result[0]
|
||||
|
||||
return result
|
||||
|
|
@ -89,7 +89,7 @@ class FANExtractor(object):
|
|||
low2 = self.b2_plus(low1)
|
||||
low3 = self.b3(low2)
|
||||
|
||||
up2 = nn.tf_upsample2d(low3)
|
||||
up2 = nn.upsample2d(low3)
|
||||
|
||||
return up1+up2
|
||||
|
||||
|
|
|
|||
|
|
@ -10,7 +10,7 @@ class FaceEnhancer(object):
|
|||
"""
|
||||
x4 face enhancer
|
||||
"""
|
||||
def __init__(self, place_model_on_cpu=False):
|
||||
def __init__(self, place_model_on_cpu=False, run_on_cpu=False):
|
||||
nn.initialize(data_format="NHWC")
|
||||
tf = nn.tf
|
||||
|
||||
|
|
@ -111,23 +111,23 @@ class FaceEnhancer(object):
|
|||
x = tf.nn.leaky_relu(self.center_conv2(x), 0.1)
|
||||
x = tf.nn.leaky_relu(self.center_conv3(x), 0.1)
|
||||
|
||||
x = tf.concat( [nn.tf_resize2d_bilinear(x), e4], -1 )
|
||||
x = tf.concat( [nn.resize2d_bilinear(x), e4], -1 )
|
||||
x = tf.nn.leaky_relu(self.d4_conv0(x), 0.1)
|
||||
x = tf.nn.leaky_relu(self.d4_conv1(x), 0.1)
|
||||
|
||||
x = tf.concat( [nn.tf_resize2d_bilinear(x), e3], -1 )
|
||||
x = tf.concat( [nn.resize2d_bilinear(x), e3], -1 )
|
||||
x = tf.nn.leaky_relu(self.d3_conv0(x), 0.1)
|
||||
x = tf.nn.leaky_relu(self.d3_conv1(x), 0.1)
|
||||
|
||||
x = tf.concat( [nn.tf_resize2d_bilinear(x), e2], -1 )
|
||||
x = tf.concat( [nn.resize2d_bilinear(x), e2], -1 )
|
||||
x = tf.nn.leaky_relu(self.d2_conv0(x), 0.1)
|
||||
x = tf.nn.leaky_relu(self.d2_conv1(x), 0.1)
|
||||
|
||||
x = tf.concat( [nn.tf_resize2d_bilinear(x), e1], -1 )
|
||||
x = tf.concat( [nn.resize2d_bilinear(x), e1], -1 )
|
||||
x = tf.nn.leaky_relu(self.d1_conv0(x), 0.1)
|
||||
x = tf.nn.leaky_relu(self.d1_conv1(x), 0.1)
|
||||
|
||||
x = tf.concat( [nn.tf_resize2d_bilinear(x), e0], -1 )
|
||||
x = tf.concat( [nn.resize2d_bilinear(x), e0], -1 )
|
||||
x = tf.nn.leaky_relu(self.d0_conv0(x), 0.1)
|
||||
x = d0 = tf.nn.leaky_relu(self.d0_conv1(x), 0.1)
|
||||
|
||||
|
|
@ -138,22 +138,22 @@ class FaceEnhancer(object):
|
|||
x = d0
|
||||
x = tf.nn.leaky_relu(self.dec2x_conv0(x), 0.1)
|
||||
x = tf.nn.leaky_relu(self.dec2x_conv1(x), 0.1)
|
||||
x = d2x = nn.tf_resize2d_bilinear(x)
|
||||
x = d2x = nn.resize2d_bilinear(x)
|
||||
|
||||
x = tf.nn.leaky_relu(self.out2x_conv0(x), 0.1)
|
||||
x = self.out2x_conv1(x)
|
||||
|
||||
out2x = nn.tf_resize2d_bilinear(out1x) + tf.nn.tanh(x)
|
||||
out2x = nn.resize2d_bilinear(out1x) + tf.nn.tanh(x)
|
||||
|
||||
x = d2x
|
||||
x = tf.nn.leaky_relu(self.dec4x_conv0(x), 0.1)
|
||||
x = tf.nn.leaky_relu(self.dec4x_conv1(x), 0.1)
|
||||
x = d4x = nn.tf_resize2d_bilinear(x)
|
||||
x = d4x = nn.resize2d_bilinear(x)
|
||||
|
||||
x = tf.nn.leaky_relu(self.out4x_conv0(x), 0.1)
|
||||
x = self.out4x_conv1(x)
|
||||
|
||||
out4x = nn.tf_resize2d_bilinear(out2x) + tf.nn.tanh(x)
|
||||
out4x = nn.resize2d_bilinear(out2x) + tf.nn.tanh(x)
|
||||
|
||||
return out4x
|
||||
|
||||
|
|
@ -161,17 +161,15 @@ class FaceEnhancer(object):
|
|||
if not model_path.exists():
|
||||
raise Exception("Unable to load FaceEnhancer.npy")
|
||||
|
||||
e = tf.device("/CPU:0") if place_model_on_cpu else None
|
||||
if e is not None: e.__enter__()
|
||||
self.model = FaceEnhancer()
|
||||
self.model.load_weights (model_path)
|
||||
if e is not None: e.__exit__(None,None,None)
|
||||
|
||||
self.model.build_for_run ([ (tf.float32, nn.get4Dshape (192,192,3) ),
|
||||
(tf.float32, (None,1,) ),
|
||||
(tf.float32, (None,1,) ),
|
||||
])
|
||||
|
||||
with tf.device ('/CPU:0' if place_model_on_cpu else '/GPU:0'):
|
||||
self.model = FaceEnhancer()
|
||||
self.model.load_weights (model_path)
|
||||
|
||||
with tf.device ('/CPU:0' if run_on_cpu else '/GPU:0'):
|
||||
self.model.build_for_run ([ (tf.float32, nn.get4Dshape (192,192,3) ),
|
||||
(tf.float32, (None,1,) ),
|
||||
(tf.float32, (None,1,) ),
|
||||
])
|
||||
|
||||
def enhance (self, inp_img, is_tanh=False, preserve_size=True):
|
||||
if not is_tanh:
|
||||
|
|
|
|||
|
|
@ -21,7 +21,7 @@ class S3FDExtractor(object):
|
|||
super().__init__(**kwargs)
|
||||
|
||||
def build_weights(self):
|
||||
self.weight = tf.get_variable ("weight", (1, 1, 1, self.n_channels), dtype=nn.tf_floatx, initializer=tf.initializers.ones )
|
||||
self.weight = tf.get_variable ("weight", (1, 1, 1, self.n_channels), dtype=nn.floatx, initializer=tf.initializers.ones )
|
||||
|
||||
def get_weights(self):
|
||||
return [self.weight]
|
||||
|
|
@ -36,7 +36,7 @@ class S3FDExtractor(object):
|
|||
super().__init__(name='S3FD')
|
||||
|
||||
def on_build(self):
|
||||
self.minus = tf.constant([104,117,123], dtype=nn.tf_floatx )
|
||||
self.minus = tf.constant([104,117,123], dtype=nn.floatx )
|
||||
self.conv1_1 = nn.Conv2D(3, 64, kernel_size=3, strides=1, padding='SAME')
|
||||
self.conv1_2 = nn.Conv2D(64, 64, kernel_size=3, strides=1, padding='SAME')
|
||||
|
||||
|
|
|
|||
|
|
@ -9,106 +9,13 @@ import numpy as np
|
|||
from core.interact import interact as io
|
||||
from core.leras import nn
|
||||
|
||||
"""
|
||||
Dataset used to train located in official DFL mega.nz folder
|
||||
https://mega.nz/#F!b9MzCK4B!zEAG9txu7uaRUjXz9PtBqg
|
||||
|
||||
using https://github.com/ternaus/TernausNet
|
||||
TernausNet: U-Net with VGG11 Encoder Pre-Trained on ImageNet for Image Segmentation
|
||||
"""
|
||||
|
||||
class TernausNet(object):
|
||||
VERSION = 1
|
||||
def __init__ (self, name, resolution, face_type_str=None, load_weights=True, weights_file_root=None, training=False, place_model_on_cpu=False, data_format="NHWC"):
|
||||
|
||||
def __init__ (self, name, resolution, load_weights=True, weights_file_root=None, training=False, place_model_on_cpu=False, run_on_cpu=False, optimizer=None, data_format="NHWC"):
|
||||
nn.initialize(data_format=data_format)
|
||||
tf = nn.tf
|
||||
|
||||
class Ternaus(nn.ModelBase):
|
||||
def on_build(self, in_ch, base_ch):
|
||||
|
||||
self.features_0 = nn.Conv2D (in_ch, base_ch, kernel_size=3, padding='SAME')
|
||||
self.blurpool_0 = nn.BlurPool (filt_size=3)
|
||||
|
||||
self.features_3 = nn.Conv2D (base_ch, base_ch*2, kernel_size=3, padding='SAME')
|
||||
self.blurpool_3 = nn.BlurPool (filt_size=3)
|
||||
|
||||
self.features_6 = nn.Conv2D (base_ch*2, base_ch*4, kernel_size=3, padding='SAME')
|
||||
self.features_8 = nn.Conv2D (base_ch*4, base_ch*4, kernel_size=3, padding='SAME')
|
||||
self.blurpool_8 = nn.BlurPool (filt_size=3)
|
||||
|
||||
self.features_11 = nn.Conv2D (base_ch*4, base_ch*8, kernel_size=3, padding='SAME')
|
||||
self.features_13 = nn.Conv2D (base_ch*8, base_ch*8, kernel_size=3, padding='SAME')
|
||||
self.blurpool_13 = nn.BlurPool (filt_size=3)
|
||||
|
||||
self.features_16 = nn.Conv2D (base_ch*8, base_ch*8, kernel_size=3, padding='SAME')
|
||||
self.features_18 = nn.Conv2D (base_ch*8, base_ch*8, kernel_size=3, padding='SAME')
|
||||
self.blurpool_18 = nn.BlurPool (filt_size=3)
|
||||
|
||||
self.conv_center = nn.Conv2D (base_ch*8, base_ch*8, kernel_size=3, padding='SAME')
|
||||
|
||||
self.conv1_up = nn.Conv2DTranspose (base_ch*8, base_ch*4, kernel_size=3, padding='SAME')
|
||||
self.conv1 = nn.Conv2D (base_ch*12, base_ch*8, kernel_size=3, padding='SAME')
|
||||
|
||||
self.conv2_up = nn.Conv2DTranspose (base_ch*8, base_ch*4, kernel_size=3, padding='SAME')
|
||||
self.conv2 = nn.Conv2D (base_ch*12, base_ch*8, kernel_size=3, padding='SAME')
|
||||
|
||||
self.conv3_up = nn.Conv2DTranspose (base_ch*8, base_ch*2, kernel_size=3, padding='SAME')
|
||||
self.conv3 = nn.Conv2D (base_ch*6, base_ch*4, kernel_size=3, padding='SAME')
|
||||
|
||||
self.conv4_up = nn.Conv2DTranspose (base_ch*4, base_ch, kernel_size=3, padding='SAME')
|
||||
self.conv4 = nn.Conv2D (base_ch*3, base_ch*2, kernel_size=3, padding='SAME')
|
||||
|
||||
self.conv5_up = nn.Conv2DTranspose (base_ch*2, base_ch//2, kernel_size=3, padding='SAME')
|
||||
self.conv5 = nn.Conv2D (base_ch//2+base_ch, base_ch, kernel_size=3, padding='SAME')
|
||||
|
||||
self.out_conv = nn.Conv2D (base_ch, 1, kernel_size=3, padding='SAME')
|
||||
|
||||
def forward(self, inp):
|
||||
x, = inp
|
||||
|
||||
x = x0 = tf.nn.relu(self.features_0(x))
|
||||
x = self.blurpool_0(x)
|
||||
|
||||
x = x1 = tf.nn.relu(self.features_3(x))
|
||||
x = self.blurpool_3(x)
|
||||
|
||||
x = tf.nn.relu(self.features_6(x))
|
||||
x = x2 = tf.nn.relu(self.features_8(x))
|
||||
x = self.blurpool_8(x)
|
||||
|
||||
x = tf.nn.relu(self.features_11(x))
|
||||
x = x3 = tf.nn.relu(self.features_13(x))
|
||||
x = self.blurpool_13(x)
|
||||
|
||||
x = tf.nn.relu(self.features_16(x))
|
||||
x = x4 = tf.nn.relu(self.features_18(x))
|
||||
x = self.blurpool_18(x)
|
||||
|
||||
x = self.conv_center(x)
|
||||
|
||||
x = tf.nn.relu(self.conv1_up(x))
|
||||
x = tf.concat( [x,x4], nn.conv2d_ch_axis)
|
||||
x = tf.nn.relu(self.conv1(x))
|
||||
|
||||
x = tf.nn.relu(self.conv2_up(x))
|
||||
x = tf.concat( [x,x3], nn.conv2d_ch_axis)
|
||||
x = tf.nn.relu(self.conv2(x))
|
||||
|
||||
x = tf.nn.relu(self.conv3_up(x))
|
||||
x = tf.concat( [x,x2], nn.conv2d_ch_axis)
|
||||
x = tf.nn.relu(self.conv3(x))
|
||||
|
||||
x = tf.nn.relu(self.conv4_up(x))
|
||||
x = tf.concat( [x,x1], nn.conv2d_ch_axis)
|
||||
x = tf.nn.relu(self.conv4(x))
|
||||
|
||||
x = tf.nn.relu(self.conv5_up(x))
|
||||
x = tf.concat( [x,x0], nn.conv2d_ch_axis)
|
||||
x = tf.nn.relu(self.conv5(x))
|
||||
|
||||
logits = self.out_conv(x)
|
||||
return logits, tf.nn.sigmoid(logits)
|
||||
|
||||
if weights_file_root is not None:
|
||||
weights_file_root = Path(weights_file_root)
|
||||
else:
|
||||
|
|
@ -117,39 +24,42 @@ class TernausNet(object):
|
|||
|
||||
with tf.device ('/CPU:0'):
|
||||
#Place holders on CPU
|
||||
self.input_t = tf.placeholder (nn.tf_floatx, nn.get4Dshape(resolution,resolution,3) )
|
||||
self.target_t = tf.placeholder (nn.tf_floatx, nn.get4Dshape(resolution,resolution,1) )
|
||||
self.input_t = tf.placeholder (nn.floatx, nn.get4Dshape(resolution,resolution,3) )
|
||||
self.target_t = tf.placeholder (nn.floatx, nn.get4Dshape(resolution,resolution,1) )
|
||||
|
||||
# Initializing model classes
|
||||
with tf.device ('/CPU:0' if place_model_on_cpu else '/GPU:0'):
|
||||
self.net = Ternaus(3, 64, name='Ternaus')
|
||||
self.net = nn.Ternaus(3, 64, name='Ternaus')
|
||||
self.net_weights = self.net.get_weights()
|
||||
|
||||
model_name = f'{name}_{resolution}'
|
||||
if face_type_str is not None:
|
||||
model_name += f'_{face_type_str}'
|
||||
model_name = f'{name}_{resolution}'
|
||||
|
||||
self.model_filename_list = [ [self.net, f'{model_name}.npy'] ]
|
||||
|
||||
if training:
|
||||
self.opt = nn.TFRMSpropOptimizer(lr=0.0001, name='opt')
|
||||
if optimizer is None:
|
||||
raise ValueError("Optimizer should be provided for traning mode.")
|
||||
|
||||
self.opt = optimizer
|
||||
self.opt.initialize_variables (self.net_weights, vars_on_cpu=place_model_on_cpu)
|
||||
self.model_filename_list += [ [self.opt, f'{model_name}_opt.npy' ] ]
|
||||
else:
|
||||
_, pred = self.net([self.input_t])
|
||||
with tf.device ('/CPU:0' if run_on_cpu else '/GPU:0'):
|
||||
_, pred = self.net([self.input_t])
|
||||
|
||||
def net_run(input_np):
|
||||
return nn.tf_sess.run ( [pred], feed_dict={self.input_t :input_np})[0]
|
||||
self.net_run = net_run
|
||||
|
||||
# Loading/initializing all models/optimizers weights
|
||||
for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"):
|
||||
for model, filename in self.model_filename_list:
|
||||
do_init = not load_weights
|
||||
|
||||
if not do_init:
|
||||
do_init = not model.load_weights( self.weights_file_root / filename )
|
||||
|
||||
if do_init:
|
||||
model.init_weights()
|
||||
model.init_weights()
|
||||
if model == self.net:
|
||||
try:
|
||||
with open( Path(__file__).parent / 'vgg11_enc_weights.npy', 'rb' ) as f:
|
||||
|
|
@ -177,7 +87,7 @@ class TernausNet(object):
|
|||
|
||||
return result
|
||||
|
||||
"""
|
||||
"""
|
||||
if load_weights:
|
||||
self.net.load_weights (self.weights_path)
|
||||
else:
|
||||
|
|
|
|||
|
|
@ -2,4 +2,5 @@ from .FaceType import FaceType
|
|||
from .S3FDExtractor import S3FDExtractor
|
||||
from .FANExtractor import FANExtractor
|
||||
from .FaceEnhancer import FaceEnhancer
|
||||
from .TernausNet import TernausNet
|
||||
from .TernausNet import TernausNet
|
||||
from .DFLSegNet import DFLSegNet
|
||||
Loading…
Add table
Add a link
Reference in a new issue