mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-20 21:43:21 -07:00
updated pdf manuals for AVATAR model.
Avatar converter: added super resolution option. All converters: super resolution DCSCN network is now replaced by RankSRGAN
This commit is contained in:
iperov
parent
19c66286da
commit
c39ed9d9c9
15 changed files with 161 additions and 188 deletions
Binary file not shown.
|
|
@ -1,164 +0,0 @@
|
|||
import numpy as np
|
||||
import cv2
|
||||
from pathlib import Path
|
||||
from nnlib import nnlib
|
||||
from interact import interact as io
|
||||
|
||||
class DCSCN():
|
||||
def __init__(self):
|
||||
exec( nnlib.import_all(), locals(), globals() )
|
||||
|
||||
inp_x = KL.Input([None, None, 1])
|
||||
inp_x2 = KL.Input([None, None, 1])
|
||||
|
||||
x = inp_x
|
||||
layers_count = 12
|
||||
layers = []
|
||||
for i in range(1,layers_count+1):
|
||||
if i == 1:
|
||||
output_feature_num = 196
|
||||
else:
|
||||
x1 = (i-1) / float(layers_count - 1)
|
||||
y1 = x1 ** (1.0 / 1.5)
|
||||
output_feature_num = int((196 - 48) * (1 - y1) + 48)
|
||||
x = Conv2D(output_feature_num, kernel_size=3, strides=1, padding='same', name='CNN%d' % (i) ) (x)
|
||||
x = PReLU(shared_axes=[1,2], name='CNN%d_prelu' % (i) ) (x)
|
||||
layers.append(x)
|
||||
|
||||
x_concat = KL.Concatenate()(layers)
|
||||
|
||||
A1 = Conv2D(64, kernel_size=1, strides=1, padding='same', name='A1' ) (x_concat)
|
||||
A1 = PReLU(shared_axes=[1,2], name='A1_prelu') (A1)
|
||||
|
||||
B1 = Conv2D(32, kernel_size=1, strides=1, padding='same', name='B1' ) (x_concat)
|
||||
B1 = PReLU(shared_axes=[1,2], name='B1_prelu') (B1)
|
||||
|
||||
B2 = Conv2D(32, kernel_size=3, strides=1, padding='same', name='B2' ) (B1)
|
||||
B2 = PReLU(shared_axes=[1,2], name='B2_prelu') (B2)
|
||||
|
||||
x = KL.Concatenate()([B2,A1])
|
||||
x = Conv2D(96*4, kernel_size=3, strides=1, padding='same', name='Up_PS' )(x)
|
||||
x = PixelShuffler()(x)
|
||||
x = Conv2D(1, kernel_size=3, strides=1, padding='same', name='R_CNN1', use_bias=False )(x)
|
||||
x = KL.Add()([x, inp_x2])
|
||||
self.model = keras.models.Model ([inp_x, inp_x2], [x])
|
||||
self.model.load_weights ( Path(__file__).parent / 'DCSCN.h5' )
|
||||
|
||||
def upscale(self, img, is_bgr=True, is_float=True):
|
||||
if is_bgr:
|
||||
img = img[...,::-1]
|
||||
|
||||
if is_float:
|
||||
img = np.clip (img*255, 0, 255)
|
||||
|
||||
img_shape_len = len(img.shape)
|
||||
h, w = img.shape[:2]
|
||||
ch = img.shape[2] if len(img.shape) >= 3 else 1
|
||||
|
||||
nh, nw = h*2, w*2
|
||||
|
||||
img_x = self.convert_rgb_to_y(img)
|
||||
|
||||
img_bx = cv2.resize(img_x, (nh, nw), cv2.INTER_CUBIC)
|
||||
|
||||
ensemble = 8
|
||||
|
||||
output = np.zeros([nh,nw,1], dtype=np.float32)
|
||||
|
||||
for i in range(ensemble):
|
||||
x = np.reshape( self.flip(img_x, i), (1,h,w,1) )
|
||||
bx = np.reshape( self.flip(img_bx, i), (1,nh,nw,1) )
|
||||
y = self.model.predict([x,bx])[0]
|
||||
y = self.flip(y, i, invert=True)
|
||||
output += y
|
||||
|
||||
output /= ensemble
|
||||
|
||||
bimg = cv2.resize(img, (nh, nw), cv2.INTER_CUBIC)
|
||||
bimg_ycbcr = self.convert_rgb_to_ycbcr(bimg)
|
||||
|
||||
if ch > 1:
|
||||
output = self.convert_y_and_cbcr_to_rgb(output, bimg_ycbcr[:, :, 1:3])
|
||||
|
||||
if is_float:
|
||||
output = np.clip (output/255.0, 0, 1.0)
|
||||
|
||||
if is_bgr:
|
||||
output = output[...,::-1]
|
||||
|
||||
return output
|
||||
|
||||
def convert_rgb_to_y(self, image):
|
||||
if len(image.shape) <= 2 or image.shape[2] == 1:
|
||||
return image
|
||||
|
||||
xform = np.array([[65.738 / 256.0, 129.057 / 256.0, 25.064 / 256.0]], dtype=np.float32)
|
||||
y_image = image.dot(xform.T) + 16.0
|
||||
|
||||
return y_image
|
||||
|
||||
|
||||
def convert_rgb_to_ycbcr(self, image):
|
||||
if len(image.shape) <= 2 or image.shape[2] == 1:
|
||||
return image
|
||||
|
||||
xform = np.array(
|
||||
[[65.738 / 256.0, 129.057 / 256.0, 25.064 / 256.0],
|
||||
[- 37.945 / 256.0, - 74.494 / 256.0, 112.439 / 256.0],
|
||||
[112.439 / 256.0, - 94.154 / 256.0, - 18.285 / 256.0]], dtype=np.float32)
|
||||
|
||||
ycbcr_image = image.dot(xform.T)
|
||||
ycbcr_image[:, :, 0] += 16.0
|
||||
ycbcr_image[:, :, [1, 2]] += 128.0
|
||||
|
||||
return ycbcr_image
|
||||
|
||||
def convert_ycbcr_to_rgb(self,ycbcr_image):
|
||||
rgb_image = np.zeros([ycbcr_image.shape[0], ycbcr_image.shape[1], 3], dtype=np.float32)
|
||||
|
||||
rgb_image[:, :, 0] = ycbcr_image[:, :, 0] - 16.0
|
||||
rgb_image[:, :, [1, 2]] = ycbcr_image[:, :, [1, 2]] - 128.0
|
||||
xform = np.array(
|
||||
[[298.082 / 256.0, 0, 408.583 / 256.0],
|
||||
[298.082 / 256.0, -100.291 / 256.0, -208.120 / 256.0],
|
||||
[298.082 / 256.0, 516.412 / 256.0, 0]], dtype=np.float32)
|
||||
rgb_image = rgb_image.dot(xform.T)
|
||||
|
||||
return rgb_image
|
||||
|
||||
def convert_y_and_cbcr_to_rgb(self,y_image, cbcr_image):
|
||||
if len(y_image.shape) <= 2:
|
||||
y_image = y_image.reshape[y_image.shape[0], y_image.shape[1], 1]
|
||||
|
||||
if len(y_image.shape) == 3 and y_image.shape[2] == 3:
|
||||
y_image = y_image[:, :, 0:1]
|
||||
|
||||
ycbcr_image = np.zeros([y_image.shape[0], y_image.shape[1], 3], dtype=np.float32)
|
||||
ycbcr_image[:, :, 0] = y_image[:, :, 0]
|
||||
ycbcr_image[:, :, 1:3] = cbcr_image[:, :, 0:2]
|
||||
|
||||
return self.convert_ycbcr_to_rgb(ycbcr_image)
|
||||
|
||||
def flip(self, image, flip_type, invert=False):
|
||||
if flip_type == 0:
|
||||
return image
|
||||
elif flip_type == 1:
|
||||
return np.flipud(image)
|
||||
elif flip_type == 2:
|
||||
return np.fliplr(image)
|
||||
elif flip_type == 3:
|
||||
return np.flipud(np.fliplr(image))
|
||||
elif flip_type == 4:
|
||||
return np.rot90(image, 1 if invert is False else -1)
|
||||
elif flip_type == 5:
|
||||
return np.rot90(image, -1 if invert is False else 1)
|
||||
elif flip_type == 6:
|
||||
if invert is False:
|
||||
return np.flipud(np.rot90(image))
|
||||
else:
|
||||
return np.rot90(np.flipud(image), -1)
|
||||
elif flip_type == 7:
|
||||
if invert is False:
|
||||
return np.flipud(np.rot90(image, -1))
|
||||
else:
|
||||
return np.rot90(np.flipud(image), 1)
|
||||
BIN
imagelib/RankSRGAN.h5
Normal file
BIN
imagelib/RankSRGAN.h5
Normal file
Binary file not shown.
109
imagelib/RankSRGAN.py
Normal file
109
imagelib/RankSRGAN.py
Normal file
|
|
@ -0,0 +1,109 @@
|
|||
import numpy as np
|
||||
import cv2
|
||||
from pathlib import Path
|
||||
from nnlib import nnlib
|
||||
from interact import interact as io
|
||||
|
||||
class RankSRGAN():
|
||||
def __init__(self):
|
||||
exec( nnlib.import_all(), locals(), globals() )
|
||||
|
||||
class PixelShufflerTorch(KL.Layer):
|
||||
def __init__(self, size=(2, 2), data_format='channels_last', **kwargs):
|
||||
super(PixelShufflerTorch, self).__init__(**kwargs)
|
||||
self.data_format = data_format
|
||||
self.size = size
|
||||
|
||||
def call(self, inputs):
|
||||
input_shape = K.shape(inputs)
|
||||
if K.int_shape(input_shape)[0] != 4:
|
||||
raise ValueError('Inputs should have rank 4; Received input shape:', str(K.int_shape(inputs)))
|
||||
|
||||
batch_size, h, w, c = input_shape[0], input_shape[1], input_shape[2], K.int_shape(inputs)[-1]
|
||||
rh, rw = self.size
|
||||
oh, ow = h * rh, w * rw
|
||||
oc = c // (rh * rw)
|
||||
|
||||
out = inputs
|
||||
out = K.permute_dimensions(out, (0, 3, 1, 2)) #NCHW
|
||||
|
||||
out = K.reshape(out, (batch_size, oc, rh, rw, h, w))
|
||||
out = K.permute_dimensions(out, (0, 1, 4, 2, 5, 3))
|
||||
out = K.reshape(out, (batch_size, oc, oh, ow))
|
||||
|
||||
out = K.permute_dimensions(out, (0, 2, 3, 1))
|
||||
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))
|
||||
|
||||
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(PixelShufflerTorch, self).get_config()
|
||||
|
||||
return dict(list(base_config.items()) + list(config.items()))
|
||||
|
||||
def res_block(inp, name_prefix):
|
||||
x = inp
|
||||
x = Conv2D (ndf, kernel_size=3, strides=1, padding='same', activation="relu", name=name_prefix+"0")(x)
|
||||
x = Conv2D (ndf, kernel_size=3, strides=1, padding='same', name=name_prefix+"2")(x)
|
||||
return Add()([inp,x])
|
||||
|
||||
ndf = 64
|
||||
nb = 16
|
||||
inp = Input ( (None, None,3) )
|
||||
x = inp
|
||||
|
||||
x = x0 = Conv2D (ndf, kernel_size=3, strides=1, padding='same', name="model0")(x)
|
||||
for i in range(nb):
|
||||
x = res_block(x, "model1%.2d" %i )
|
||||
x = Conv2D (ndf, kernel_size=3, strides=1, padding='same', name="model1160")(x)
|
||||
x = Add()([x0,x])
|
||||
|
||||
x = ReLU() ( PixelShufflerTorch() ( Conv2D (ndf*4, kernel_size=3, strides=1, padding='same', name="model2")(x) ) )
|
||||
x = ReLU() ( PixelShufflerTorch() ( Conv2D (ndf*4, kernel_size=3, strides=1, padding='same', name="model5")(x) ) )
|
||||
|
||||
x = Conv2D (ndf, kernel_size=3, strides=1, padding='same', activation="relu", name="model8")(x)
|
||||
x = Conv2D (3, kernel_size=3, strides=1, padding='same', name="model10")(x)
|
||||
self.model = Model(inp, x )
|
||||
self.model.load_weights ( Path(__file__).parent / 'RankSRGAN.h5')
|
||||
|
||||
def upscale(self, img, scale=2, is_bgr=True, is_float=True):
|
||||
if scale not in [2,4]:
|
||||
raise ValueError ("RankSRGAN: supported scale are 2 or 4.")
|
||||
|
||||
if not is_bgr:
|
||||
img = img[...,::-1]
|
||||
|
||||
if not is_float:
|
||||
img /= 255.0
|
||||
|
||||
h, w = img.shape[:2]
|
||||
ch = img.shape[2] if len(img.shape) >= 3 else 1
|
||||
|
||||
output = self.model.predict([img[None,...]])[0]
|
||||
|
||||
if scale == 2:
|
||||
output = cv2.resize (output, (w*scale, h*scale), cv2.INTER_CUBIC)
|
||||
|
||||
if not is_float:
|
||||
output = np.clip (output * 255.0, 0, 255.0)
|
||||
|
||||
if not is_bgr:
|
||||
output = output[...,::-1]
|
||||
|
||||
return output
|
||||
|
|
@ -13,7 +13,7 @@ from .reduce_colors import reduce_colors
|
|||
|
||||
from .color_transfer import color_hist_match, reinhard_color_transfer, linear_color_transfer
|
||||
|
||||
from .DCSCN import DCSCN
|
||||
from .RankSRGAN import RankSRGAN
|
||||
|
||||
from .common import normalize_channels, overlay_alpha_image
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue