mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-20 13:33:24 -07:00
initial
This commit is contained in:
parent
73de93b4f1
commit
6bd5a44264
71 changed files with 8448 additions and 0 deletions
198
nnlib/__init__.py
Normal file
198
nnlib/__init__.py
Normal file
|
@ -0,0 +1,198 @@
|
|||
def tf_image_histogram (tf, input):
|
||||
x = input
|
||||
x += 1 / 255.0
|
||||
|
||||
output = []
|
||||
for i in range(256, 0, -1):
|
||||
v = i / 255.0
|
||||
y = (x - v) * 1000
|
||||
|
||||
y = tf.clip_by_value (y, -1.0, 0.0) + 1
|
||||
|
||||
output.append ( tf.reduce_sum (y) )
|
||||
x -= y*v
|
||||
|
||||
return tf.stack ( output[::-1] )
|
||||
|
||||
def tf_dssim(tf, t1, t2):
|
||||
return (1.0 - tf.image.ssim (t1, t2, 1.0)) / 2.0
|
||||
|
||||
def tf_ssim(tf, t1, t2):
|
||||
return tf.image.ssim (t1, t2, 1.0)
|
||||
|
||||
def DSSIMMaskLossClass(tf):
|
||||
class DSSIMMaskLoss(object):
|
||||
def __init__(self, mask_list, is_tanh=False):
|
||||
self.mask_list = mask_list
|
||||
self.is_tanh = is_tanh
|
||||
|
||||
def __call__(self,y_true, y_pred):
|
||||
total_loss = None
|
||||
for mask in self.mask_list:
|
||||
|
||||
if not self.is_tanh:
|
||||
loss = (1.0 - tf.image.ssim (y_true*mask, y_pred*mask, 1.0)) / 2.0
|
||||
else:
|
||||
loss = (1.0 - tf.image.ssim ( (y_true/2+0.5)*(mask/2+0.5), (y_pred/2+0.5)*(mask/2+0.5), 1.0)) / 2.0
|
||||
|
||||
if total_loss is None:
|
||||
total_loss = loss
|
||||
else:
|
||||
total_loss += loss
|
||||
|
||||
return total_loss
|
||||
|
||||
return DSSIMMaskLoss
|
||||
|
||||
def MSEMaskLossClass(keras):
|
||||
class MSEMaskLoss(object):
|
||||
def __init__(self, mask_list, is_tanh=False):
|
||||
self.mask_list = mask_list
|
||||
self.is_tanh = is_tanh
|
||||
|
||||
def __call__(self,y_true, y_pred):
|
||||
K = keras.backend
|
||||
|
||||
total_loss = None
|
||||
for mask in self.mask_list:
|
||||
|
||||
if not self.is_tanh:
|
||||
loss = K.mean(K.square(y_true*mask - y_pred*mask))
|
||||
else:
|
||||
loss = K.mean(K.square( (y_true/2+0.5)*(mask/2+0.5) - (y_pred/2+0.5)*(mask/2+0.5) ))
|
||||
|
||||
if total_loss is None:
|
||||
total_loss = loss
|
||||
else:
|
||||
total_loss += loss
|
||||
|
||||
return total_loss
|
||||
|
||||
return MSEMaskLoss
|
||||
|
||||
def PixelShufflerClass(keras):
|
||||
class PixelShuffler(keras.engine.topology.Layer):
|
||||
def __init__(self, size=(2, 2), data_format=None, **kwargs):
|
||||
super(PixelShuffler, self).__init__(**kwargs)
|
||||
self.data_format = keras.utils.conv_utils.normalize_data_format(data_format)
|
||||
self.size = keras.utils.conv_utils.normalize_tuple(size, 2, 'size')
|
||||
|
||||
def call(self, inputs):
|
||||
|
||||
input_shape = keras.backend.int_shape(inputs)
|
||||
if len(input_shape) != 4:
|
||||
raise ValueError('Inputs should have rank ' +
|
||||
str(4) +
|
||||
'; Received input shape:', str(input_shape))
|
||||
|
||||
if self.data_format == 'channels_first':
|
||||
batch_size, c, h, w = input_shape
|
||||
if batch_size is None:
|
||||
batch_size = -1
|
||||
rh, rw = self.size
|
||||
oh, ow = h * rh, w * rw
|
||||
oc = c // (rh * rw)
|
||||
|
||||
out = keras.backend.reshape(inputs, (batch_size, rh, rw, oc, h, w))
|
||||
out = keras.backend.permute_dimensions(out, (0, 3, 4, 1, 5, 2))
|
||||
out = keras.backend.reshape(out, (batch_size, oc, oh, ow))
|
||||
return out
|
||||
|
||||
elif self.data_format == 'channels_last':
|
||||
batch_size, h, w, c = input_shape
|
||||
if batch_size is None:
|
||||
batch_size = -1
|
||||
rh, rw = self.size
|
||||
oh, ow = h * rh, w * rw
|
||||
oc = c // (rh * rw)
|
||||
|
||||
out = keras.backend.reshape(inputs, (batch_size, h, w, rh, rw, oc))
|
||||
out = keras.backend.permute_dimensions(out, (0, 1, 3, 2, 4, 5))
|
||||
out = keras.backend.reshape(out, (batch_size, oh, ow, oc))
|
||||
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))
|
||||
|
||||
if self.data_format == 'channels_first':
|
||||
height = input_shape[2] * self.size[0] if input_shape[2] is not None else None
|
||||
width = input_shape[3] * self.size[1] if input_shape[3] is not None else None
|
||||
channels = input_shape[1] // self.size[0] // self.size[1]
|
||||
|
||||
if channels * self.size[0] * self.size[1] != input_shape[1]:
|
||||
raise ValueError('channels of input and size are incompatible')
|
||||
|
||||
return (input_shape[0],
|
||||
channels,
|
||||
height,
|
||||
width)
|
||||
|
||||
elif self.data_format == 'channels_last':
|
||||
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(PixelShuffler, self).get_config()
|
||||
|
||||
return dict(list(base_config.items()) + list(config.items()))
|
||||
return PixelShuffler
|
||||
|
||||
def conv(keras, input_tensor, filters):
|
||||
x = input_tensor
|
||||
x = keras.layers.convolutional.Conv2D(filters, kernel_size=5, strides=2, padding='same')(x)
|
||||
x = keras.layers.advanced_activations.LeakyReLU(0.1)(x)
|
||||
return x
|
||||
|
||||
def upscale(keras, input_tensor, filters, k_size=3):
|
||||
x = input_tensor
|
||||
x = keras.layers.convolutional.Conv2D(filters * 4, kernel_size=k_size, padding='same')(x)
|
||||
x = keras.layers.advanced_activations.LeakyReLU(0.1)(x)
|
||||
x = PixelShufflerClass(keras)()(x)
|
||||
return x
|
||||
|
||||
def upscale4(keras, input_tensor, filters):
|
||||
x = input_tensor
|
||||
x = keras.layers.convolutional.Conv2D(filters * 16, kernel_size=3, padding='same')(x)
|
||||
x = keras.layers.advanced_activations.LeakyReLU(0.1)(x)
|
||||
x = PixelShufflerClass(keras)(size=(4, 4))(x)
|
||||
return x
|
||||
|
||||
def res(keras, input_tensor, filters):
|
||||
x = input_tensor
|
||||
x = keras.layers.convolutional.Conv2D(filters, kernel_size=3, kernel_initializer=keras.initializers.RandomNormal(0, 0.02), use_bias=False, padding="same")(x)
|
||||
x = keras.layers.advanced_activations.LeakyReLU(alpha=0.2)(x)
|
||||
x = keras.layers.convolutional.Conv2D(filters, kernel_size=3, kernel_initializer=keras.initializers.RandomNormal(0, 0.02), use_bias=False, padding="same")(x)
|
||||
x = keras.layers.Add()([x, input_tensor])
|
||||
x = keras.layers.advanced_activations.LeakyReLU(alpha=0.2)(x)
|
||||
return x
|
||||
|
||||
def resize_like(tf, keras, ref_tensor, input_tensor):
|
||||
def func(input_tensor, ref_tensor):
|
||||
H, W = ref_tensor.get_shape()[1], ref_tensor.get_shape()[2]
|
||||
return tf.image.resize_bilinear(input_tensor, [H.value, W.value])
|
||||
|
||||
return keras.layers.Lambda(func, arguments={'ref_tensor':ref_tensor})(input_tensor)
|
||||
|
||||
def total_variation_loss(keras, x):
|
||||
K = keras.backend
|
||||
assert K.ndim(x) == 4
|
||||
B,H,W,C = K.int_shape(x)
|
||||
a = K.square(x[:, :H - 1, :W - 1, :] - x[:, 1:, :W - 1, :])
|
||||
b = K.square(x[:, :H - 1, :W - 1, :] - x[:, :H - 1, 1:, :])
|
||||
|
||||
return K.mean (a+b)
|
Loading…
Add table
Add a link
Reference in a new issue