mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-07-10 15:23:32 -07:00
added fanseg for future WF segmentation model
This commit is contained in:
Colombo
parent
3b6ad4abf9
commit
143792fd31
6 changed files with 429 additions and 112 deletions
|
|
@ -24,44 +24,44 @@ class TernausNet(object):
|
|||
tf = nn.tf
|
||||
|
||||
class Ternaus(nn.ModelBase):
|
||||
def on_build(self, in_ch, ch):
|
||||
def on_build(self, in_ch, base_ch):
|
||||
|
||||
self.features_0 = nn.Conv2D (in_ch, ch, kernel_size=3, padding='SAME')
|
||||
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 (ch, ch*2, kernel_size=3, padding='SAME')
|
||||
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 (ch*2, ch*4, kernel_size=3, padding='SAME')
|
||||
self.features_8 = nn.Conv2D (ch*4, ch*4, kernel_size=3, padding='SAME')
|
||||
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 (ch*4, ch*8, kernel_size=3, padding='SAME')
|
||||
self.features_13 = nn.Conv2D (ch*8, ch*8, kernel_size=3, padding='SAME')
|
||||
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 (ch*8, ch*8, kernel_size=3, padding='SAME')
|
||||
self.features_18 = nn.Conv2D (ch*8, ch*8, kernel_size=3, padding='SAME')
|
||||
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 (ch*8, ch*8, kernel_size=3, padding='SAME')
|
||||
self.conv_center = nn.Conv2D (base_ch*8, base_ch*8, kernel_size=3, padding='SAME')
|
||||
|
||||
self.conv1_up = nn.Conv2DTranspose (ch*8, ch*4, kernel_size=3, padding='SAME')
|
||||
self.conv1 = nn.Conv2D (ch*12, 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 (ch*8, ch*4, kernel_size=3, padding='SAME')
|
||||
self.conv2 = nn.Conv2D (ch*12, 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 (ch*8, ch*2, kernel_size=3, padding='SAME')
|
||||
self.conv3 = nn.Conv2D (ch*6, ch*4, 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 (ch*4, ch, kernel_size=3, padding='SAME')
|
||||
self.conv4 = nn.Conv2D (ch*3, ch*2, 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 (ch*2, ch//2, kernel_size=3, padding='SAME')
|
||||
self.conv5 = nn.Conv2D (ch//2+ch, ch, 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 (ch, 1, kernel_size=3, padding='SAME')
|
||||
self.out_conv = nn.Conv2D (base_ch, 1, kernel_size=3, padding='SAME')
|
||||
|
||||
def forward(self, inp):
|
||||
x, = inp
|
||||
|
|
@ -106,92 +106,120 @@ class TernausNet(object):
|
|||
x = tf.concat( [x,x0], -1)
|
||||
x = tf.nn.relu(self.conv5(x))
|
||||
|
||||
x = tf.nn.sigmoid(self.out_conv(x))
|
||||
return 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:
|
||||
weights_file_root = Path(__file__).parent
|
||||
self.weights_path = weights_file_root / ('%s_%d_%s.npy' % (name, resolution, face_type_str) )
|
||||
self.weights_file_root = weights_file_root
|
||||
|
||||
e = tf.device('/CPU:0') if place_model_on_cpu else None
|
||||
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) )
|
||||
|
||||
if e is not None: e.__enter__()
|
||||
self.net = Ternaus(3, 64, name='Ternaus')
|
||||
if load_weights:
|
||||
self.net.load_weights (self.weights_path)
|
||||
# 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_weights = self.net.get_weights()
|
||||
|
||||
self.model_filename_list = [ [self.net, '%s_%d_%s.npy' % (name, resolution, face_type_str) ] ]
|
||||
|
||||
if training:
|
||||
self.opt = nn.TFRMSpropOptimizer(lr=0.0001, name='opt')
|
||||
self.opt.initialize_variables (self.net_weights, vars_on_cpu=place_model_on_cpu)
|
||||
self.model_filename_list += [ [self.opt, '%s_%d_%s_opt.npy' % (name, resolution, face_type_str) ] ]
|
||||
else:
|
||||
self.net.init_weights()
|
||||
if e is not None: e.__exit__(None,None,None)
|
||||
_, 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
|
||||
|
||||
self.net.build_for_run ( [(tf.float32, nn.get4Dshape (resolution,resolution,3) )] )
|
||||
# Loading/initializing all models/optimizers weights
|
||||
for model, filename in io.progress_bar_generator(self.model_filename_list, "Initializing models"):
|
||||
do_init = not load_weights
|
||||
|
||||
if training:
|
||||
raise Exception("training not supported yet")
|
||||
if not do_init:
|
||||
do_init = not model.load_weights( self.weights_file_root / filename )
|
||||
|
||||
if do_init:
|
||||
model.init_weights()
|
||||
if model == self.net:
|
||||
try:
|
||||
with open( Path(__file__).parent / 'vgg11_enc_weights.npy', 'rb' ) as f:
|
||||
d = pickle.loads (f.read())
|
||||
|
||||
"""
|
||||
if training:
|
||||
try:
|
||||
with open( Path(__file__).parent / 'vgg11_enc_weights.npy', 'rb' ) as f:
|
||||
d = pickle.loads (f.read())
|
||||
|
||||
for i in [0,3,6,8,11,13,16,18]:
|
||||
s = 'features.%d' % i
|
||||
|
||||
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.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], self.model.trainable_weights) )
|
||||
"""
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type=None, exc_value=None, traceback=None):
|
||||
return False #pass exception between __enter__ and __exit__ to outter level
|
||||
for i in [0,3,6,8,11,13,16,18]:
|
||||
model.get_layer_by_name ('features_%d' % i).set_weights ( d['features.%d' % i] )
|
||||
except:
|
||||
io.log_err("Unable to load VGG11 pretrained weights from vgg11_enc_weights.npy")
|
||||
|
||||
def save_weights(self):
|
||||
self.net.save_weights (str(self.weights_path))
|
||||
|
||||
def train(self, inp, real):
|
||||
loss, = self.train_func ([inp, real])
|
||||
return loss
|
||||
for model, filename in io.progress_bar_generator(self.model_filename_list, "Saving"):
|
||||
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[np.newaxis,...]
|
||||
input_image = input_image[None,...]
|
||||
|
||||
result = np.clip ( self.net.run([input_image]), 0, 1.0 )
|
||||
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
|
||||
|
||||
"""
|
||||
if load_weights:
|
||||
self.net.load_weights (self.weights_path)
|
||||
else:
|
||||
self.net.init_weights()
|
||||
|
||||
if load_weights:
|
||||
self.opt.load_weights (self.opt_path)
|
||||
else:
|
||||
self.opt.init_weights()
|
||||
"""
|
||||
"""
|
||||
if training:
|
||||
try:
|
||||
with open( Path(__file__).parent / 'vgg11_enc_weights.npy', 'rb' ) as f:
|
||||
d = pickle.loads (f.read())
|
||||
|
||||
for i in [0,3,6,8,11,13,16,18]:
|
||||
s = 'features.%d' % i
|
||||
|
||||
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.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], self.model.trainable_weights) )
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue