mirror of
https://github.com/iperov/DeepFaceLab.git
synced 2025-08-19 21:13:20 -07:00
refactoring. Added RecycleGAN for testing.
This commit is contained in:
iperov
parent
8686309417
commit
f8824f9601
24 changed files with 1661 additions and 1505 deletions
|
|
@ -6,6 +6,7 @@ You can implement your own Converter, check example ConverterMasked.py
|
|||
class ConverterBase(object):
|
||||
MODE_FACE = 0
|
||||
MODE_IMAGE = 1
|
||||
MODE_IMAGE_WITH_LANDMARKS = 2
|
||||
|
||||
#overridable
|
||||
def __init__(self, predictor):
|
||||
|
|
|
|||
|
|
@ -34,7 +34,7 @@ class ConverterImage(ConverterBase):
|
|||
self.predictor ( np.zeros ( (self.predictor_input_size, self.predictor_input_size,3), dtype=np.float32) )
|
||||
|
||||
#override
|
||||
def convert_image (self, img_bgr, img_landmarks, debug):
|
||||
def convert_image (self, img_bgr, debug):
|
||||
img_size = img_bgr.shape[1], img_bgr.shape[0]
|
||||
|
||||
predictor_input_bgr = cv2.resize ( img_bgr, (self.predictor_input_size, self.predictor_input_size), cv2.INTER_LANCZOS4 )
|
||||
|
|
@ -42,5 +42,5 @@ class ConverterImage(ConverterBase):
|
|||
|
||||
output = cv2.resize ( predicted_bgr, (self.output_size, self.output_size), cv2.INTER_LANCZOS4 )
|
||||
if debug:
|
||||
return (img_bgr,output,)
|
||||
return (predictor_input_bgr,output,)
|
||||
return output
|
||||
|
|
|
|||
|
|
@ -9,9 +9,8 @@ from utils import std_utils
|
|||
from utils import image_utils
|
||||
import numpy as np
|
||||
import cv2
|
||||
import gpufmkmgr
|
||||
from samples import SampleGeneratorBase
|
||||
|
||||
from nnlib import nnlib
|
||||
'''
|
||||
You can implement your own model. Check examples.
|
||||
'''
|
||||
|
|
@ -63,27 +62,22 @@ class ModelBase(object):
|
|||
if self.epoch == 0:
|
||||
for filename in Path_utils.get_image_paths(self.preview_history_path):
|
||||
Path(filename).unlink()
|
||||
|
||||
|
||||
self.gpu_config = gpufmkmgr.GPUConfig(allow_growth=False, **in_options)
|
||||
self.gpu_total_vram_gb = self.gpu_config.gpu_total_vram_gb
|
||||
|
||||
self.device_config = nnlib.DeviceConfig(allow_growth=False, **in_options)
|
||||
|
||||
if self.epoch == 0:
|
||||
#first run
|
||||
self.options['created_vram_gb'] = self.gpu_total_vram_gb
|
||||
self.created_vram_gb = self.gpu_total_vram_gb
|
||||
self.options['created_vram_gb'] = self.device_config.gpu_total_vram_gb
|
||||
self.created_vram_gb = self.device_config.gpu_total_vram_gb
|
||||
else:
|
||||
#not first run
|
||||
if 'created_vram_gb' in self.options.keys():
|
||||
self.created_vram_gb = self.options['created_vram_gb']
|
||||
else:
|
||||
self.options['created_vram_gb'] = self.gpu_total_vram_gb
|
||||
self.created_vram_gb = self.gpu_total_vram_gb
|
||||
|
||||
self.tf = gpufmkmgr.import_tf( self.gpu_config )
|
||||
self.tf_sess = gpufmkmgr.get_tf_session()
|
||||
self.keras = gpufmkmgr.import_keras()
|
||||
self.keras_contrib = gpufmkmgr.import_keras_contrib()
|
||||
self.options['created_vram_gb'] = self.device_config.gpu_total_vram_gb
|
||||
self.created_vram_gb = self.device_config.gpu_total_vram_gb
|
||||
|
||||
nnlib.import_all (self.device_config)
|
||||
|
||||
self.onInitialize(**in_options)
|
||||
|
||||
|
|
@ -108,18 +102,18 @@ class ModelBase(object):
|
|||
print ("==")
|
||||
print ("== Options:")
|
||||
print ("== |== batch_size : %s " % (self.batch_size) )
|
||||
print ("== |== multi_gpu : %s " % (self.gpu_config.multi_gpu) )
|
||||
print ("== |== multi_gpu : %s " % (self.device_config.multi_gpu) )
|
||||
for key in self.options.keys():
|
||||
print ("== |== %s : %s" % (key, self.options[key]) )
|
||||
|
||||
print ("== Running on:")
|
||||
if self.gpu_config.cpu_only:
|
||||
if self.device_config.cpu_only:
|
||||
print ("== |== [CPU]")
|
||||
else:
|
||||
for idx in self.gpu_config.gpu_idxs:
|
||||
print ("== |== [%d : %s]" % (idx, gpufmkmgr.getDeviceName(idx)) )
|
||||
for idx in self.device_config.gpu_idxs:
|
||||
print ("== |== [%d : %s]" % (idx, nnlib.device.getDeviceName(idx)) )
|
||||
|
||||
if not self.gpu_config.cpu_only and self.gpu_total_vram_gb == 2:
|
||||
if not self.device_config.cpu_only and self.device_config.gpu_total_vram_gb == 2:
|
||||
print ("==")
|
||||
print ("== WARNING: You are using 2GB GPU. Result quality may be significantly decreased.")
|
||||
print ("== If training does not start, close all programs and try again.")
|
||||
|
|
@ -168,18 +162,18 @@ class ModelBase(object):
|
|||
return ConverterBase(self, **in_options)
|
||||
|
||||
def to_multi_gpu_model_if_possible (self, models_list):
|
||||
if len(self.gpu_config.gpu_idxs) > 1:
|
||||
if len(self.device_config.gpu_idxs) > 1:
|
||||
#make batch_size to divide on GPU count without remainder
|
||||
self.batch_size = int( self.batch_size / len(self.gpu_config.gpu_idxs) )
|
||||
self.batch_size = int( self.batch_size / len(self.device_config.gpu_idxs) )
|
||||
if self.batch_size == 0:
|
||||
self.batch_size = 1
|
||||
self.batch_size *= len(self.gpu_config.gpu_idxs)
|
||||
self.batch_size *= len(self.device_config.gpu_idxs)
|
||||
|
||||
result = []
|
||||
for model in models_list:
|
||||
for i in range( len(model.output_names) ):
|
||||
model.output_names = 'output_%d' % (i)
|
||||
result += [ self.keras.utils.multi_gpu_model( model, self.gpu_config.gpu_idxs ) ]
|
||||
result += [ nnlib.keras.utils.multi_gpu_model( model, self.device_config.gpu_idxs ) ]
|
||||
|
||||
return result
|
||||
else:
|
||||
|
|
@ -259,12 +253,12 @@ class ModelBase(object):
|
|||
cv2.imwrite ( str (self.preview_history_path / ('%.6d.jpg' %( self.epoch) )), img )
|
||||
|
||||
self.epoch += 1
|
||||
|
||||
#............."Saving...
|
||||
|
||||
if epoch_time >= 10000:
|
||||
loss_string = "Training [#{0:06d}][{1:03d}s]".format ( self.epoch, epoch_time / 1000 )
|
||||
#............."Saving...
|
||||
loss_string = "Training [#{0:06d}][{1:.5s}s]".format ( self.epoch, '{:0.4f}'.format(epoch_time / 1000) )
|
||||
else:
|
||||
loss_string = "Training [#{0:06d}][{1:04d}ms]".format ( self.epoch, int(epoch_time*1000) % 10000 )
|
||||
loss_string = "Training [#{0:06d}][{1:04d}ms]".format ( self.epoch, int(epoch_time*1000) )
|
||||
for (loss_name, loss_value) in losses:
|
||||
loss_string += " %s:%.3f" % (loss_name, loss_value)
|
||||
|
||||
|
|
@ -274,7 +268,7 @@ class ModelBase(object):
|
|||
self.last_sample = self.generate_next_sample()
|
||||
|
||||
def finalize(self):
|
||||
gpufmkmgr.finalize_keras()
|
||||
nnlib.finalize_all()
|
||||
|
||||
def is_first_run(self):
|
||||
return self.epoch == 0
|
||||
|
|
@ -282,6 +276,12 @@ class ModelBase(object):
|
|||
def is_debug(self):
|
||||
return self.debug
|
||||
|
||||
def set_batch_size(self, batch_size):
|
||||
self.batch_size = batch_size
|
||||
|
||||
def get_batch_size(self):
|
||||
return self.batch_size
|
||||
|
||||
def get_epoch(self):
|
||||
return self.epoch
|
||||
|
||||
|
|
@ -301,16 +301,16 @@ class ModelBase(object):
|
|||
#example d = {2:2,3:4,4:8,5:16,6:32,7:32,8:32,9:48}
|
||||
keys = [x for x in d.keys()]
|
||||
|
||||
if self.gpu_config.cpu_only:
|
||||
if self.device_config.cpu_only:
|
||||
if self.batch_size == 0:
|
||||
self.batch_size = 2
|
||||
else:
|
||||
if self.gpu_total_vram_gb < keys[0]:
|
||||
if self.device_config.gpu_total_vram_gb < keys[0]:
|
||||
raise Exception ('Sorry, this model works only on %dGB+ GPU' % ( keys[0] ) )
|
||||
|
||||
if self.batch_size == 0:
|
||||
for x in keys:
|
||||
if self.gpu_total_vram_gb <= x:
|
||||
if self.device_config.gpu_total_vram_gb <= x:
|
||||
self.batch_size = d[x]
|
||||
break
|
||||
|
||||
|
|
|
|||
|
|
@ -2,10 +2,7 @@ import numpy as np
|
|||
import cv2
|
||||
from models import ModelBase
|
||||
from samples import *
|
||||
from nnlib import tf_dssim
|
||||
from nnlib import DSSIMLossClass
|
||||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from nnlib import nnlib
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
|
|
@ -17,9 +14,7 @@ class Model(ModelBase):
|
|||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
tf = self.tf
|
||||
keras = self.keras
|
||||
K = keras.backend
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
|
||||
self.set_vram_batch_requirements( {3.5:8,4:8,5:12,6:16,7:24,8:32,9:48} )
|
||||
if self.batch_size < 4:
|
||||
|
|
@ -34,39 +29,39 @@ class Model(ModelBase):
|
|||
self.encoder256.load_weights (self.get_strpath_storage_for_file(self.encoder256H5))
|
||||
self.decoder256.load_weights (self.get_strpath_storage_for_file(self.decoder256H5))
|
||||
|
||||
if self.is_training_mode:
|
||||
self.encoder64, self.decoder64_src, self.decoder64_dst, self.encoder256, self.decoder256 = self.to_multi_gpu_model_if_possible ( [self.encoder64, self.decoder64_src, self.decoder64_dst, self.encoder256, self.decoder256] )
|
||||
#if self.is_training_mode:
|
||||
# self.encoder64, self.decoder64_src, self.decoder64_dst, self.encoder256, self.decoder256 = self.to_multi_gpu_model_if_possible ( [self.encoder64, self.decoder64_src, self.decoder64_dst, self.encoder256, self.decoder256] )
|
||||
|
||||
input_A_warped64 = keras.layers.Input(img_shape64)
|
||||
input_B_warped64 = keras.layers.Input(img_shape64)
|
||||
input_A_warped64 = Input(img_shape64)
|
||||
input_B_warped64 = Input(img_shape64)
|
||||
A_rec64 = self.decoder64_src(self.encoder64(input_A_warped64))
|
||||
B_rec64 = self.decoder64_dst(self.encoder64(input_B_warped64))
|
||||
self.ae64 = self.keras.models.Model([input_A_warped64, input_B_warped64], [A_rec64, B_rec64] )
|
||||
self.ae64 = Model([input_A_warped64, input_B_warped64], [A_rec64, B_rec64] )
|
||||
|
||||
if self.is_training_mode:
|
||||
self.ae64, = self.to_multi_gpu_model_if_possible ( [self.ae64,] )
|
||||
|
||||
self.ae64.compile(optimizer=self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999),
|
||||
loss=[DSSIMLossClass(self.tf)(), DSSIMLossClass(self.tf)()] )
|
||||
self.ae64.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999),
|
||||
loss=[DSSIMLoss(), DSSIMLoss()] )
|
||||
|
||||
self.A64_view = K.function ([input_A_warped64], [A_rec64])
|
||||
self.B64_view = K.function ([input_B_warped64], [B_rec64])
|
||||
|
||||
input_A_warped64 = keras.layers.Input(img_shape64)
|
||||
input_A_target256 = keras.layers.Input(img_shape256)
|
||||
input_A_warped64 = Input(img_shape64)
|
||||
input_A_target256 = Input(img_shape256)
|
||||
A_rec256 = self.decoder256( self.encoder256(input_A_warped64) )
|
||||
|
||||
input_B_warped64 = keras.layers.Input(img_shape64)
|
||||
input_B_warped64 = Input(img_shape64)
|
||||
BA_rec64 = self.decoder64_src( self.encoder64(input_B_warped64) )
|
||||
BA_rec256 = self.decoder256( self.encoder256(BA_rec64) )
|
||||
|
||||
self.ae256 = self.keras.models.Model([input_A_warped64], [A_rec256] )
|
||||
self.ae256 = Model([input_A_warped64], [A_rec256] )
|
||||
|
||||
if self.is_training_mode:
|
||||
self.ae256, = self.to_multi_gpu_model_if_possible ( [self.ae256,] )
|
||||
|
||||
self.ae256.compile(optimizer=self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999),
|
||||
loss=[DSSIMLossClass(self.tf)()])
|
||||
self.ae256.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999),
|
||||
loss=[DSSIMLoss()])
|
||||
|
||||
self.A256_view = K.function ([input_A_warped64], [A_rec256])
|
||||
self.BA256_view = K.function ([input_B_warped64], [BA_rec256])
|
||||
|
|
@ -153,62 +148,67 @@ class Model(ModelBase):
|
|||
return ConverterAvatar(self.predictor_func, predictor_input_size=64, output_size=256, **in_options)
|
||||
|
||||
def Build(self):
|
||||
keras, K = self.keras, self.keras.backend
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
||||
img_shape64 = (64,64,3)
|
||||
img_shape256 = (256,256,3)
|
||||
|
||||
def upscale (dim):
|
||||
def func(x):
|
||||
return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
|
||||
return func
|
||||
|
||||
def Encoder(_input):
|
||||
x = _input
|
||||
x = self.keras.layers.convolutional.Conv2D(90, kernel_size=5, strides=1, padding='same')(x)
|
||||
x = self.keras.layers.convolutional.Conv2D(90, kernel_size=5, strides=1, padding='same')(x)
|
||||
x = self.keras.layers.MaxPooling2D(pool_size=(3, 3), strides=2, padding='same')(x)
|
||||
x = Conv2D(90, kernel_size=5, strides=1, padding='same')(x)
|
||||
x = Conv2D(90, kernel_size=5, strides=1, padding='same')(x)
|
||||
x = MaxPooling2D(pool_size=(3, 3), strides=2, padding='same')(x)
|
||||
|
||||
x = self.keras.layers.convolutional.Conv2D(180, kernel_size=3, strides=1, padding='same')(x)
|
||||
x = self.keras.layers.convolutional.Conv2D(180, kernel_size=3, strides=1, padding='same')(x)
|
||||
x = self.keras.layers.MaxPooling2D(pool_size=(3, 3), strides=2, padding='same')(x)
|
||||
x = Conv2D(180, kernel_size=3, strides=1, padding='same')(x)
|
||||
x = Conv2D(180, kernel_size=3, strides=1, padding='same')(x)
|
||||
x = MaxPooling2D(pool_size=(3, 3), strides=2, padding='same')(x)
|
||||
|
||||
x = self.keras.layers.convolutional.Conv2D(360, kernel_size=3, strides=1, padding='same')(x)
|
||||
x = self.keras.layers.convolutional.Conv2D(360, kernel_size=3, strides=1, padding='same')(x)
|
||||
x = self.keras.layers.MaxPooling2D(pool_size=(3, 3), strides=2, padding='same')(x)
|
||||
x = Conv2D(360, kernel_size=3, strides=1, padding='same')(x)
|
||||
x = Conv2D(360, kernel_size=3, strides=1, padding='same')(x)
|
||||
x = MaxPooling2D(pool_size=(3, 3), strides=2, padding='same')(x)
|
||||
|
||||
x = self.keras.layers.Dense (1024)(x)
|
||||
x = self.keras.layers.advanced_activations.LeakyReLU(0.1)(x)
|
||||
x = self.keras.layers.Dropout(0.5)(x)
|
||||
x = Dense (1024)(x)
|
||||
x = LeakyReLU(0.1)(x)
|
||||
x = Dropout(0.5)(x)
|
||||
|
||||
x = self.keras.layers.Dense (1024)(x)
|
||||
x = self.keras.layers.advanced_activations.LeakyReLU(0.1)(x)
|
||||
x = self.keras.layers.Dropout(0.5)(x)
|
||||
x = self.keras.layers.Flatten()(x)
|
||||
x = self.keras.layers.Dense (64)(x)
|
||||
x = Dense (1024)(x)
|
||||
x = LeakyReLU(0.1)(x)
|
||||
x = Dropout(0.5)(x)
|
||||
x = Flatten()(x)
|
||||
x = Dense (64)(x)
|
||||
|
||||
return keras.models.Model (_input, x)
|
||||
|
||||
encoder256 = Encoder( keras.layers.Input (img_shape64) )
|
||||
encoder64 = Encoder( keras.layers.Input (img_shape64) )
|
||||
encoder256 = Encoder( Input (img_shape64) )
|
||||
encoder64 = Encoder( Input (img_shape64) )
|
||||
|
||||
def decoder256(encoder):
|
||||
decoder_input = keras.layers.Input ( K.int_shape(encoder.outputs[0])[1:] )
|
||||
decoder_input = Input ( K.int_shape(encoder.outputs[0])[1:] )
|
||||
x = decoder_input
|
||||
x = self.keras.layers.Dense(16 * 16 * 720)(x)
|
||||
x = keras.layers.Reshape ( (16, 16, 720) )(x)
|
||||
x = upscale(keras, x, 720)
|
||||
x = upscale(keras, x, 360)
|
||||
x = upscale(keras, x, 180)
|
||||
x = upscale(keras, x, 90)
|
||||
x = keras.layers.convolutional.Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
x = Dense(16 * 16 * 720)(x)
|
||||
x = Reshape ( (16, 16, 720) )(x)
|
||||
x = upscale(720)(x)
|
||||
x = upscale(360)(x)
|
||||
x = upscale(180)(x)
|
||||
x = upscale(90)(x)
|
||||
x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
return keras.models.Model(decoder_input, x)
|
||||
|
||||
def decoder64(encoder):
|
||||
decoder_input = keras.layers.Input ( K.int_shape(encoder.outputs[0])[1:] )
|
||||
decoder_input = Input ( K.int_shape(encoder.outputs[0])[1:] )
|
||||
x = decoder_input
|
||||
x = self.keras.layers.Dense(8 * 8 * 720)(x)
|
||||
x = keras.layers.Reshape ( (8, 8, 720) )(x)
|
||||
x = upscale(keras, x, 360)
|
||||
x = upscale(keras, x, 180)
|
||||
x = upscale(keras, x, 90)
|
||||
x = keras.layers.convolutional.Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
return keras.models.Model(decoder_input, x)
|
||||
x = Dense(8 * 8 * 720)(x)
|
||||
x = Reshape ( (8, 8, 720) )(x)
|
||||
x = upscale(360)(x)
|
||||
x = upscale(180)(x)
|
||||
x = upscale(90)(x)
|
||||
x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
return Model(decoder_input, x)
|
||||
|
||||
return img_shape64, img_shape256, encoder64, decoder64(encoder64), decoder64(encoder64), encoder256, decoder256(encoder256)
|
||||
|
||||
|
|
@ -230,7 +230,7 @@ class ConverterAvatar(ConverterBase):
|
|||
|
||||
#override
|
||||
def get_mode(self):
|
||||
return ConverterBase.MODE_IMAGE
|
||||
return ConverterBase.MODE_IMAGE_WITH_LANDMARKS
|
||||
|
||||
#override
|
||||
def dummy_predict(self):
|
||||
|
|
|
|||
|
|
@ -1,12 +1,10 @@
|
|||
from models import ModelBase
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
from nnlib import DSSIMMaskLossClass
|
||||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from nnlib import nnlib
|
||||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
encoderH5 = 'encoder.h5'
|
||||
|
|
@ -15,30 +13,27 @@ class Model(ModelBase):
|
|||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {4.5:16,5:16,6:16,7:16,8:24,9:24,10:32,11:32,12:32,13:48} )
|
||||
|
||||
ae_input_layer = self.keras.layers.Input(shape=(128, 128, 3))
|
||||
mask_layer = self.keras.layers.Input(shape=(128, 128, 1)) #same as output
|
||||
|
||||
self.encoder = self.Encoder(ae_input_layer)
|
||||
self.decoder_src = self.Decoder()
|
||||
self.decoder_dst = self.Decoder()
|
||||
ae_input_layer = Input(shape=(128, 128, 3))
|
||||
mask_layer = Input(shape=(128, 128, 1)) #same as output
|
||||
|
||||
self.encoder, self.decoder_src, self.decoder_dst = self.Build(ae_input_layer)
|
||||
|
||||
if not self.is_first_run():
|
||||
self.encoder.load_weights (self.get_strpath_storage_for_file(self.encoderH5))
|
||||
self.decoder_src.load_weights (self.get_strpath_storage_for_file(self.decoder_srcH5))
|
||||
self.decoder_dst.load_weights (self.get_strpath_storage_for_file(self.decoder_dstH5))
|
||||
|
||||
self.autoencoder_src = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder_src(self.encoder(ae_input_layer)))
|
||||
self.autoencoder_dst = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder_dst(self.encoder(ae_input_layer)))
|
||||
self.autoencoder_src = Model([ae_input_layer,mask_layer], self.decoder_src(self.encoder(ae_input_layer)))
|
||||
self.autoencoder_dst = Model([ae_input_layer,mask_layer], self.decoder_dst(self.encoder(ae_input_layer)))
|
||||
|
||||
if self.is_training_mode:
|
||||
self.autoencoder_src, self.autoencoder_dst = self.to_multi_gpu_model_if_possible ( [self.autoencoder_src, self.autoencoder_dst] )
|
||||
|
||||
optimizer = self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999)
|
||||
dssimloss = DSSIMMaskLossClass(self.tf)([mask_layer])
|
||||
self.autoencoder_src.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] )
|
||||
|
||||
if self.is_training_mode:
|
||||
f = SampleProcessor.TypeFlags
|
||||
|
|
@ -123,33 +118,48 @@ class Model(ModelBase):
|
|||
|
||||
return ConverterMasked(self.predictor_func, predictor_input_size=128, output_size=128, face_type=FaceType.FULL, clip_border_mask_per=0.046875, **in_options)
|
||||
|
||||
def Encoder(self, input_layer):
|
||||
x = input_layer
|
||||
x = conv(self.keras, x, 128)
|
||||
x = conv(self.keras, x, 256)
|
||||
x = conv(self.keras, x, 512)
|
||||
x = conv(self.keras, x, 1024)
|
||||
|
||||
x = self.keras.layers.Dense(512)(self.keras.layers.Flatten()(x))
|
||||
x = self.keras.layers.Dense(8 * 8 * 512)(x)
|
||||
x = self.keras.layers.Reshape((8, 8, 512))(x)
|
||||
x = upscale(self.keras, x, 512)
|
||||
def Build(self, input_layer):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
||||
def downscale (dim):
|
||||
def func(x):
|
||||
return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x))
|
||||
return func
|
||||
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
|
||||
def Decoder(self):
|
||||
input_ = self.keras.layers.Input(shape=(16, 16, 512))
|
||||
x = input_
|
||||
x = upscale(self.keras, x, 512)
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = upscale(self.keras, x, 128)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(self.keras, y, 512)
|
||||
y = upscale(self.keras, y, 256)
|
||||
y = upscale(self.keras, y, 128)
|
||||
def upscale (dim):
|
||||
def func(x):
|
||||
return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
|
||||
return func
|
||||
|
||||
x = self.keras.layers.convolutional.Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
y = self.keras.layers.convolutional.Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
def Encoder(input_layer):
|
||||
x = input_layer
|
||||
x = downscale(128)(x)
|
||||
x = downscale(256)(x)
|
||||
x = downscale(512)(x)
|
||||
x = downscale(1024)(x)
|
||||
|
||||
x = Dense(512)(Flatten()(x))
|
||||
x = Dense(8 * 8 * 512)(x)
|
||||
x = Reshape((8, 8, 512))(x)
|
||||
x = upscale(512)(x)
|
||||
|
||||
return Model(input_layer, x)
|
||||
|
||||
def Decoder():
|
||||
input_ = Input(shape=(16, 16, 512))
|
||||
x = input_
|
||||
x = upscale(512)(x)
|
||||
x = upscale(256)(x)
|
||||
x = upscale(128)(x)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(512)(y)
|
||||
y = upscale(256)(y)
|
||||
y = upscale(128)(y)
|
||||
|
||||
x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
|
||||
return Model(input_, [x,y])
|
||||
|
||||
return self.keras.models.Model(input_, [x,y])
|
||||
return Encoder(input_layer), Decoder(), Decoder()
|
||||
|
|
@ -1,9 +1,6 @@
|
|||
import numpy as np
|
||||
import cv2
|
||||
|
||||
from nnlib import DSSIMMaskLossClass
|
||||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from nnlib import nnlib
|
||||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
|
|
@ -16,10 +13,8 @@ class Model(ModelBase):
|
|||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
tf = self.tf
|
||||
keras = self.keras
|
||||
K = keras.backend
|
||||
self.set_vram_batch_requirements( {2.5:2,3:2,4:2,4:4,5:8,6:8,7:16,8:16,9:24,10:24,11:32,12:32,13:48} )
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {2.5:2,3:2,4:2,4:4,5:8,6:12,7:16,8:16,9:24,10:24,11:32,12:32,13:48} )
|
||||
|
||||
bgr_shape, mask_shape, self.encoder, self.decoder_src, self.decoder_dst = self.Build(self.created_vram_gb)
|
||||
if not self.is_first_run():
|
||||
|
|
@ -27,21 +22,21 @@ class Model(ModelBase):
|
|||
self.decoder_src.load_weights (self.get_strpath_storage_for_file(self.decoder_srcH5))
|
||||
self.decoder_dst.load_weights (self.get_strpath_storage_for_file(self.decoder_dstH5))
|
||||
|
||||
input_src_bgr = self.keras.layers.Input(bgr_shape)
|
||||
input_src_mask = self.keras.layers.Input(mask_shape)
|
||||
input_dst_bgr = self.keras.layers.Input(bgr_shape)
|
||||
input_dst_mask = self.keras.layers.Input(mask_shape)
|
||||
input_src_bgr = Input(bgr_shape)
|
||||
input_src_mask = Input(mask_shape)
|
||||
input_dst_bgr = Input(bgr_shape)
|
||||
input_dst_mask = Input(mask_shape)
|
||||
|
||||
rec_src_bgr, rec_src_mask = self.decoder_src( self.encoder(input_src_bgr) )
|
||||
rec_dst_bgr, rec_dst_mask = self.decoder_dst( self.encoder(input_dst_bgr) )
|
||||
|
||||
self.ae = self.keras.models.Model([input_src_bgr,input_src_mask,input_dst_bgr,input_dst_mask], [rec_src_bgr, rec_src_mask, rec_dst_bgr, rec_dst_mask] )
|
||||
self.ae = Model([input_src_bgr,input_src_mask,input_dst_bgr,input_dst_mask], [rec_src_bgr, rec_src_mask, rec_dst_bgr, rec_dst_mask] )
|
||||
|
||||
if self.is_training_mode:
|
||||
self.ae, = self.to_multi_gpu_model_if_possible ( [self.ae,] )
|
||||
|
||||
self.ae.compile(optimizer=self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999),
|
||||
loss=[ DSSIMMaskLossClass(self.tf)([input_src_mask]), 'mae', DSSIMMaskLossClass(self.tf)([input_dst_mask]), 'mae' ] )
|
||||
self.ae.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999),
|
||||
loss=[ DSSIMMaskLoss([input_src_mask]), 'mae', DSSIMMaskLoss([input_dst_mask]), 'mae' ] )
|
||||
|
||||
self.src_view = K.function([input_src_bgr],[rec_src_bgr, rec_src_mask])
|
||||
self.dst_view = K.function([input_dst_bgr],[rec_dst_bgr, rec_dst_mask])
|
||||
|
|
@ -129,61 +124,73 @@ class Model(ModelBase):
|
|||
return ConverterMasked(self.predictor_func, predictor_input_size=128, output_size=128, face_type=FaceType.HALF, **in_options)
|
||||
|
||||
def Build(self, created_vram_gb):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
||||
bgr_shape = (128, 128, 3)
|
||||
mask_shape = (128, 128, 1)
|
||||
|
||||
def downscale (dim):
|
||||
def func(x):
|
||||
return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x))
|
||||
return func
|
||||
|
||||
def upscale (dim):
|
||||
def func(x):
|
||||
return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
|
||||
return func
|
||||
|
||||
def Encoder(input_shape):
|
||||
input_layer = self.keras.layers.Input(input_shape)
|
||||
input_layer = Input(input_shape)
|
||||
x = input_layer
|
||||
if created_vram_gb >= 5:
|
||||
x = conv(self.keras, x, 128)
|
||||
x = conv(self.keras, x, 256)
|
||||
x = conv(self.keras, x, 512)
|
||||
x = conv(self.keras, x, 1024)
|
||||
x = self.keras.layers.Dense(512)(self.keras.layers.Flatten()(x))
|
||||
x = self.keras.layers.Dense(8 * 8 * 512)(x)
|
||||
x = self.keras.layers.Reshape((8, 8, 512))(x)
|
||||
x = upscale(self.keras, x, 512)
|
||||
x = downscale(128)(x)
|
||||
x = downscale(256)(x)
|
||||
x = downscale(512)(x)
|
||||
x = downscale(1024)(x)
|
||||
x = Dense(512)(Flatten()(x))
|
||||
x = Dense(8 * 8 * 512)(x)
|
||||
x = Reshape((8, 8, 512))(x)
|
||||
x = upscale(512)(x)
|
||||
else:
|
||||
x = conv(self.keras, x, 128)
|
||||
x = conv(self.keras, x, 256)
|
||||
x = conv(self.keras, x, 512)
|
||||
x = conv(self.keras, x, 1024)
|
||||
x = self.keras.layers.Dense(256)(self.keras.layers.Flatten()(x))
|
||||
x = self.keras.layers.Dense(8 * 8 * 256)(x)
|
||||
x = self.keras.layers.Reshape((8, 8, 256))(x)
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = downscale(128)(x)
|
||||
x = downscale(256)(x)
|
||||
x = downscale(512)(x)
|
||||
x = downscale(1024)(x)
|
||||
x = Dense(256)(Flatten()(x))
|
||||
x = Dense(8 * 8 * 256)(x)
|
||||
x = Reshape((8, 8, 256))(x)
|
||||
x = upscale(256)(x)
|
||||
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
return Model(input_layer, x)
|
||||
|
||||
def Decoder():
|
||||
if created_vram_gb >= 5:
|
||||
input_ = self.keras.layers.Input(shape=(16, 16, 512))
|
||||
input_ = Input(shape=(16, 16, 512))
|
||||
x = input_
|
||||
x = upscale(self.keras, x, 512)
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = upscale(self.keras, x, 128)
|
||||
x = upscale(512)(x)
|
||||
x = upscale(256)(x)
|
||||
x = upscale(128)(x)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(self.keras, y, 512)
|
||||
y = upscale(self.keras, y, 256)
|
||||
y = upscale(self.keras, y, 128)
|
||||
y = upscale(512)(y)
|
||||
y = upscale(256)(y)
|
||||
y = upscale(128)(y)
|
||||
else:
|
||||
input_ = self.keras.layers.Input(shape=(16, 16, 256))
|
||||
x = input_
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = upscale(self.keras, x, 128)
|
||||
x = upscale(self.keras, x, 64)
|
||||
x = upscale(256)(x)
|
||||
x = upscale(128)(x)
|
||||
x = upscale(64)(x)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(self.keras, y, 256)
|
||||
y = upscale(self.keras, y, 128)
|
||||
y = upscale(self.keras, y, 64)
|
||||
y = upscale(256)(y)
|
||||
y = upscale(128)(y)
|
||||
y = upscale(64)(y)
|
||||
|
||||
x = self.keras.layers.convolutional.Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
y = self.keras.layers.convolutional.Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
|
||||
|
||||
return self.keras.models.Model(input_, [x,y])
|
||||
return Model(input_, [x,y])
|
||||
|
||||
return bgr_shape, mask_shape, Encoder(bgr_shape), Decoder(), Decoder()
|
||||
|
|
|
|||
|
|
@ -1,11 +1,9 @@
|
|||
from models import ModelBase
|
||||
import numpy as np
|
||||
from samples import *
|
||||
|
||||
from nnlib import DSSIMMaskLossClass
|
||||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from nnlib import nnlib
|
||||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
|
|
@ -15,9 +13,7 @@ class Model(ModelBase):
|
|||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
tf = self.tf
|
||||
keras = self.keras
|
||||
K = keras.backend
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {1.5:2,2:2,3:8,4:16,5:24,6:32,7:40,8:48} )
|
||||
|
||||
bgr_shape, mask_shape, self.encoder, self.decoder_src, self.decoder_dst = self.Build(self.created_vram_gb)
|
||||
|
|
@ -26,21 +22,21 @@ class Model(ModelBase):
|
|||
self.decoder_src.load_weights (self.get_strpath_storage_for_file(self.decoder_srcH5))
|
||||
self.decoder_dst.load_weights (self.get_strpath_storage_for_file(self.decoder_dstH5))
|
||||
|
||||
input_src_bgr = self.keras.layers.Input(bgr_shape)
|
||||
input_src_mask = self.keras.layers.Input(mask_shape)
|
||||
input_dst_bgr = self.keras.layers.Input(bgr_shape)
|
||||
input_dst_mask = self.keras.layers.Input(mask_shape)
|
||||
input_src_bgr = Input(bgr_shape)
|
||||
input_src_mask = Input(mask_shape)
|
||||
input_dst_bgr = Input(bgr_shape)
|
||||
input_dst_mask = Input(mask_shape)
|
||||
|
||||
rec_src_bgr, rec_src_mask = self.decoder_src( self.encoder(input_src_bgr) )
|
||||
rec_dst_bgr, rec_dst_mask = self.decoder_dst( self.encoder(input_dst_bgr) )
|
||||
|
||||
self.ae = self.keras.models.Model([input_src_bgr,input_src_mask,input_dst_bgr,input_dst_mask], [rec_src_bgr, rec_src_mask, rec_dst_bgr, rec_dst_mask] )
|
||||
self.ae = Model([input_src_bgr,input_src_mask,input_dst_bgr,input_dst_mask], [rec_src_bgr, rec_src_mask, rec_dst_bgr, rec_dst_mask] )
|
||||
|
||||
if self.is_training_mode:
|
||||
self.ae, = self.to_multi_gpu_model_if_possible ( [self.ae,] )
|
||||
|
||||
self.ae.compile(optimizer=self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999),
|
||||
loss=[ DSSIMMaskLossClass(self.tf)([input_src_mask]), 'mae', DSSIMMaskLossClass(self.tf)([input_dst_mask]), 'mae' ] )
|
||||
self.ae.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999),
|
||||
loss=[ DSSIMMaskLoss([input_src_mask]), 'mae', DSSIMMaskLoss([input_dst_mask]), 'mae' ] )
|
||||
|
||||
self.src_view = K.function([input_src_bgr],[rec_src_bgr, rec_src_mask])
|
||||
self.dst_view = K.function([input_dst_bgr],[rec_dst_bgr, rec_dst_mask])
|
||||
|
|
@ -130,58 +126,69 @@ class Model(ModelBase):
|
|||
return ConverterMasked(self.predictor_func, predictor_input_size=64, output_size=64, face_type=FaceType.HALF, **in_options)
|
||||
|
||||
def Build(self, created_vram_gb):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
||||
bgr_shape = (64, 64, 3)
|
||||
mask_shape = (64, 64, 1)
|
||||
|
||||
def downscale (dim):
|
||||
def func(x):
|
||||
return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x))
|
||||
return func
|
||||
|
||||
def upscale (dim):
|
||||
def func(x):
|
||||
return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
|
||||
return func
|
||||
|
||||
def Encoder(input_shape):
|
||||
input_layer = self.keras.layers.Input(input_shape)
|
||||
input_layer = Input(input_shape)
|
||||
x = input_layer
|
||||
if created_vram_gb >= 4:
|
||||
x = conv(self.keras, x, 128)
|
||||
x = conv(self.keras, x, 256)
|
||||
x = conv(self.keras, x, 512)
|
||||
x = conv(self.keras, x, 1024)
|
||||
x = self.keras.layers.Dense(1024)(self.keras.layers.Flatten()(x))
|
||||
x = self.keras.layers.Dense(4 * 4 * 1024)(x)
|
||||
x = self.keras.layers.Reshape((4, 4, 1024))(x)
|
||||
x = upscale(self.keras, x, 512)
|
||||
x = downscale(128)(x)
|
||||
x = downscale(256)(x)
|
||||
x = downscale(512)(x)
|
||||
x = downscale(1024)(x)
|
||||
x = Dense(1024)(Flatten()(x))
|
||||
x = Dense(4 * 4 * 1024)(x)
|
||||
x = Reshape((4, 4, 1024))(x)
|
||||
x = upscale(512)(x)
|
||||
else:
|
||||
x = conv(self.keras, x, 128 )
|
||||
x = conv(self.keras, x, 256 )
|
||||
x = conv(self.keras, x, 512 )
|
||||
x = conv(self.keras, x, 768 )
|
||||
x = self.keras.layers.Dense(512)(self.keras.layers.Flatten()(x))
|
||||
x = self.keras.layers.Dense(4 * 4 * 512)(x)
|
||||
x = self.keras.layers.Reshape((4, 4, 512))(x)
|
||||
x = upscale(self.keras, x, 256)
|
||||
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
x = downscale(128)(x)
|
||||
x = downscale(256)(x)
|
||||
x = downscale(512)(x)
|
||||
x = downscale(768)(x)
|
||||
x = Dense(512)(Flatten()(x))
|
||||
x = Dense(4 * 4 * 512)(x)
|
||||
x = Reshape((4, 4, 512))(x)
|
||||
x = upscale(256)(x)
|
||||
return Model(input_layer, x)
|
||||
|
||||
def Decoder():
|
||||
if created_vram_gb >= 4:
|
||||
input_ = self.keras.layers.Input(shape=(8, 8, 512))
|
||||
input_ = Input(shape=(8, 8, 512))
|
||||
x = input_
|
||||
x = upscale(self.keras, x, 512)
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = upscale(self.keras, x, 128)
|
||||
|
||||
else:
|
||||
input_ = self.keras.layers.Input(shape=(8, 8, 256))
|
||||
x = upscale(512)(x)
|
||||
x = upscale(256)(x)
|
||||
x = upscale(128)(x)
|
||||
|
||||
x = input_
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = upscale(self.keras, x, 128)
|
||||
x = upscale(self.keras, x, 64)
|
||||
else:
|
||||
input_ = Input(shape=(8, 8, 256))
|
||||
|
||||
x = input_
|
||||
x = upscale(256)(x)
|
||||
x = upscale(128)(x)
|
||||
x = upscale(64)(x)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(self.keras, y, 256)
|
||||
y = upscale(self.keras, y, 128)
|
||||
y = upscale(self.keras, y, 64)
|
||||
y = upscale(256)(y)
|
||||
y = upscale(128)(y)
|
||||
y = upscale(64)(y)
|
||||
|
||||
x = self.keras.layers.convolutional.Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
y = self.keras.layers.convolutional.Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
|
||||
|
||||
return self.keras.models.Model(input_, [x,y])
|
||||
return Model(input_, [x,y])
|
||||
|
||||
return bgr_shape, mask_shape, Encoder(bgr_shape), Decoder(), Decoder()
|
||||
|
|
@ -1,10 +1,7 @@
|
|||
from models import ModelBase
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
from nnlib import DSSIMMaskLossClass
|
||||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from nnlib import nnlib
|
||||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
|
||||
|
|
@ -17,16 +14,14 @@ class Model(ModelBase):
|
|||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {4.5:4,5:4,6:8,7:12,8:16,9:20,10:24,11:24,12:32,13:48} )
|
||||
|
||||
ae_input_layer = self.keras.layers.Input(shape=(128, 128, 3))
|
||||
mask_layer = self.keras.layers.Input(shape=(128, 128, 1)) #same as output
|
||||
ae_input_layer = Input(shape=(128, 128, 3))
|
||||
mask_layer = Input(shape=(128, 128, 1)) #same as output
|
||||
|
||||
self.encoder = self.Encoder(ae_input_layer)
|
||||
self.decoder = self.Decoder()
|
||||
self.inter_B = self.Intermediate ()
|
||||
self.inter_AB = self.Intermediate ()
|
||||
|
||||
self.encoder, self.decoder, self.inter_B, self.inter_AB = self.Build(ae_input_layer)
|
||||
|
||||
if not self.is_first_run():
|
||||
self.encoder.load_weights (self.get_strpath_storage_for_file(self.encoderH5))
|
||||
self.decoder.load_weights (self.get_strpath_storage_for_file(self.decoderH5))
|
||||
|
|
@ -36,16 +31,14 @@ class Model(ModelBase):
|
|||
code = self.encoder(ae_input_layer)
|
||||
AB = self.inter_AB(code)
|
||||
B = self.inter_B(code)
|
||||
self.autoencoder_src = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder(self.keras.layers.Concatenate()([AB, AB])) )
|
||||
self.autoencoder_dst = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder(self.keras.layers.Concatenate()([B, AB])) )
|
||||
self.autoencoder_src = Model([ae_input_layer,mask_layer], self.decoder(Concatenate()([AB, AB])) )
|
||||
self.autoencoder_dst = Model([ae_input_layer,mask_layer], self.decoder(Concatenate()([B, AB])) )
|
||||
|
||||
if self.is_training_mode:
|
||||
self.autoencoder_src, self.autoencoder_dst = self.to_multi_gpu_model_if_possible ( [self.autoencoder_src, self.autoencoder_dst] )
|
||||
|
||||
optimizer = self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999)
|
||||
dssimloss = DSSIMMaskLossClass(self.tf)([mask_layer])
|
||||
self.autoencoder_src.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
|
||||
self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] )
|
||||
|
||||
if self.is_training_mode:
|
||||
f = SampleProcessor.TypeFlags
|
||||
|
|
@ -131,37 +124,52 @@ class Model(ModelBase):
|
|||
in_options['blur_mask_modifier'] = 0
|
||||
|
||||
return ConverterMasked(self.predictor_func, predictor_input_size=128, output_size=128, face_type=FaceType.FULL, clip_border_mask_per=0.046875, **in_options)
|
||||
|
||||
def Build(self, input_layer):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
||||
def Encoder(self, input_layer,):
|
||||
x = input_layer
|
||||
x = conv(self.keras, x, 128)
|
||||
x = conv(self.keras, x, 256)
|
||||
x = conv(self.keras, x, 512)
|
||||
x = conv(self.keras, x, 1024)
|
||||
x = self.keras.layers.Flatten()(x)
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
def downscale (dim):
|
||||
def func(x):
|
||||
return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x))
|
||||
return func
|
||||
|
||||
def upscale (dim):
|
||||
def func(x):
|
||||
return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
|
||||
return func
|
||||
|
||||
def Encoder():
|
||||
x = input_layer
|
||||
x = downscale(128)(x)
|
||||
x = downscale(256)(x)
|
||||
x = downscale(512)(x)
|
||||
x = downscale(1024)(x)
|
||||
x = Flatten()(x)
|
||||
return Model(input_layer, x)
|
||||
|
||||
def Intermediate(self):
|
||||
input_layer = self.keras.layers.Input(shape=(None, 8 * 8 * 1024))
|
||||
x = input_layer
|
||||
x = self.keras.layers.Dense(256)(x)
|
||||
x = self.keras.layers.Dense(8 * 8 * 512)(x)
|
||||
x = self.keras.layers.Reshape((8, 8, 512))(x)
|
||||
x = upscale(self.keras, x, 512)
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
def Intermediate():
|
||||
input_layer = Input(shape=(None, 8 * 8 * 1024))
|
||||
x = input_layer
|
||||
x = Dense(256)(x)
|
||||
x = Dense(8 * 8 * 512)(x)
|
||||
x = Reshape((8, 8, 512))(x)
|
||||
x = upscale(512)(x)
|
||||
return Model(input_layer, x)
|
||||
|
||||
def Decoder(self):
|
||||
input_ = self.keras.layers.Input(shape=(16, 16, 1024))
|
||||
x = input_
|
||||
x = upscale(self.keras, x, 512)
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = upscale(self.keras, x, 128)
|
||||
x = self.keras.layers.convolutional.Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(self.keras, y, 512)
|
||||
y = upscale(self.keras, y, 256)
|
||||
y = upscale(self.keras, y, 128)
|
||||
y = self.keras.layers.convolutional.Conv2D(1, kernel_size=5, padding='same', activation='sigmoid' )(y)
|
||||
|
||||
return self.keras.models.Model(input_, [x,y])
|
||||
def Decoder():
|
||||
input_ = Input(shape=(16, 16, 1024))
|
||||
x = input_
|
||||
x = upscale(512)(x)
|
||||
x = upscale(256)(x)
|
||||
x = upscale(128)(x)
|
||||
x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(512)(y)
|
||||
y = upscale(256)(y)
|
||||
y = upscale(128)(y)
|
||||
y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid' )(y)
|
||||
|
||||
return Model(input_, [x,y])
|
||||
|
||||
return Encoder(), Decoder(), Intermediate(), Intermediate()
|
||||
|
|
@ -1,13 +1,10 @@
|
|||
from models import ModelBase
|
||||
import numpy as np
|
||||
import cv2
|
||||
|
||||
from nnlib import DSSIMMaskLossClass
|
||||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from nnlib import nnlib
|
||||
from models import ModelBase
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
encoderH5 = 'encoder.h5'
|
||||
|
|
@ -17,16 +14,14 @@ class Model(ModelBase):
|
|||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {4.5:4,5:4,6:8,7:12,8:16,9:20,10:24,11:24,12:32,13:48} )
|
||||
|
||||
ae_input_layer = self.keras.layers.Input(shape=(128, 128, 3))
|
||||
mask_layer = self.keras.layers.Input(shape=(128, 128, 1)) #same as output
|
||||
ae_input_layer = Input(shape=(128, 128, 3))
|
||||
mask_layer = Input(shape=(128, 128, 1)) #same as output
|
||||
|
||||
self.encoder = self.Encoder(ae_input_layer)
|
||||
self.decoder = self.Decoder()
|
||||
self.inter_B = self.Intermediate ()
|
||||
self.inter_AB = self.Intermediate ()
|
||||
|
||||
self.encoder, self.decoder, self.inter_B, self.inter_AB = self.Build(ae_input_layer)
|
||||
|
||||
if not self.is_first_run():
|
||||
self.encoder.load_weights (self.get_strpath_storage_for_file(self.encoderH5))
|
||||
self.decoder.load_weights (self.get_strpath_storage_for_file(self.decoderH5))
|
||||
|
|
@ -36,16 +31,14 @@ class Model(ModelBase):
|
|||
code = self.encoder(ae_input_layer)
|
||||
AB = self.inter_AB(code)
|
||||
B = self.inter_B(code)
|
||||
self.autoencoder_src = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder(self.keras.layers.Concatenate()([AB, AB])) )
|
||||
self.autoencoder_dst = self.keras.models.Model([ae_input_layer,mask_layer], self.decoder(self.keras.layers.Concatenate()([B, AB])) )
|
||||
self.autoencoder_src = Model([ae_input_layer,mask_layer], self.decoder(Concatenate()([AB, AB])) )
|
||||
self.autoencoder_dst = Model([ae_input_layer,mask_layer], self.decoder(Concatenate()([B, AB])) )
|
||||
|
||||
if self.is_training_mode:
|
||||
self.autoencoder_src, self.autoencoder_dst = self.to_multi_gpu_model_if_possible ( [self.autoencoder_src, self.autoencoder_dst] )
|
||||
|
||||
optimizer = self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999)
|
||||
dssimloss = DSSIMMaskLossClass(self.tf)([mask_layer])
|
||||
self.autoencoder_src.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
|
||||
self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] )
|
||||
|
||||
if self.is_training_mode:
|
||||
f = SampleProcessor.TypeFlags
|
||||
|
|
@ -131,37 +124,52 @@ class Model(ModelBase):
|
|||
in_options['blur_mask_modifier'] = 0
|
||||
|
||||
return ConverterMasked(self.predictor_func, predictor_input_size=128, output_size=128, face_type=FaceType.FULL, clip_border_mask_per=0.046875, **in_options)
|
||||
|
||||
def Build(self, input_layer):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
||||
def Encoder(self, input_layer,):
|
||||
x = input_layer
|
||||
x = conv(self.keras, x, 128)
|
||||
x = conv(self.keras, x, 256)
|
||||
x = conv(self.keras, x, 512)
|
||||
x = conv(self.keras, x, 1024)
|
||||
x = self.keras.layers.Flatten()(x)
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
def downscale (dim):
|
||||
def func(x):
|
||||
return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x))
|
||||
return func
|
||||
|
||||
def upscale (dim):
|
||||
def func(x):
|
||||
return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
|
||||
return func
|
||||
|
||||
def Encoder():
|
||||
x = input_layer
|
||||
x = downscale(128)(x)
|
||||
x = downscale(256)(x)
|
||||
x = downscale(512)(x)
|
||||
x = downscale(1024)(x)
|
||||
x = Flatten()(x)
|
||||
return Model(input_layer, x)
|
||||
|
||||
def Intermediate(self):
|
||||
input_layer = self.keras.layers.Input(shape=(None, 8 * 8 * 1024))
|
||||
x = input_layer
|
||||
x = self.keras.layers.Dense(256)(x)
|
||||
x = self.keras.layers.Dense(8 * 8 * 512)(x)
|
||||
x = self.keras.layers.Reshape((8, 8, 512))(x)
|
||||
x = upscale(self.keras, x, 512)
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
def Intermediate():
|
||||
input_layer = Input(shape=(None, 8 * 8 * 1024))
|
||||
x = input_layer
|
||||
x = Dense(256)(x)
|
||||
x = Dense(8 * 8 * 512)(x)
|
||||
x = Reshape((8, 8, 512))(x)
|
||||
x = upscale(512)(x)
|
||||
return Model(input_layer, x)
|
||||
|
||||
def Decoder(self):
|
||||
input_ = self.keras.layers.Input(shape=(16, 16, 1024))
|
||||
x = input_
|
||||
x = upscale(self.keras, x, 512)
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = upscale(self.keras, x, 128)
|
||||
x = self.keras.layers.convolutional.Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(self.keras, y, 512)
|
||||
y = upscale(self.keras, y, 256)
|
||||
y = upscale(self.keras, y, 128)
|
||||
y = self.keras.layers.convolutional.Conv2D(1, kernel_size=5, padding='same', activation='sigmoid' )(y)
|
||||
|
||||
return self.keras.models.Model(input_, [x,y])
|
||||
def Decoder():
|
||||
input_ = Input(shape=(16, 16, 1024))
|
||||
x = input_
|
||||
x = upscale(512)(x)
|
||||
x = upscale(256)(x)
|
||||
x = upscale(128)(x)
|
||||
x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
|
||||
y = input_ #mask decoder
|
||||
y = upscale(512)(y)
|
||||
y = upscale(256)(y)
|
||||
y = upscale(128)(y)
|
||||
y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid' )(y)
|
||||
|
||||
return Model(input_, [x,y])
|
||||
|
||||
return Encoder(), Decoder(), Intermediate(), Intermediate()
|
||||
|
|
|
|||
|
|
@ -1,10 +1,7 @@
|
|||
import numpy as np
|
||||
import cv2
|
||||
|
||||
from nnlib import nnlib
|
||||
from models import ModelBase
|
||||
from nnlib import DSSIMMaskLossClass
|
||||
from nnlib import conv
|
||||
from nnlib import upscale
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
|
||||
|
|
@ -21,20 +18,15 @@ class Model(ModelBase):
|
|||
|
||||
#override
|
||||
def onInitialize(self, **in_options):
|
||||
exec(nnlib.import_all(), locals(), globals())
|
||||
self.set_vram_batch_requirements( {4.5:4,5:4,6:8,7:12,8:16,9:20,10:24,11:24,12:32,13:48} )
|
||||
|
||||
ae_input_layer = self.keras.layers.Input(shape=(128, 128, 3))
|
||||
mask_layer = self.keras.layers.Input(shape=(128, 128, 1)) #same as output
|
||||
ae_input_layer = Input(shape=(128, 128, 3))
|
||||
mask_layer = Input(shape=(128, 128, 1)) #same as output
|
||||
|
||||
self.encoder = self.Encoder(ae_input_layer)
|
||||
self.decoderMask = self.DecoderMask()
|
||||
self.decoderCommonA = self.DecoderCommon()
|
||||
self.decoderCommonB = self.DecoderCommon()
|
||||
self.decoderRGB = self.DecoderRGB()
|
||||
self.decoderBW = self.DecoderBW()
|
||||
self.inter_A = self.Intermediate ()
|
||||
self.inter_B = self.Intermediate ()
|
||||
|
||||
self.encoder, self.decoderMask, self.decoderCommonA, self.decoderCommonB, self.decoderRGB, \
|
||||
self.decoderBW, self.inter_A, self.inter_B = self.Build(ae_input_layer)
|
||||
|
||||
if not self.is_first_run():
|
||||
self.encoder.load_weights (self.get_strpath_storage_for_file(self.encoderH5))
|
||||
self.decoderMask.load_weights (self.get_strpath_storage_for_file(self.decoderMaskH5))
|
||||
|
|
@ -49,37 +41,35 @@ class Model(ModelBase):
|
|||
A = self.inter_A(code)
|
||||
B = self.inter_B(code)
|
||||
|
||||
inter_A_A = self.keras.layers.Concatenate()([A, A])
|
||||
inter_B_A = self.keras.layers.Concatenate()([B, A])
|
||||
inter_A_A = Concatenate()([A, A])
|
||||
inter_B_A = Concatenate()([B, A])
|
||||
|
||||
x1,m1 = self.decoderCommonA (inter_A_A)
|
||||
x2,m2 = self.decoderCommonA (inter_A_A)
|
||||
self.autoencoder_src = self.keras.models.Model([ae_input_layer,mask_layer],
|
||||
[ self.decoderBW (self.keras.layers.Concatenate()([x1,x2]) ),
|
||||
self.decoderMask(self.keras.layers.Concatenate()([m1,m2]) )
|
||||
self.autoencoder_src = Model([ae_input_layer,mask_layer],
|
||||
[ self.decoderBW (Concatenate()([x1,x2]) ),
|
||||
self.decoderMask(Concatenate()([m1,m2]) )
|
||||
])
|
||||
|
||||
x1,m1 = self.decoderCommonA (inter_A_A)
|
||||
x2,m2 = self.decoderCommonB (inter_A_A)
|
||||
self.autoencoder_src_RGB = self.keras.models.Model([ae_input_layer,mask_layer],
|
||||
[ self.decoderRGB (self.keras.layers.Concatenate()([x1,x2]) ),
|
||||
self.decoderMask (self.keras.layers.Concatenate()([m1,m2]) )
|
||||
self.autoencoder_src_RGB = Model([ae_input_layer,mask_layer],
|
||||
[ self.decoderRGB (Concatenate()([x1,x2]) ),
|
||||
self.decoderMask (Concatenate()([m1,m2]) )
|
||||
])
|
||||
|
||||
x1,m1 = self.decoderCommonA (inter_B_A)
|
||||
x2,m2 = self.decoderCommonB (inter_B_A)
|
||||
self.autoencoder_dst = self.keras.models.Model([ae_input_layer,mask_layer],
|
||||
[ self.decoderRGB (self.keras.layers.Concatenate()([x1,x2]) ),
|
||||
self.decoderMask (self.keras.layers.Concatenate()([m1,m2]) )
|
||||
self.autoencoder_dst = Model([ae_input_layer,mask_layer],
|
||||
[ self.decoderRGB (Concatenate()([x1,x2]) ),
|
||||
self.decoderMask (Concatenate()([m1,m2]) )
|
||||
])
|
||||
|
||||
if self.is_training_mode:
|
||||
self.autoencoder_src, self.autoencoder_dst = self.to_multi_gpu_model_if_possible ( [self.autoencoder_src, self.autoencoder_dst] )
|
||||
|
||||
optimizer = self.keras.optimizers.Adam(lr=5e-5, beta_1=0.5, beta_2=0.999)
|
||||
dssimloss = DSSIMMaskLossClass(self.tf)([mask_layer])
|
||||
self.autoencoder_src.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=optimizer, loss=[dssimloss, 'mse'] )
|
||||
|
||||
self.autoencoder_src.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] )
|
||||
self.autoencoder_dst.compile(optimizer=Adam(lr=5e-5, beta_1=0.5, beta_2=0.999), loss=[DSSIMMaskLoss([mask_layer]), 'mse'] )
|
||||
|
||||
if self.is_training_mode:
|
||||
f = SampleProcessor.TypeFlags
|
||||
|
|
@ -169,53 +159,67 @@ class Model(ModelBase):
|
|||
|
||||
return ConverterMasked(self.predictor_func, predictor_input_size=128, output_size=128, face_type=FaceType.FULL, clip_border_mask_per=0.046875, **in_options)
|
||||
|
||||
|
||||
def Encoder(self, input_layer,):
|
||||
x = input_layer
|
||||
x = conv(self.keras, x, 128)
|
||||
x = conv(self.keras, x, 256)
|
||||
x = conv(self.keras, x, 512)
|
||||
x = conv(self.keras, x, 1024)
|
||||
x = self.keras.layers.Flatten()(x)
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
|
||||
def Intermediate(self):
|
||||
input_layer = self.keras.layers.Input(shape=(None, 8 * 8 * 1024))
|
||||
x = input_layer
|
||||
x = self.keras.layers.Dense(256)(x)
|
||||
x = self.keras.layers.Dense(8 * 8 * 512)(x)
|
||||
x = self.keras.layers.Reshape((8, 8, 512))(x)
|
||||
x = upscale(self.keras, x, 512)
|
||||
return self.keras.models.Model(input_layer, x)
|
||||
|
||||
def DecoderCommon(self):
|
||||
input_ = self.keras.layers.Input(shape=(16, 16, 1024))
|
||||
x = input_
|
||||
x = upscale(self.keras, x, 512)
|
||||
x = upscale(self.keras, x, 256)
|
||||
x = upscale(self.keras, x, 128)
|
||||
def Build(self, input_layer):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
|
||||
y = input_
|
||||
y = upscale(self.keras, y, 256)
|
||||
y = upscale(self.keras, y, 128)
|
||||
y = upscale(self.keras, y, 64)
|
||||
|
||||
return self.keras.models.Model(input_, [x,y])
|
||||
|
||||
def DecoderRGB(self):
|
||||
input_ = self.keras.layers.Input(shape=(128, 128, 256))
|
||||
x = input_
|
||||
x = self.keras.layers.convolutional.Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
return self.keras.models.Model(input_, [x])
|
||||
def downscale (dim):
|
||||
def func(x):
|
||||
return LeakyReLU(0.1)(Conv2D(dim, 5, strides=2, padding='same')(x))
|
||||
return func
|
||||
|
||||
def upscale (dim):
|
||||
def func(x):
|
||||
return PixelShuffler()(LeakyReLU(0.1)(Conv2D(dim * 4, 3, strides=1, padding='same')(x)))
|
||||
return func
|
||||
|
||||
def Encoder():
|
||||
x = input_layer
|
||||
x = downscale(128)(x)
|
||||
x = downscale(256)(x)
|
||||
x = downscale(512)(x)
|
||||
x = downscale(1024)(x)
|
||||
x = Flatten()(x)
|
||||
return Model(input_layer, x)
|
||||
|
||||
def DecoderBW(self):
|
||||
input_ = self.keras.layers.Input(shape=(128, 128, 256))
|
||||
x = input_
|
||||
x = self.keras.layers.convolutional.Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
return self.keras.models.Model(input_, [x])
|
||||
|
||||
def DecoderMask(self):
|
||||
input_ = self.keras.layers.Input(shape=(128, 128, 128))
|
||||
y = input_
|
||||
y = self.keras.layers.convolutional.Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
return self.keras.models.Model(input_, [y])
|
||||
def Intermediate():
|
||||
input_layer = Input(shape=(None, 8 * 8 * 1024))
|
||||
x = input_layer
|
||||
x = Dense(256)(x)
|
||||
x = Dense(8 * 8 * 512)(x)
|
||||
x = Reshape((8, 8, 512))(x)
|
||||
x = upscale(512)(x)
|
||||
return Model(input_layer, x)
|
||||
|
||||
def DecoderCommon():
|
||||
input_ = Input(shape=(16, 16, 1024))
|
||||
x = input_
|
||||
x = upscale(512)(x)
|
||||
x = upscale(256)(x)
|
||||
x = upscale(128)(x)
|
||||
|
||||
y = input_
|
||||
y = upscale(256)(y)
|
||||
y = upscale(128)(y)
|
||||
y = upscale(64)(y)
|
||||
|
||||
return Model(input_, [x,y])
|
||||
|
||||
def DecoderRGB():
|
||||
input_ = Input(shape=(128, 128, 256))
|
||||
x = input_
|
||||
x = Conv2D(3, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
return Model(input_, [x])
|
||||
|
||||
def DecoderBW():
|
||||
input_ = Input(shape=(128, 128, 256))
|
||||
x = input_
|
||||
x = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(x)
|
||||
return Model(input_, [x])
|
||||
|
||||
def DecoderMask():
|
||||
input_ = Input(shape=(128, 128, 128))
|
||||
y = input_
|
||||
y = Conv2D(1, kernel_size=5, padding='same', activation='sigmoid')(y)
|
||||
return Model(input_, [y])
|
||||
|
||||
return Encoder(), DecoderMask(), DecoderCommon(), DecoderCommon(), DecoderRGB(), DecoderBW(), Intermediate(), Intermediate()
|
||||
243
models/Model_RecycleGAN/Model.py
Normal file
243
models/Model_RecycleGAN/Model.py
Normal file
|
|
@ -0,0 +1,243 @@
|
|||
from models import ModelBase
|
||||
import numpy as np
|
||||
import cv2
|
||||
from mathlib import get_power_of_two
|
||||
from nnlib import nnlib
|
||||
|
||||
|
||||
from facelib import FaceType
|
||||
from samples import *
|
||||
|
||||
class Model(ModelBase):
|
||||
|
||||
GAH5 = 'GA.h5'
|
||||
PAH5 = 'PA.h5'
|
||||
DAH5 = 'DA.h5'
|
||||
GBH5 = 'GB.h5'
|
||||
DBH5 = 'DB.h5'
|
||||
PBH5 = 'PB.h5'
|
||||
|
||||
#override
|
||||
def onInitialize(self, batch_size=-1, **in_options):
|
||||
exec(nnlib.code_import_all, locals(), globals())
|
||||
self.set_vram_batch_requirements( {6:6} )
|
||||
|
||||
|
||||
|
||||
created_batch_size = self.get_batch_size()
|
||||
if self.epoch == 0:
|
||||
#first run
|
||||
|
||||
print ("\nModel first run. Enter options.")
|
||||
|
||||
try:
|
||||
input_created_batch_size = int ( input ("Batch_size (default - based on VRAM) : ") )
|
||||
except:
|
||||
input_created_batch_size = 0
|
||||
|
||||
if input_created_batch_size != 0:
|
||||
created_batch_size = input_created_batch_size
|
||||
|
||||
self.options['created_batch_size'] = created_batch_size
|
||||
self.created_vram_gb = self.device_config.gpu_total_vram_gb
|
||||
else:
|
||||
#not first run
|
||||
if 'created_batch_size' in self.options.keys():
|
||||
created_batch_size = self.options['created_batch_size']
|
||||
else:
|
||||
raise Exception("Continue traning, but created_batch_size not found.")
|
||||
|
||||
resolution = 128
|
||||
bgr_shape = (resolution, resolution, 3)
|
||||
ngf = 64
|
||||
npf = 64
|
||||
ndf = 64
|
||||
lambda_A = 10
|
||||
lambda_B = 10
|
||||
|
||||
self.set_batch_size(created_batch_size)
|
||||
|
||||
use_batch_norm = created_batch_size > 1
|
||||
self.GA = modelify(ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=False))(Input(bgr_shape))
|
||||
self.GB = modelify(ResNet (bgr_shape[2], use_batch_norm, n_blocks=6, ngf=ngf, use_dropout=False))(Input(bgr_shape))
|
||||
#self.GA = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape))
|
||||
#self.GB = modelify(UNet (bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=ngf, use_dropout=True))(Input(bgr_shape))
|
||||
|
||||
self.PA = modelify(UNetTemporalPredictor(bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=npf, use_dropout=True))([Input(bgr_shape), Input(bgr_shape)])
|
||||
self.PB = modelify(UNetTemporalPredictor(bgr_shape[2], use_batch_norm, num_downs=get_power_of_two(resolution)-1, ngf=npf, use_dropout=True))([Input(bgr_shape), Input(bgr_shape)])
|
||||
|
||||
self.DA = modelify(NLayerDiscriminator(use_batch_norm, ndf=ndf, n_layers=3) ) (Input(bgr_shape))
|
||||
self.DB = modelify(NLayerDiscriminator(use_batch_norm, ndf=ndf, n_layers=3) ) (Input(bgr_shape))
|
||||
|
||||
if not self.is_first_run():
|
||||
self.GA.load_weights (self.get_strpath_storage_for_file(self.GAH5))
|
||||
self.DA.load_weights (self.get_strpath_storage_for_file(self.DAH5))
|
||||
self.PA.load_weights (self.get_strpath_storage_for_file(self.PAH5))
|
||||
self.GB.load_weights (self.get_strpath_storage_for_file(self.GBH5))
|
||||
self.DB.load_weights (self.get_strpath_storage_for_file(self.DBH5))
|
||||
self.PB.load_weights (self.get_strpath_storage_for_file(self.PBH5))
|
||||
|
||||
real_A0 = Input(bgr_shape, name="real_A0")
|
||||
real_A1 = Input(bgr_shape, name="real_A1")
|
||||
real_A2 = Input(bgr_shape, name="real_A2")
|
||||
|
||||
real_B0 = Input(bgr_shape, name="real_B0")
|
||||
real_B1 = Input(bgr_shape, name="real_B1")
|
||||
real_B2 = Input(bgr_shape, name="real_B2")
|
||||
|
||||
DA_ones = K.ones ( K.int_shape(self.DA.outputs[0])[1:] )
|
||||
DA_zeros = K.zeros ( K.int_shape(self.DA.outputs[0])[1:] )
|
||||
DB_ones = K.ones ( K.int_shape(self.DB.outputs[0])[1:] )
|
||||
DB_zeros = K.zeros ( K.int_shape(self.DB.outputs[0])[1:] )
|
||||
|
||||
def CycleLoss (t1,t2):
|
||||
return K.mean(K.square(t1 - t2))
|
||||
|
||||
def RecurrentLOSS(t1,t2):
|
||||
return K.mean(K.square(t1 - t2))
|
||||
|
||||
def RecycleLOSS(t1,t2):
|
||||
return K.mean(K.square(t1 - t2))
|
||||
|
||||
fake_B0 = self.GA(real_A0)
|
||||
fake_B1 = self.GA(real_A1)
|
||||
|
||||
fake_A0 = self.GB(real_B0)
|
||||
fake_A1 = self.GB(real_B1)
|
||||
|
||||
#rec_FB0 = self.GA(fake_A0)
|
||||
#rec_FB1 = self.GA(fake_A1)
|
||||
|
||||
#rec_FA0 = self.GB(fake_B0)
|
||||
#rec_FA1 = self.GB(fake_B1)
|
||||
|
||||
pred_A2 = self.PA ( [real_A0, real_A1])
|
||||
pred_B2 = self.PB ( [real_B0, real_B1])
|
||||
rec_A2 = self.GB ( self.PB ( [fake_B0, fake_B1]) )
|
||||
rec_B2 = self.GA ( self.PA ( [fake_A0, fake_A1]))
|
||||
|
||||
loss_G = K.mean(K.square(self.DB(fake_B0) - DB_ones)) + \
|
||||
K.mean(K.square(self.DB(fake_B1) - DB_ones)) + \
|
||||
K.mean(K.square(self.DA(fake_A0) - DA_ones)) + \
|
||||
K.mean(K.square(self.DA(fake_A1) - DA_ones)) + \
|
||||
lambda_A * ( #CycleLoss(rec_FA0, real_A0) + \
|
||||
#CycleLoss(rec_FA1, real_A1) + \
|
||||
RecurrentLOSS(pred_A2, real_A2) + \
|
||||
RecycleLOSS(rec_A2, real_A2) ) + \
|
||||
lambda_B * ( #CycleLoss(rec_FB0, real_B0) + \
|
||||
#CycleLoss(rec_FB1, real_B1) + \
|
||||
RecurrentLOSS(pred_B2, real_B2) + \
|
||||
RecycleLOSS(rec_B2, real_B2) )
|
||||
|
||||
weights_G = self.GA.trainable_weights + self.GB.trainable_weights + self.PA.trainable_weights + self.PB.trainable_weights
|
||||
|
||||
self.G_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_G],
|
||||
Adam(lr=2e-4, beta_1=0.5, beta_2=0.999).get_updates(loss_G, weights_G) )
|
||||
|
||||
###########
|
||||
|
||||
loss_D_A0 = ( K.mean(K.square( self.DA(real_A0) - DA_ones)) + \
|
||||
K.mean(K.square( self.DA(fake_A0) - DA_zeros)) ) * 0.5
|
||||
|
||||
loss_D_A1 = ( K.mean(K.square( self.DA(real_A1) - DA_ones)) + \
|
||||
K.mean(K.square( self.DA(fake_A1) - DA_zeros)) ) * 0.5
|
||||
|
||||
loss_D_A = loss_D_A0 + loss_D_A1
|
||||
|
||||
self.DA_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_D_A],
|
||||
Adam(lr=2e-4, beta_1=0.5, beta_2=0.999).get_updates(loss_D_A, self.DA.trainable_weights) )
|
||||
|
||||
############
|
||||
|
||||
loss_D_B0 = ( K.mean(K.square( self.DB(real_B0) - DB_ones)) + \
|
||||
K.mean(K.square( self.DB(fake_B0) - DB_zeros)) ) * 0.5
|
||||
|
||||
loss_D_B1 = ( K.mean(K.square( self.DB(real_B1) - DB_ones)) + \
|
||||
K.mean(K.square( self.DB(fake_B1) - DB_zeros)) ) * 0.5
|
||||
|
||||
loss_D_B = loss_D_B0 + loss_D_B1
|
||||
|
||||
self.DB_train = K.function ([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[loss_D_B],
|
||||
Adam(lr=2e-4, beta_1=0.5, beta_2=0.999).get_updates(loss_D_B, self.DB.trainable_weights) )
|
||||
|
||||
############
|
||||
|
||||
|
||||
self.G_view = K.function([real_A0, real_A1, real_A2, real_B0, real_B1, real_B2],[fake_A0, fake_A1, pred_A2, rec_A2, fake_B0, fake_B1, pred_B2, rec_B2 ])
|
||||
self.G_convert = K.function([real_B0],[fake_A0])
|
||||
|
||||
|
||||
if self.is_training_mode:
|
||||
f = SampleProcessor.TypeFlags
|
||||
self.set_training_data_generators ([
|
||||
SampleGeneratorImageTemporal(self.training_data_src_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
temporal_image_count=3,
|
||||
sample_process_options=SampleProcessor.Options(random_flip = False, normalize_tanh = True),
|
||||
output_sample_types=[ [f.SOURCE | f.MODE_BGR, resolution] ] ),
|
||||
|
||||
SampleGeneratorImageTemporal(self.training_data_dst_path, debug=self.is_debug(), batch_size=self.batch_size,
|
||||
temporal_image_count=3,
|
||||
sample_process_options=SampleProcessor.Options(random_flip = False, normalize_tanh = True),
|
||||
output_sample_types=[ [f.SOURCE | f.MODE_BGR, resolution] ] ),
|
||||
])
|
||||
|
||||
#import code
|
||||
#code.interact(local=dict(globals(), **locals()))
|
||||
self.supress_std_once = False
|
||||
|
||||
#override
|
||||
def onSave(self):
|
||||
self.save_weights_safe( [[self.GA, self.get_strpath_storage_for_file(self.GAH5)],
|
||||
[self.GB, self.get_strpath_storage_for_file(self.GBH5)],
|
||||
[self.DA, self.get_strpath_storage_for_file(self.DAH5)],
|
||||
[self.DB, self.get_strpath_storage_for_file(self.DBH5)],
|
||||
[self.PA, self.get_strpath_storage_for_file(self.PAH5)],
|
||||
[self.PB, self.get_strpath_storage_for_file(self.PBH5)] ])
|
||||
|
||||
#override
|
||||
def onTrainOneEpoch(self, sample):
|
||||
source_src_0, source_src_1, source_src_2, = sample[0]
|
||||
source_dst_0, source_dst_1, source_dst_2, = sample[1]
|
||||
|
||||
feed = [source_src_0, source_src_1, source_src_2, source_dst_0, source_dst_1, source_dst_2]
|
||||
|
||||
loss_G, = self.G_train ( feed )
|
||||
loss_DA, = self.DA_train( feed )
|
||||
loss_DB, = self.DB_train( feed )
|
||||
|
||||
return ( ('G', loss_G), ('DA', loss_DA), ('DB', loss_DB) )
|
||||
|
||||
#override
|
||||
def onGetPreview(self, sample):
|
||||
test_A0 = sample[0][0]
|
||||
test_A1 = sample[0][1]
|
||||
test_A2 = sample[0][2]
|
||||
|
||||
test_B0 = sample[1][0]
|
||||
test_B1 = sample[1][1]
|
||||
test_B2 = sample[1][2]
|
||||
|
||||
G_view_result = self.G_view([test_A0, test_A1, test_A2, test_B0, test_B1, test_B2])
|
||||
|
||||
fake_A0, fake_A1, pred_A2, rec_A2, fake_B0, fake_B1, pred_B2, rec_B2 = [ x[0] / 2 + 0.5 for x in G_view_result]
|
||||
test_A0, test_A1, test_A2, test_B0, test_B1, test_B2 = [ x[0] / 2 + 0.5 for x in [test_A0, test_A1, test_A2, test_B0, test_B1, test_B2] ]
|
||||
|
||||
|
||||
r = np.concatenate ((np.concatenate ( (test_A0, test_A1, test_A2, pred_A2, fake_B0, fake_B1, rec_A2), axis=1),
|
||||
np.concatenate ( (test_B0, test_B1, test_B2, pred_B2, fake_A0, fake_A1, rec_B2), axis=1)
|
||||
), axis=0)
|
||||
|
||||
return [ ('RecycleGAN, A0-A1-A2-PA2-FB0-FB1-RA2, B0-B1-B2-PB2-FA0-FA1-RB2, ', r ) ]
|
||||
|
||||
def predictor_func (self, face):
|
||||
x = self.G_convert ( [ np.expand_dims(face *2 - 1,0)] )[0]
|
||||
return x[0] / 2 + 0.5
|
||||
|
||||
#override
|
||||
def get_converter(self, **in_options):
|
||||
from models import ConverterImage
|
||||
|
||||
return ConverterImage(self.predictor_func, predictor_input_size=128, output_size=128, **in_options)
|
||||
|
||||
|
||||
|
||||
1
models/Model_RecycleGAN/__init__.py
Normal file
1
models/Model_RecycleGAN/__init__.py
Normal file
|
|
@ -0,0 +1 @@
|
|||
from .Model import Model
|
||||
Loading…
Add table
Add a link
Reference in a new issue